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.
This commit is contained in:
Yuya Nishihara 2023-05-20 16:36:26 +09:00
parent 522308bb9c
commit 3ba544414c
3 changed files with 30 additions and 56 deletions

View file

@ -14,54 +14,32 @@
use std::collections::HashSet; use std::collections::HashSet;
use std::hash::Hash; use std::hash::Hash;
use std::iter::Iterator; use std::iter;
pub struct BfsIter<'id_fn, 'neighbors_fn, T, ID, NI> { pub fn bfs<T, ID, II, 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>(
start: II, start: II,
id_fn: Box<dyn Fn(&T) -> ID + 'id_fn>, id_fn: impl Fn(&T) -> ID,
neighbors_fn: Box<dyn FnMut(&T) -> NI + 'neighbors_fn>, mut neighbors_fn: impl FnMut(&T) -> NI,
) -> BfsIter<'id_fn, 'neighbors_fn, T, ID, NI> ) -> impl Iterator<Item = T>
where where
ID: Hash + Eq, ID: Hash + Eq,
II: IntoIterator<Item = T>, II: IntoIterator<Item = T>,
NI: IntoIterator<Item = T>, NI: IntoIterator<Item = T>,
{ {
BfsIter { let mut work: Vec<T> = start.into_iter().collect();
id_fn, let mut visited: HashSet<ID> = HashSet::new();
neighbors_fn, iter::from_fn(move || loop {
work: start.into_iter().collect(), let c = work.pop()?;
visited: Default::default(), 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. /// Returns neighbors before the node itself.
@ -158,17 +136,13 @@ where
{ {
let start: Vec<T> = start.into_iter().collect(); let start: Vec<T> = start.into_iter().collect();
let mut reachable: HashSet<T> = start.iter().cloned().collect(); let mut reachable: HashSet<T> = start.iter().cloned().collect();
for _node in bfs( for _node in bfs(start.into_iter(), id_fn, |node| {
start.into_iter(),
Box::new(id_fn),
Box::new(|node| {
let neighbors: Vec<T> = neighbors_fn(node).into_iter().collect(); let neighbors: Vec<T> = neighbors_fn(node).into_iter().collect();
for neighbor in &neighbors { for neighbor in &neighbors {
reachable.remove(neighbor); reachable.remove(neighbor);
} }
neighbors neighbors
}), }) {}
) {}
reachable reachable
} }

View file

@ -111,8 +111,8 @@ impl DefaultIndexStore {
let mut parent_op_id: Option<OperationId> = None; let mut parent_op_id: Option<OperationId> = None;
for op in dag_walk::bfs( for op in dag_walk::bfs(
vec![operation.clone()], vec![operation.clone()],
Box::new(|op: &Operation| op.id().clone()), |op: &Operation| op.id().clone(),
Box::new(|op: &Operation| op.parents()), |op: &Operation| op.parents(),
) { ) {
if operations_dir.join(op.id().hex()).is_file() { if operations_dir.join(op.id().hex()).is_file() {
if parent_op_id.is_none() { if parent_op_id.is_none() {

View file

@ -28,8 +28,8 @@ fn count_non_merge_operations(repo: &Arc<ReadonlyRepo>) -> usize {
for op_id in dag_walk::bfs( for op_id in dag_walk::bfs(
vec![op_id], vec![op_id],
Box::new(|op_id| op_id.clone()), |op_id| op_id.clone(),
Box::new(|op_id| op_store.read_operation(op_id).unwrap().parents), |op_id| op_store.read_operation(op_id).unwrap().parents,
) { ) {
if op_store.read_operation(&op_id).unwrap().parents.len() <= 1 { if op_store.read_operation(&op_id).unwrap().parents.len() <= 1 {
num_ops += 1; num_ops += 1;