revset: implement a Subgraph expression

CHANGELOG.md

@@ -102,6 +102,9 @@ to avoid letting the user edit the immutable one.
 * `jj rebase -r` now accepts `--insert-after` and `--insert-before` options to
   customize the location of the rebased revisions.
 
+* A new revset `subgraph(srcs, boundary_heads)` will return all commits that are
+  reachable from `srcs` without traversing ancestors of `boundary_heads`.
+
 ### Fixed bugs
 
 * Revsets now support `\`-escapes in string literal.

cli/tests/test_revset_output.rs

@@ -278,7 +278,7 @@ fn test_function_name_hint() {
       | ^----^
       |
       = Function "branch" doesn't exist
-    Hint: Did you mean "branches"?
+    Hint: Did you mean "branches", "subgraph"?
     "###);
 
     // Both builtin function and function alias should be suggested
@@ -308,7 +308,7 @@ fn test_function_name_hint() {
       | ^----^
       |
       = Function "branch" doesn't exist
-    Hint: Did you mean "branches"?
+    Hint: Did you mean "branches", "subgraph"?
     "###);
 }
 

docs/revsets.md

@@ -87,6 +87,10 @@ revsets (expressions) as arguments.
 
 * `descendants(x)`: Same as `x::`.
 
+* `subgraph(s, b)`: All commits reachable from `s` without entering `::b`.
+  Equivalent to `(b..s)::` for simple graph structures, but distinct when more
+  complex merge commit graphs are involved.
+
 * `connected(x)`: Same as `x::x`. Useful when `x` includes several commits.
 
 * `all()`: All visible commits in the repo.

lib/src/default_index/revset_engine.rs

@@ -16,7 +16,7 @@
 
 use std::cell::RefCell;
 use std::cmp::{Ordering, Reverse};
-use std::collections::{BTreeSet, BinaryHeap, HashSet};
+use std::collections::{BTreeSet, BinaryHeap, HashMap, HashSet};
 use std::ops::Range;
 use std::rc::Rc;
 use std::sync::Arc;
@@ -709,6 +709,96 @@ pub fn evaluate<I: AsCompositeIndex + Clone>(
     Ok(RevsetImpl::new(internal_revset, index))
 }
 
+#[derive(Clone, Eq, PartialEq)]
+struct UnionFindNode {
+    root: IndexPosition,
+    size: u32,
+}
+
+// A disjoint-set representation of connected index positions.
+struct UnionFind {
+    map: HashMap<IndexPosition, UnionFindNode>,
+}
+
+impl UnionFind {
+    fn new(index: &CompositeIndex, positions: &[IndexPosition]) -> Self {
+        let mut union_find = Self {
+            map: HashMap::new(),
+        };
+        let domain: HashSet<_> = positions.iter().collect();
+        for pos in positions {
+            for parent in index.entry_by_pos(*pos).parent_positions() {
+                if domain.contains(&parent) {
+                    union_find.union(*pos, parent);
+                }
+            }
+        }
+
+        // Make sure all the roots are stable
+        for pos in positions {
+            union_find.find_mut(*pos);
+        }
+
+        union_find
+    }
+
+    fn find(&self, pos: IndexPosition) -> &UnionFindNode {
+        self.map.get(&pos).unwrap()
+    }
+
+    fn find_mut(&mut self, pos: IndexPosition) -> UnionFindNode {
+        match self.map.get(&pos) {
+            Some(node) => {
+                if node.root != pos {
+                    let new_root = self.find_mut(node.root);
+                    self.map.insert(pos, new_root.clone());
+                    new_root
+                } else {
+                    node.clone()
+                }
+            }
+            None => {
+                let node = UnionFindNode { root: pos, size: 1 };
+                self.map.insert(pos, node.clone());
+                node
+            }
+        }
+    }
+
+    fn union(&mut self, a: IndexPosition, b: IndexPosition) {
+        let a = self.find_mut(a);
+        let b = self.find_mut(b);
+        if a.root == b.root {
+            return;
+        }
+
+        // Merge the roots.
+        let new_root = UnionFindNode {
+            root: if a.size < b.size { b.root } else { a.root },
+            size: a.size + b.size,
+        };
+        self.map.insert(a.root, new_root.clone());
+        self.map.insert(b.root, new_root);
+    }
+
+    fn all_with_roots(self, roots: &HashSet<IndexPosition>) -> Vec<IndexPosition> {
+        let mut vec = self
+            .map
+            .into_iter()
+            .filter_map(|(pos, node)| {
+                if roots.contains(&node.root) {
+                    Some(pos)
+                } else {
+                    None
+                }
+            })
+            .sorted()
+            .collect_vec();
+        vec.reverse();
+        vec
+    }
+}
+
 struct EvaluationContext<'index> {
     store: Arc<Store>,
     index: &'index CompositeIndex,
@@ -829,6 +919,36 @@ impl<'index> EvaluationContext<'index> {
                     Ok(Box::new(EagerRevset { positions }))
                 }
             }
+            ResolvedExpression::Subgraph {
+                sources,
+                boundary_heads,
+                visible_heads,
+            } => {
+                let domain = ResolvedExpression::Range {
+                    roots: boundary_heads.clone(),
+                    heads: visible_heads.clone(),
+                    generation: GENERATION_RANGE_FULL,
+                };
+                let candidates = self
+                    .evaluate(&domain)?
+                    .positions()
+                    .attach(index)
+                    .collect_vec();
+                let union_find = UnionFind::new(index, &candidates);
+
+                let subgraph_roots: HashSet<_> = self
+                    .evaluate(&ResolvedExpression::Intersection(
+                        sources.clone(),
+                        domain.into(),
+                    ))?
+                    .positions()
+                    .attach(index)
+                    .map(|pos| union_find.find(pos).root)
+                    .collect();
+                Ok(Box::new(EagerRevset {
+                    positions: union_find.all_with_roots(&subgraph_roots),
+                }))
+            }
             ResolvedExpression::Heads(candidates) => {
                 let candidate_set = self.evaluate(candidates)?;
                 let head_positions: BTreeSet<_> =

lib/src/revset.rs

@@ -190,6 +190,11 @@ pub enum RevsetExpression {
         heads: Rc<RevsetExpression>,
         // TODO: maybe add generation_from_roots/heads?
     },
+    // Commits reachable from "sources" without traversing ancestors of "boundary_heads"
+    Subgraph {
+        sources: Rc<RevsetExpression>,
+        boundary_heads: Rc<RevsetExpression>,
+    },
     Heads(Rc<RevsetExpression>),
     Roots(Rc<RevsetExpression>),
     Latest {
@@ -374,6 +379,18 @@ impl RevsetExpression {
         self.dag_range_to(self)
     }
 
+    /// Commits connected to `sources` without traversing ancestors of
+    /// `boundary_heads`.
+    pub fn subgraph(
+        sources: Rc<RevsetExpression>,
+        boundary_heads: Rc<RevsetExpression>,
+    ) -> Rc<RevsetExpression> {
+        Rc::new(RevsetExpression::Subgraph {
+            sources,
+            boundary_heads,
+        })
+    }
+
     /// Commits reachable from `heads` but not from `self`.
     pub fn range(
         self: &Rc<RevsetExpression>,
@@ -502,6 +519,13 @@ pub enum ResolvedExpression {
         heads: Box<ResolvedExpression>,
         generation_from_roots: Range<u64>,
     },
+    /// Commits reachable from `sources` without traversing ancestors of
+    /// `boundary_heads`.
+    Subgraph {
+        sources: Box<ResolvedExpression>,
+        boundary_heads: Box<ResolvedExpression>,
+        visible_heads: Box<ResolvedExpression>,
+    },
     Heads(Box<ResolvedExpression>),
     Roots(Box<ResolvedExpression>),
     Latest {
@@ -630,6 +654,12 @@ static BUILTIN_FUNCTION_MAP: Lazy<HashMap<&'static str, RevsetFunction>> = Lazy:
         let candidates = parse_expression_rule(arg.into_inner(), state)?;
         Ok(candidates.connected())
     });
+    map.insert("subgraph", |name, arguments_pair, state| {
+        let ([source_arg, boundary_heads_arg], []) = expect_arguments(name, arguments_pair)?;
+        let sources = parse_expression_rule(source_arg.into_inner(), state)?;
+        let boundary_heads = parse_expression_rule(boundary_heads_arg.into_inner(), state)?;
+        Ok(RevsetExpression::subgraph(sources, boundary_heads))
+    });
     map.insert("none", |name, arguments_pair, _state| {
         expect_no_arguments(name, arguments_pair)?;
         Ok(RevsetExpression::none())
@@ -958,6 +988,15 @@ fn try_transform_expression<E>(
                 transform_rec_pair((roots, heads), pre, post)?
                     .map(|(roots, heads)| RevsetExpression::DagRange { roots, heads })
             }
+            RevsetExpression::Subgraph {
+                sources,
+                boundary_heads,
+            } => transform_rec_pair((sources, boundary_heads), pre, post)?.map(
+                |(sources, boundary_heads)| RevsetExpression::Subgraph {
+                    sources,
+                    boundary_heads,
+                },
+            ),
             RevsetExpression::Heads(candidates) => {
                 transform_rec(candidates, pre, post)?.map(RevsetExpression::Heads)
             }
@@ -1746,6 +1785,14 @@ impl VisibilityResolutionContext<'_> {
                 heads: self.resolve(heads).into(),
                 generation_from_roots: GENERATION_RANGE_FULL,
             },
+            RevsetExpression::Subgraph {
+                sources,
+                boundary_heads,
+            } => ResolvedExpression::Subgraph {
+                sources: self.resolve(sources).into(),
+                boundary_heads: self.resolve(boundary_heads).into(),
+                visible_heads: self.resolve_visible_heads().into(),
+            },
             RevsetExpression::Heads(candidates) => {
                 ResolvedExpression::Heads(self.resolve(candidates).into())
             }
@@ -1831,6 +1878,7 @@ impl VisibilityResolutionContext<'_> {
             | RevsetExpression::Descendants { .. }
             | RevsetExpression::Range { .. }
             | RevsetExpression::DagRange { .. }
+            | RevsetExpression::Subgraph { .. }
             | RevsetExpression::Heads(_)
             | RevsetExpression::Roots(_)
             | RevsetExpression::Latest { .. } => {