dag_walk: heads: use bfs and terminate if frontier has single item

We use `heads_ok()` for finding the head operations when there are
multiple current op heads. The current DFS-based algortihm needs to
always walk all the way to the root. That can be expensive when the
operations are slow to retrieve. In the common case where there are
two operations close to each other in the graph, we should be able to
terminate the search once we've reached the common ancestor. This
patch replaces the DFS by a BFS and adds the early termination.

lib/src/dag_walk.rs

@@ -488,14 +488,28 @@ where
 {
     let start: Vec<T> = start.into_iter().try_collect()?;
     let mut heads: HashSet<T> = start.iter().cloned().collect();
-    dfs_ok(start.into_iter().map(Ok), id_fn, |node| {
+    // Do a BFS until we have only one item left in the frontier. That frontier must
-        let neighbors: Vec<Result<T, E>> = neighbors_fn(node).into_iter().collect();
+    // have originated from one of the heads, and since there can't be cycles,
-        for neighbor in neighbors.iter().filter_map(|x| x.as_ref().ok()) {
+    // it won't be able to eliminate any other heads.
-            heads.remove(neighbor);
+    let mut frontier: Vec<T> = heads.iter().cloned().collect();
+    let mut visited: HashSet<ID> = heads.iter().map(&id_fn).collect();
+    let mut root_reached = false;
+    while frontier.len() > 1 || (!frontier.is_empty() && root_reached) {
+        frontier = frontier
+            .iter()
+            .flat_map(|node| {
+                let neighbors = neighbors_fn(node).into_iter().collect_vec();
+                if neighbors.is_empty() {
+                    root_reached = true;
+                }
+                neighbors
+            })
+            .try_collect()?;
+        for node in &frontier {
+            heads.remove(node);
         }
-        neighbors
+        frontier.retain(|node| visited.insert(id_fn(node)));
-    })
+    }
-    .try_for_each(|r| r.map(|_| ()))?;
     Ok(heads)
 }
 
@@ -1211,10 +1225,14 @@ mod tests {
         let neighbors_fn = |node: &char| neighbors[node].clone();
 
         let result = heads_ok([Ok('C')], id_fn, neighbors_fn);
-        assert_eq!(result, Err('X'));
+        assert_eq!(result, Ok(hashset! {'C'}));
         let result = heads_ok([Ok('B')], id_fn, neighbors_fn);
-        assert_eq!(result, Err('X'));
+        assert_eq!(result, Ok(hashset! {'B'}));
         let result = heads_ok([Ok('A')], id_fn, neighbors_fn);
         assert_eq!(result, Ok(hashset! {'A'}));
+        let result = heads_ok([Ok('C'), Ok('B')], id_fn, neighbors_fn);
+        assert_eq!(result, Err('X'));
+        let result = heads_ok([Ok('C'), Ok('A')], id_fn, neighbors_fn);
+        assert_eq!(result, Err('X'));
     }
 }