From 832cf8d838c30ce460a56d175c0bdfc45bddff5b Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Tue, 22 Nov 2022 16:03:30 +0900
Subject: [PATCH] index: add RevWalk variant that tracks generation

Even though I couldn't determine if RevWalkGenerationRange has a measurable
cost compared to RevWalk, I'm not comfortable with enabling generation
tracking by default. So this patch adds a separate struct. I duplicated
Iterator::next() method as it seemed rather complicated to extract a common
iterator wrapper.

Actual filtering function and tests will be added by the next commit.
---
 lib/src/index.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/lib/src/index.rs b/lib/src/index.rs
index 13e68de1b..82d9f7ff4 100644
--- a/lib/src/index.rs
+++ b/lib/src/index.rs
@@ -1102,6 +1102,52 @@ impl<'a> Iterator for RevWalk<'a> {
     }
 }
 
+#[derive(Clone)]
+pub struct RevWalkGenerationRange<'a> {
+    queue: RevWalkQueue<'a, u32>,
+}
+
+impl<'a> Iterator for RevWalkGenerationRange<'a> {
+    type Item = IndexEntry<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(item) = self.queue.pop() {
+            if let RevWalkWorkItemState::Wanted(mut known_gen) = item.state {
+                self.queue.push_wanted_parents(&item.entry.0, known_gen + 1);
+                while let Some(x) = self.queue.pop_eq(&item.entry.0) {
+                    // For wanted item, simply track all generation chains. This can
+                    // be optimized if the wanted range is just upper/lower bounded.
+                    // If the range is fully bounded and if the range is wide, we
+                    // can instead extend 'gen' to a range of the same width, and
+                    // merge overlapping generation ranges.
+                    match x.state {
+                        RevWalkWorkItemState::Wanted(gen) if known_gen != gen => {
+                            self.queue.push_wanted_parents(&item.entry.0, gen + 1);
+                            known_gen = gen;
+                        }
+                        RevWalkWorkItemState::Wanted(_) => {}
+                        RevWalkWorkItemState::Unwanted => unreachable!(),
+                    }
+                }
+                return Some(item.entry.0);
+            } else if self.queue.items.len() == self.queue.unwanted_count {
+                // No more wanted entries to walk
+                debug_assert!(!self.queue.items.iter().any(|x| x.is_wanted()));
+                return None;
+            } else {
+                self.queue.skip_while_eq(&item.entry.0);
+                self.queue.push_unwanted_parents(&item.entry.0);
+            }
+        }
+
+        debug_assert_eq!(
+            self.queue.items.iter().filter(|x| !x.is_wanted()).count(),
+            self.queue.unwanted_count
+        );
+        None
+    }
+}
+
 impl IndexSegment for ReadonlyIndex {
     fn segment_num_parent_commits(&self) -> u32 {
         self.num_parent_commits