From 20aa31336e2bb250e938d8d1c8bf1da05f5f46d3 Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Thu, 6 Apr 2023 21:33:11 +0900
Subject: [PATCH] revset: extract dag-range calculation to function

The returned reachable set can be reused to calculate roots() expression.
---
 lib/src/default_revset_engine.rs | 53 ++++++++++++++++++++------------
 1 file changed, 34 insertions(+), 19 deletions(-)

diff --git a/lib/src/default_revset_engine.rs b/lib/src/default_revset_engine.rs
index 9e7f5eb2a..2e9282bd7 100644
--- a/lib/src/default_revset_engine.rs
+++ b/lib/src/default_revset_engine.rs
@@ -674,25 +674,8 @@ impl<'index, 'heads> EvaluationContext<'index, 'heads> {
             RevsetExpression::DagRange { roots, heads } => {
                 let root_set = self.evaluate(roots)?;
                 let head_set = self.evaluate(heads)?;
-                let mut reachable: HashSet<_> =
-                    root_set.iter().map(|entry| entry.position()).collect();
-                let mut result = vec![];
-                let candidates = self.walk_ancestors(&*head_set).collect_vec();
-                for candidate in candidates.into_iter().rev() {
-                    if reachable.contains(&candidate.position())
-                        || candidate
-                            .parent_positions()
-                            .iter()
-                            .any(|parent_pos| reachable.contains(parent_pos))
-                    {
-                        reachable.insert(candidate.position());
-                        result.push(candidate);
-                    }
-                }
-                result.reverse();
-                Ok(Box::new(EagerRevset {
-                    index_entries: result,
-                }))
+                let (dag_range_set, _) = self.collect_dag_range(&*root_set, &*head_set);
+                Ok(Box::new(dag_range_set))
             }
             RevsetExpression::VisibleHeads => Ok(self.revset_for_commit_ids(self.visible_heads)),
             RevsetExpression::Heads(candidates) => {
@@ -828,6 +811,38 @@ impl<'index, 'heads> EvaluationContext<'index, 'heads> {
         self.composite_index.walk_revs(&head_ids, &[])
     }
 
+    /// Calculates `root_set:head_set`.
+    ///
+    /// The returned `HashSet` contains all `root_set` positions no matter if
+    /// the entry isn't an ancestor of the `head_set`. The `EagerRevset`
+    /// only contains entries reachable from the `head_set`.
+    fn collect_dag_range<'a, 'b, S, T>(
+        &self,
+        root_set: &S,
+        head_set: &T,
+    ) -> (EagerRevset<'index>, HashSet<IndexPosition>)
+    where
+        S: InternalRevset<'a> + ?Sized,
+        T: InternalRevset<'b> + ?Sized,
+    {
+        let mut reachable: HashSet<_> = root_set.iter().map(|entry| entry.position()).collect();
+        let mut index_entries = vec![];
+        let candidates = self.walk_ancestors(head_set).collect_vec();
+        for candidate in candidates.into_iter().rev() {
+            if reachable.contains(&candidate.position())
+                || candidate
+                    .parent_positions()
+                    .iter()
+                    .any(|parent_pos| reachable.contains(parent_pos))
+            {
+                reachable.insert(candidate.position());
+                index_entries.push(candidate);
+            }
+        }
+        index_entries.reverse();
+        (EagerRevset { index_entries }, reachable)
+    }
+
     fn revset_for_commit_ids(
         &self,
         commit_ids: &[CommitId],