From 6cda4cf612e7a4c1c710b50812502c52fb0173b4 Mon Sep 17 00:00:00 2001
From: James Sully <sullyj3@gmail.com>
Date: Sat, 2 Sep 2023 16:02:35 +1000
Subject: [PATCH] revset: Add optional argument n to ancestors() in revset
 language

---
 CHANGELOG.md             |  4 ++++
 docs/revsets.md          |  3 ++-
 lib/src/revset.rs        | 12 +++++++++---
 lib/tests/test_revset.rs | 27 +++++++++++++++++++++++++++
 4 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 85192a89dd..43be66d363 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### New features
 
+* The `ancestors()` revset function now takes an optional `depth` argument 
+  to specify how many ancestors to return. For example, use `jj log -r 'ancestors(@, 5)`
+  to view the last 5 commits.
+
 ### Fixed bugs
 
 ## [0.9.0] - 2023-09-06
diff --git a/docs/revsets.md b/docs/revsets.md
index b6ee6f97fd..620bff8415 100644
--- a/docs/revsets.md
+++ b/docs/revsets.md
@@ -79,7 +79,8 @@ revsets (expressions) as arguments.
 
 * `parents(x)`: Same as `x-`.
 * `children(x)`: Same as `x+`.
-* `ancestors(x)`: Same as `:x`.
+* `ancestors(x[, depth])`: Last `depth` ancestors of `x` (including `x` 
+  itself). `ancestors(x)` is the same as `:x`, returning all ancestors.
 * `descendants(x)`: Same as `x::`.
 * `connected(x)`: Same as `x::x`. Useful when `x` includes several commits.
 * `all()`: All visible commits in the repo.
diff --git a/lib/src/revset.rs b/lib/src/revset.rs
index 098e1ac21e..4c13f79c20 100644
--- a/lib/src/revset.rs
+++ b/lib/src/revset.rs
@@ -1168,9 +1168,15 @@ static BUILTIN_FUNCTION_MAP: Lazy<HashMap<&'static str, RevsetFunction>> = Lazy:
         Ok(expression.children())
     });
     map.insert("ancestors", |name, arguments_pair, state| {
-        let arg = expect_one_argument(name, arguments_pair)?;
-        let expression = parse_expression_rule(arg.into_inner(), state)?;
-        Ok(expression.ancestors())
+        let ([descendents_arg], [depth_opt_arg]) = expect_arguments(name, arguments_pair)?;
+        let descendents = parse_expression_rule(descendents_arg.into_inner(), state)?;
+        let generation = if let Some(depth_arg) = depth_opt_arg {
+            let depth = parse_function_argument_as_literal("integer", name, depth_arg, state)?;
+            0..depth
+        } else {
+            GENERATION_RANGE_FULL
+        };
+        Ok(descendents.ancestors_range(generation))
     });
     map.insert("descendants", |name, arguments_pair, state| {
         let arg = expect_one_argument(name, arguments_pair)?;
diff --git a/lib/tests/test_revset.rs b/lib/tests/test_revset.rs
index e6696f1ae7..b1ca5ac063 100644
--- a/lib/tests/test_revset.rs
+++ b/lib/tests/test_revset.rs
@@ -1205,6 +1205,33 @@ fn test_evaluate_expression_ancestors(use_git: bool) {
             root_commit.id().clone(),
         ]
     );
+
+    // Can find last n ancestors of a commit
+    assert_eq!(
+        resolve_commit_ids(
+            mut_repo,
+            &format!("ancestors({}, 0)", commit2.id().hex())
+        ),
+        vec![]
+    );
+    assert_eq!(
+        resolve_commit_ids(
+            mut_repo,
+            &format!("ancestors({}, 1)", commit3.id().hex())
+        ),
+        vec![commit3.id().clone()]
+    );
+    assert_eq!(
+        resolve_commit_ids(
+            mut_repo,
+            &format!("ancestors({}, 3)", commit3.id().hex())
+        ),
+        vec![
+            commit3.id().clone(),
+            commit2.id().clone(),
+            commit1.id().clone(),
+        ]
+    );
 }
 
 #[test_case(false ; "local backend")]