From b9fc6d420338b03f3812095223b48710da1c5a38 Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Wed, 25 Jan 2023 18:46:38 +0900
Subject: [PATCH] templater: rewrite divergent property by leveraging IdIndex

---
 lib/src/index.rs |  7 +++++++
 lib/src/repo.rs  | 10 ++++++++++
 src/templater.rs | 36 ++++++++++--------------------------
 3 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/lib/src/index.rs b/lib/src/index.rs
index 961de765c..34da05388 100644
--- a/lib/src/index.rs
+++ b/lib/src/index.rs
@@ -295,6 +295,13 @@ impl HexPrefix {
         })
     }
 
+    pub fn from_bytes(bytes: &[u8]) -> Self {
+        HexPrefix {
+            min_prefix_bytes: bytes.to_owned(),
+            has_odd_byte: false,
+        }
+    }
+
     pub fn hex(&self) -> String {
         let mut hex_string = hex::encode(&self.min_prefix_bytes);
         if self.has_odd_byte {
diff --git a/lib/src/repo.rs b/lib/src/repo.rs
index cf476d6d2..4cb83353b 100644
--- a/lib/src/repo.rs
+++ b/lib/src/repo.rs
@@ -92,6 +92,16 @@ impl<'a> RepoRef<'a> {
         }
     }
 
+    pub fn resolve_change_id(&self, change_id: &ChangeId) -> Option<Vec<IndexEntry<'a>>> {
+        // Replace this if we added more efficient lookup method.
+        let prefix = HexPrefix::from_bytes(change_id.as_bytes());
+        match self.resolve_change_id_prefix(&prefix) {
+            PrefixResolution::NoMatch => None,
+            PrefixResolution::SingleMatch(entries) => Some(entries),
+            PrefixResolution::AmbiguousMatch => panic!("complete change_id should be unambiguous"),
+        }
+    }
+
     pub fn resolve_change_id_prefix(
         &self,
         prefix: &HexPrefix,
diff --git a/src/templater.rs b/src/templater.rs
index 7b086252c..2f6b141e6 100644
--- a/src/templater.rs
+++ b/src/templater.rs
@@ -12,16 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::{HashMap, HashSet};
 use std::io;
-use std::ops::AddAssign;
 
 use itertools::Itertools;
-use jujutsu_lib::backend::{ChangeId, ObjectId, Signature, Timestamp};
+use jujutsu_lib::backend::{ObjectId, Signature, Timestamp};
 use jujutsu_lib::commit::Commit;
 use jujutsu_lib::op_store::WorkspaceId;
 use jujutsu_lib::repo::RepoRef;
-use jujutsu_lib::revset::RevsetExpression;
 use jujutsu_lib::rewrite::merge_commit_trees;
 
 use crate::formatter::Formatter;
@@ -347,36 +344,23 @@ impl TemplateProperty<Commit> for IsGitHeadProperty<'_> {
     }
 }
 
-pub struct DivergentProperty {
-    divergent_changes: HashSet<ChangeId>,
+pub struct DivergentProperty<'a> {
+    repo: RepoRef<'a>,
 }
 
-impl DivergentProperty {
-    pub fn new(repo: RepoRef) -> Self {
-        // TODO: Create a persistent index from change id to commit ids.
-        let mut commit_count_by_change: HashMap<ChangeId, i32> = HashMap::new();
-        for index_entry in RevsetExpression::all().evaluate(repo, None).unwrap().iter() {
-            let change_id = index_entry.change_id();
-            commit_count_by_change
-                .entry(change_id)
-                .or_default()
-                .add_assign(1);
-        }
-        let mut divergent_changes = HashSet::new();
-        for (change_id, count) in commit_count_by_change {
-            if count > 1 {
-                divergent_changes.insert(change_id);
-            }
-        }
-        Self { divergent_changes }
+impl<'a> DivergentProperty<'a> {
+    pub fn new(repo: RepoRef<'a>) -> Self {
+        DivergentProperty { repo }
     }
 }
 
-impl TemplateProperty<Commit> for DivergentProperty {
+impl TemplateProperty<Commit> for DivergentProperty<'_> {
     type Output = bool;
 
     fn extract(&self, context: &Commit) -> Self::Output {
-        self.divergent_changes.contains(context.change_id())
+        // The given commit could be hidden in e.g. obslog.
+        let maybe_entries = self.repo.resolve_change_id(context.change_id());
+        maybe_entries.map_or(0, |entries| entries.len()) > 1
     }
 }