From 73fb922517492b7ce8a6fe310468c2950ea63729 Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Sat, 22 Apr 2023 16:16:51 +0900
Subject: [PATCH] index: reimplement collect_dag_range() of revset engine as
 iterator

I'm going to remove 'index lifetime from InternalRevset so Revset<'static>
can be easily constructed from DefaultReadonlyIndex. As the first step, this
series removes some lifetime complexity from EvaluationContext methods.

We don't need an descendant iterator API, but it helps to add separate function
to collect into HashSet<IndexPosition> instead of returning a pair of
ordered vec and set.
---
 lib/src/default_index_store.rs   | 53 ++++++++++++++++++++++++++++++++
 lib/src/default_revset_engine.rs | 20 ++----------
 2 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/lib/src/default_index_store.rs b/lib/src/default_index_store.rs
index 16236534c..22d1127a7 100644
--- a/lib/src/default_index_store.rs
+++ b/lib/src/default_index_store.rs
@@ -21,6 +21,7 @@ use std::fmt::{Debug, Formatter};
 use std::fs::File;
 use std::hash::{Hash, Hasher};
 use std::io::{Read, Write};
+use std::iter::FusedIterator;
 use std::ops::Range;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -1381,6 +1382,18 @@ impl<'a> RevWalk<'a> {
         self.take_while(move |entry| entry.position() >= bottom_position)
     }
 
+    /// Fully consumes the ancestors and walks back from `root_positions`.
+    ///
+    /// The returned iterator yields entries in order of ascending index
+    /// position.
+    pub fn descendants(self, root_positions: &[IndexPosition]) -> RevWalkDescendants<'a> {
+        RevWalkDescendants {
+            candidate_entries: self.take_until_roots(root_positions).collect(),
+            root_positions: root_positions.iter().copied().collect(),
+            reachable_positions: HashSet::new(),
+        }
+    }
+
     /// Fully consumes the ancestors and walks back from `root_positions` within
     /// `generation_range`.
     ///
@@ -1589,6 +1602,46 @@ impl RevWalkItemGenerationRange {
     }
 }
 
+/// Walks descendants from the roots, in order of ascending index position.
+#[derive(Clone)]
+pub struct RevWalkDescendants<'a> {
+    candidate_entries: Vec<IndexEntry<'a>>,
+    root_positions: HashSet<IndexPosition>,
+    reachable_positions: HashSet<IndexPosition>,
+}
+
+impl RevWalkDescendants<'_> {
+    /// Builds a set of index positions reachable from the roots.
+    ///
+    /// This is equivalent to `.map(|entry| entry.position()).collect()` on
+    /// the new iterator, but returns the internal buffer instead.
+    pub fn collect_positions_set(mut self) -> HashSet<IndexPosition> {
+        self.by_ref().for_each(drop);
+        self.reachable_positions
+    }
+}
+
+impl<'a> Iterator for RevWalkDescendants<'a> {
+    type Item = IndexEntry<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(candidate) = self.candidate_entries.pop() {
+            if self.root_positions.contains(&candidate.position())
+                || candidate
+                    .parent_positions()
+                    .iter()
+                    .any(|parent_pos| self.reachable_positions.contains(parent_pos))
+            {
+                self.reachable_positions.insert(candidate.position());
+                return Some(candidate);
+            }
+        }
+        None
+    }
+}
+
+impl FusedIterator for RevWalkDescendants<'_> {}
+
 /// Removes the greatest items (including duplicates) from the heap, returns
 /// one.
 fn dedup_pop<T: Ord>(heap: &mut BinaryHeap<T>) -> Option<T> {
diff --git a/lib/src/default_revset_engine.rs b/lib/src/default_revset_engine.rs
index fceb5bda4..9a359c0ad 100644
--- a/lib/src/default_revset_engine.rs
+++ b/lib/src/default_revset_engine.rs
@@ -725,23 +725,9 @@ impl<'index> EvaluationContext<'index> {
         T: InternalRevset<'b> + ?Sized,
     {
         let root_positions = root_set.iter().map(|entry| entry.position()).collect_vec();
-        let walk = self
-            .walk_ancestors(head_set)
-            .take_until_roots(&root_positions);
-        let root_positions: HashSet<_> = root_positions.into_iter().collect();
-        let mut reachable_positions = HashSet::new();
-        let mut index_entries = vec![];
-        for candidate in walk.collect_vec().into_iter().rev() {
-            if root_positions.contains(&candidate.position())
-                || candidate
-                    .parent_positions()
-                    .iter()
-                    .any(|parent_pos| reachable_positions.contains(parent_pos))
-            {
-                reachable_positions.insert(candidate.position());
-                index_entries.push(candidate);
-            }
-        }
+        let mut walk = self.walk_ancestors(head_set).descendants(&root_positions);
+        let mut index_entries = walk.by_ref().collect_vec();
+        let reachable_positions = walk.collect_positions_set();
         index_entries.reverse();
         (EagerRevset { index_entries }, reachable_positions)
     }