From 3ba544414cea6dbb0c66896d0b53395f9c423b94 Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Sat, 20 May 2023 16:36:26 +0900
Subject: [PATCH] dag_walk: unbox bfs() callback, use iter::from_fn() to
 implement iterator

I just wanted to remove syntactic noise from callers. iter::from_fn() helps
to avoid declaring struct with lots of type parameters.
---
 lib/src/dag_walk.rs                 | 78 ++++++++++-------------------
 lib/src/default_index_store.rs      |  4 +-
 lib/tests/test_commit_concurrent.rs |  4 +-
 3 files changed, 30 insertions(+), 56 deletions(-)

diff --git a/lib/src/dag_walk.rs b/lib/src/dag_walk.rs
index 582608b7f..fcce60671 100644
--- a/lib/src/dag_walk.rs
+++ b/lib/src/dag_walk.rs
@@ -14,54 +14,32 @@
 
 use std::collections::HashSet;
 use std::hash::Hash;
-use std::iter::Iterator;
+use std::iter;
 
-pub struct BfsIter<'id_fn, 'neighbors_fn, T, ID, NI> {
-    id_fn: Box<dyn Fn(&T) -> ID + 'id_fn>,
-    neighbors_fn: Box<dyn FnMut(&T) -> NI + 'neighbors_fn>,
-    work: Vec<T>,
-    visited: HashSet<ID>,
-}
-
-impl<T, ID, NI> Iterator for BfsIter<'_, '_, T, ID, NI>
-where
-    ID: Hash + Eq,
-    NI: IntoIterator<Item = T>,
-{
-    type Item = T;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        loop {
-            let c = self.work.pop()?;
-            let id = (self.id_fn)(&c);
-            if self.visited.contains(&id) {
-                continue;
-            }
-            for p in (self.neighbors_fn)(&c) {
-                self.work.push(p);
-            }
-            self.visited.insert(id);
-            return Some(c);
-        }
-    }
-}
-
-pub fn bfs<'id_fn, 'neighbors_fn, T, ID, II, NI>(
+pub fn bfs<T, ID, II, NI>(
     start: II,
-    id_fn: Box<dyn Fn(&T) -> ID + 'id_fn>,
-    neighbors_fn: Box<dyn FnMut(&T) -> NI + 'neighbors_fn>,
-) -> BfsIter<'id_fn, 'neighbors_fn, T, ID, NI>
+    id_fn: impl Fn(&T) -> ID,
+    mut neighbors_fn: impl FnMut(&T) -> NI,
+) -> impl Iterator<Item = T>
 where
     ID: Hash + Eq,
     II: IntoIterator<Item = T>,
     NI: IntoIterator<Item = T>,
 {
-    BfsIter {
-        id_fn,
-        neighbors_fn,
-        work: start.into_iter().collect(),
-        visited: Default::default(),
-    }
+    let mut work: Vec<T> = start.into_iter().collect();
+    let mut visited: HashSet<ID> = HashSet::new();
+    iter::from_fn(move || loop {
+        let c = work.pop()?;
+        let id = id_fn(&c);
+        if visited.contains(&id) {
+            continue;
+        }
+        for p in neighbors_fn(&c) {
+            work.push(p);
+        }
+        visited.insert(id);
+        return Some(c);
+    })
 }
 
 /// Returns neighbors before the node itself.
@@ -158,17 +136,13 @@ where
 {
     let start: Vec<T> = start.into_iter().collect();
     let mut reachable: HashSet<T> = start.iter().cloned().collect();
-    for _node in bfs(
-        start.into_iter(),
-        Box::new(id_fn),
-        Box::new(|node| {
-            let neighbors: Vec<T> = neighbors_fn(node).into_iter().collect();
-            for neighbor in &neighbors {
-                reachable.remove(neighbor);
-            }
-            neighbors
-        }),
-    ) {}
+    for _node in bfs(start.into_iter(), id_fn, |node| {
+        let neighbors: Vec<T> = neighbors_fn(node).into_iter().collect();
+        for neighbor in &neighbors {
+            reachable.remove(neighbor);
+        }
+        neighbors
+    }) {}
     reachable
 }
 
diff --git a/lib/src/default_index_store.rs b/lib/src/default_index_store.rs
index 5ebfd72b9..da37219b4 100644
--- a/lib/src/default_index_store.rs
+++ b/lib/src/default_index_store.rs
@@ -111,8 +111,8 @@ impl DefaultIndexStore {
         let mut parent_op_id: Option<OperationId> = None;
         for op in dag_walk::bfs(
             vec![operation.clone()],
-            Box::new(|op: &Operation| op.id().clone()),
-            Box::new(|op: &Operation| op.parents()),
+            |op: &Operation| op.id().clone(),
+            |op: &Operation| op.parents(),
         ) {
             if operations_dir.join(op.id().hex()).is_file() {
                 if parent_op_id.is_none() {
diff --git a/lib/tests/test_commit_concurrent.rs b/lib/tests/test_commit_concurrent.rs
index 0515377ae..78f39337e 100644
--- a/lib/tests/test_commit_concurrent.rs
+++ b/lib/tests/test_commit_concurrent.rs
@@ -28,8 +28,8 @@ fn count_non_merge_operations(repo: &Arc<ReadonlyRepo>) -> usize {
 
     for op_id in dag_walk::bfs(
         vec![op_id],
-        Box::new(|op_id| op_id.clone()),
-        Box::new(|op_id| op_store.read_operation(op_id).unwrap().parents),
+        |op_id| op_id.clone(),
+        |op_id| op_store.read_operation(op_id).unwrap().parents,
     ) {
         if op_store.read_operation(&op_id).unwrap().parents.len() <= 1 {
             num_ops += 1;