give descendants a range arg

CHANGELOG.md

@@ -28,6 +28,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * `jj fix` now supports configuring the default revset for `-s` using the
   `revsets.fix` config.
 
+* The `descendants()` revset function now accepts an optional `depth` argument;
+  like the `ancestors()` depth argument, it limits the depth of the set.
+
 ### Fixed bugs
 
 ## [0.18.0] - 2024-06-05

docs/revsets.md

@@ -87,7 +87,9 @@ revsets (expressions) as arguments.
   `ancestors(x, depth)` returns the ancestors of `x` limited to the given
   `depth`.
 
-* `descendants(x)`: Same as `x::`.
+* `descendants(x[, depth])`: `descendants(x)` is the same as `x::`.
+  `descendants(x, depth)` returns the descendants of `x` limited to the given
+  `depth`.
 
 * `reachable(srcs, domain)`: All commits reachable from `srcs` within
   `domain`, traversing all parent and child edges.

lib/src/revset.rs

@@ -339,18 +339,23 @@ impl RevsetExpression {
 
     /// Descendants of `self`, including `self`.
     pub fn descendants(self: &Rc<RevsetExpression>) -> Rc<RevsetExpression> {
-        Rc::new(RevsetExpression::Descendants {
-            roots: self.clone(),
-            generation: GENERATION_RANGE_FULL,
-        })
+        self.descendants_range(GENERATION_RANGE_FULL)
     }
 
     /// Descendants of `self` at an offset of `generation` ahead of `self`.
     /// The `generation` offset is zero-based starting from `self`.
     pub fn descendants_at(self: &Rc<RevsetExpression>, generation: u64) -> Rc<RevsetExpression> {
+        self.descendants_range(generation..(generation + 1))
+    }
+
+    /// Descendants of `self` in the given range.
+    pub fn descendants_range(
+        self: &Rc<RevsetExpression>,
+        generation_range: Range<u64>,
+    ) -> Rc<RevsetExpression> {
         Rc::new(RevsetExpression::Descendants {
             roots: self.clone(),
-            generation: generation..(generation + 1),
+            generation: generation_range,
         })
     }
 
@@ -572,9 +577,15 @@ static BUILTIN_FUNCTION_MAP: Lazy<HashMap<&'static str, RevsetFunction>> = Lazy:
         Ok(heads.ancestors_range(generation))
     });
     map.insert("descendants", |function, context| {
-        let [arg] = function.expect_exact_arguments()?;
-        let expression = lower_expression(arg, context)?;
-        Ok(expression.descendants())
+        let ([roots_arg], [depth_opt_arg]) = function.expect_arguments()?;
+        let roots = lower_expression(roots_arg, context)?;
+        let generation = if let Some(depth_arg) = depth_opt_arg {
+            let depth = expect_literal("integer", depth_arg)?;
+            0..depth
+        } else {
+            GENERATION_RANGE_FULL
+        };
+        Ok(roots.descendants_range(generation))
     });
     map.insert("connected", |function, context| {
         let [arg] = function.expect_exact_arguments()?;

lib/tests/test_revset.rs

@@ -1816,6 +1816,24 @@ fn test_evaluate_expression_descendants() {
             commit3.id().clone(),
         ]
     );
+
+    // Can find next n descendants of a commit
+    assert_eq!(
+        resolve_commit_ids(mut_repo, &format!("descendants({}, 0)", commit2.id().hex())),
+        vec![]
+    );
+    assert_eq!(
+        resolve_commit_ids(mut_repo, &format!("descendants({}, 1)", commit3.id().hex())),
+        vec![commit3.id().clone()]
+    );
+    assert_eq!(
+        resolve_commit_ids(mut_repo, &format!("descendants({}, 3)", commit3.id().hex())),
+        vec![
+            commit6.id().clone(),
+            commit5.id().clone(),
+            commit3.id().clone(),
+        ]
+    );
 }
 
 #[test]