From 30aa459d2a41905ec354762940c4ed96b5d0093e Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Thu, 15 Apr 2021 22:19:08 -0700
Subject: [PATCH] revsets: add a all_heads() revset function

This adds a `all_heads()` revset function, which contains all heads in
the view, i.e. including non-public heads and obsolete heads.
---
 lib/src/revset.pest      |  2 +-
 lib/src/revset.rs        | 21 +++++++++++++++++++++
 lib/tests/test_revset.rs | 23 +++++++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/lib/src/revset.pest b/lib/src/revset.pest
index 677c1e59c..e027dec7a 100644
--- a/lib/src/revset.pest
+++ b/lib/src/revset.pest
@@ -18,7 +18,7 @@ literal_string = { "\"" ~ (!"\"" ~ ANY)+ ~ "\"" }
 parents = { ":" }
 ancestors = { "*:" }
 
-function_name = @{ ASCII_ALPHANUMERIC+ }
+function_name = @{ (ASCII_ALPHANUMERIC | "_")+ }
 // The grammar accepts a string literal or an expression for function
 // arguments. We then decide when walking the parse tree if we
 // should interpret the string value of the literal string or the expression
diff --git a/lib/src/revset.rs b/lib/src/revset.rs
index a60c35965..4a8b3195b 100644
--- a/lib/src/revset.rs
+++ b/lib/src/revset.rs
@@ -99,6 +99,7 @@ pub enum RevsetExpression {
     Symbol(String),
     Parents(Box<RevsetExpression>),
     Ancestors(Box<RevsetExpression>),
+    AllHeads,
 }
 
 fn parse_expression_rule(mut pairs: Pairs<Rule>) -> Result<RevsetExpression, RevsetParseError> {
@@ -164,6 +165,16 @@ fn parse_function_expression(
                 })
             }
         }
+        "all_heads" => {
+            if arg_count == 0 {
+                Ok(RevsetExpression::AllHeads)
+            } else {
+                Err(RevsetParseError::InvalidFunctionArguments {
+                    name,
+                    message: "Expected 0 arguments".to_string(),
+                })
+            }
+        }
         _ => Err(RevsetParseError::NoSuchFunction(name)),
     }
 }
@@ -269,5 +280,15 @@ pub fn evaluate_expression<'repo>(
             let walk = repo.index().walk_revs(&base_ids, &[]);
             Ok(Box::new(RevWalkRevset { walk }))
         }
+        RevsetExpression::AllHeads => {
+            let index = repo.index();
+            let heads = repo.view().heads();
+            let mut index_entries: Vec<_> = heads
+                .iter()
+                .map(|id| index.entry_by_id(id).unwrap())
+                .collect();
+            index_entries.sort_by_key(|b| Reverse(b.position()));
+            Ok(Box::new(EagerRevset { index_entries }))
+        }
     }
 }
diff --git a/lib/tests/test_revset.rs b/lib/tests/test_revset.rs
index e4d3c162a..e8a619886 100644
--- a/lib/tests/test_revset.rs
+++ b/lib/tests/test_revset.rs
@@ -399,3 +399,26 @@ fn test_evaluate_expression_ancestors(use_git: bool) {
 
     tx.discard();
 }
+
+#[test_case(false ; "local store")]
+#[test_case(true ; "git store")]
+fn test_evaluate_expression_all_heads(use_git: bool) {
+    let settings = testutils::user_settings();
+    let (_temp_dir, repo) = testutils::init_repo(&settings, use_git);
+
+    let mut tx = repo.start_transaction("test");
+    let mut_repo = tx.mut_repo();
+
+    let wc_commit = repo.working_copy_locked().current_commit();
+    let commit1 = testutils::create_random_commit(&settings, &repo).write_to_repo(mut_repo);
+    let commit2 = testutils::create_random_commit(&settings, &repo)
+        .set_parents(vec![commit1.id().clone()])
+        .write_to_repo(mut_repo);
+
+    assert_eq!(
+        resolve_commit_ids(mut_repo.as_repo_ref(), "all_heads()"),
+        vec![commit2.id().clone(), wc_commit.id().clone()]
+    );
+
+    tx.discard();
+}