From b52cfc156c131a727cdcc99e46b7885a9c89abaa Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Sun, 18 Apr 2021 17:30:21 -0700
Subject: [PATCH] revsets: add a public_heads() revset function

---
 lib/src/revset.rs        | 21 +++++++++++++++++++++
 lib/tests/test_revset.rs | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/lib/src/revset.rs b/lib/src/revset.rs
index c9df8131d..bf75dcc6f 100644
--- a/lib/src/revset.rs
+++ b/lib/src/revset.rs
@@ -102,6 +102,7 @@ pub enum RevsetExpression {
     Parents(Box<RevsetExpression>),
     Ancestors(Box<RevsetExpression>),
     AllHeads,
+    PublicHeads,
     NonObsoleteHeads(Box<RevsetExpression>),
     Description {
         needle: String,
@@ -257,6 +258,16 @@ fn parse_function_expression(
                 })
             }
         }
+        "public_heads" => {
+            if arg_count == 0 {
+                Ok(RevsetExpression::PublicHeads)
+            } else {
+                Err(RevsetParseError::InvalidFunctionArguments {
+                    name,
+                    message: "Expected 0 arguments".to_string(),
+                })
+            }
+        }
         "description" => {
             if !(1..=2).contains(&arg_count) {
                 return Err(RevsetParseError::InvalidFunctionArguments {
@@ -585,6 +596,16 @@ pub fn evaluate_expression<'revset, 'repo: 'revset>(
             let base_set = evaluate_expression(repo, base_expression.as_ref())?;
             Ok(non_obsolete_heads(repo, base_set))
         }
+        RevsetExpression::PublicHeads => {
+            let index = repo.index();
+            let heads = repo.view().public_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 }))
+        }
         RevsetExpression::Description {
             needle,
             base_expression,
diff --git a/lib/tests/test_revset.rs b/lib/tests/test_revset.rs
index 3e7869997..ab466d350 100644
--- a/lib/tests/test_revset.rs
+++ b/lib/tests/test_revset.rs
@@ -519,6 +519,40 @@ fn test_evaluate_expression_all_heads(use_git: bool) {
     tx.discard();
 }
 
+#[test_case(false ; "local store")]
+#[test_case(true ; "git store")]
+fn test_evaluate_expression_public_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 commit1 = testutils::create_random_commit(&settings, &repo).write_to_repo(mut_repo);
+    let commit2 = testutils::create_random_commit(&settings, &repo).write_to_repo(mut_repo);
+
+    // Can get public heads with root commit as only public head
+    assert_eq!(
+        resolve_commit_ids(mut_repo.as_repo_ref(), "public_heads()"),
+        // TODO: This should include the root commit
+        vec![]
+    );
+    // Can get public heads with a single public head
+    mut_repo.add_public_head(&commit1);
+    assert_eq!(
+        resolve_commit_ids(mut_repo.as_repo_ref(), "public_heads()"),
+        vec![commit1.id().clone()]
+    );
+    // Can get public heads with multiple public head
+    mut_repo.add_public_head(&commit2);
+    assert_eq!(
+        resolve_commit_ids(mut_repo.as_repo_ref(), "public_heads()"),
+        vec![commit2.id().clone(), commit1.id().clone()]
+    );
+
+    tx.discard();
+}
+
 #[test_case(false ; "local store")]
 #[test_case(true ; "git store")]
 fn test_evaluate_expression_obsolete(use_git: bool) {