commit_templater: support custom cache extensions

cli/examples/custom-commit-templater/main.rs

@@ -17,19 +17,24 @@ use jj_cli::commit_templater::{CommitTemplateBuildFnTable, CommitTemplateLanguag
 use jj_cli::template_builder::TemplateLanguage;
 use jj_cli::template_parser::{self, TemplateParseError};
 use jj_cli::templater::{TemplateFunction, TemplatePropertyError};
+use jj_lib::backend::CommitId;
 use jj_lib::commit::Commit;
+use jj_lib::extensions_map::ExtensionsMap;
 use jj_lib::object_id::ObjectId;
+use jj_lib::repo::Repo;
+use jj_lib::revset::RevsetExpression;
+use once_cell::sync::OnceCell;
 
 struct HexCounter;
 
-fn num_digits_in_id(commit: Commit) -> Result<i64, TemplatePropertyError> {
+fn num_digits_in_id(id: &CommitId) -> i64 {
     let mut count = 0;
-    for ch in commit.id().hex().chars() {
+    for ch in id.hex().chars() {
         if ch.is_ascii_digit() {
             count += 1;
         }
     }
-    Ok(count)
+    count
 }
 
 fn num_char_in_id(commit: Commit, ch_match: char) -> Result<i64, TemplatePropertyError> {
@@ -42,14 +47,57 @@ fn num_char_in_id(commit: Commit, ch_match: char) -> Result<i64, TemplatePropert
     Ok(count)
 }
 
+struct MostDigitsInId {
+    count: OnceCell<i64>,
+}
+
+impl MostDigitsInId {
+    fn new() -> Self {
+        Self {
+            count: OnceCell::new(),
+        }
+    }
+
+    fn count(&self, repo: &dyn Repo) -> i64 {
+        *self.count.get_or_init(|| {
+            RevsetExpression::all()
+                .evaluate_programmatic(repo)
+                .unwrap()
+                .iter()
+                .map(|id| num_digits_in_id(&id))
+                .max()
+                .unwrap_or(0)
+        })
+    }
+}
+
 impl CommitTemplateLanguageExtension for HexCounter {
     fn build_fn_table<'repo>(&self) -> CommitTemplateBuildFnTable<'repo> {
         let mut table = CommitTemplateBuildFnTable::empty();
+        table.commit_methods.insert(
+            "has_most_digits",
+            |language, _build_context, property, call| {
+                template_parser::expect_no_arguments(call)?;
+                let most_digits = language
+                    .cache_extension::<MostDigitsInId>()
+                    .unwrap()
+                    .count(language.repo());
+                Ok(
+                    language.wrap_boolean(TemplateFunction::new(property, move |commit| {
+                        Ok(num_digits_in_id(commit.id()) == most_digits)
+                    })),
+                )
+            },
+        );
         table.commit_methods.insert(
             "num_digits_in_id",
             |language, _build_context, property, call| {
                 template_parser::expect_no_arguments(call)?;
-                Ok(language.wrap_integer(TemplateFunction::new(property, num_digits_in_id)))
+                Ok(
+                    language.wrap_integer(TemplateFunction::new(property, |commit| {
+                        Ok(num_digits_in_id(commit.id()))
+                    })),
+                )
             },
         );
         table.commit_methods.insert(
@@ -78,6 +126,10 @@ impl CommitTemplateLanguageExtension for HexCounter {
 
         table
     }
+
+    fn build_cache_extensions(&self, extensions: &mut ExtensionsMap) {
+        extensions.insert(MostDigitsInId::new());
+    }
 }
 
 fn main() -> std::process::ExitCode {

cli/src/commit_templater.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::any::Any;
 use std::cmp::max;
 use std::collections::HashMap;
 use std::io;
@@ -20,6 +21,7 @@ use std::rc::Rc;
 use itertools::Itertools as _;
 use jj_lib::backend::{ChangeId, CommitId};
 use jj_lib::commit::Commit;
+use jj_lib::extensions_map::ExtensionsMap;
 use jj_lib::hex_util::to_reverse_hex;
 use jj_lib::id_prefix::IdPrefixContext;
 use jj_lib::object_id::ObjectId as _;
@@ -42,6 +44,8 @@ use crate::text_util;
 
 pub trait CommitTemplateLanguageExtension {
     fn build_fn_table<'repo>(&self) -> CommitTemplateBuildFnTable<'repo>;
+
+    fn build_cache_extensions(&self, extensions: &mut ExtensionsMap);
 }
 
 pub struct CommitTemplateLanguage<'repo> {
@@ -50,6 +54,7 @@ pub struct CommitTemplateLanguage<'repo> {
     id_prefix_context: &'repo IdPrefixContext,
     build_fn_table: CommitTemplateBuildFnTable<'repo>,
     keyword_cache: CommitKeywordCache,
+    cache_extensions: ExtensionsMap,
 }
 
 impl<'repo> CommitTemplateLanguage<'repo> {
@@ -62,15 +67,22 @@ impl<'repo> CommitTemplateLanguage<'repo> {
         extension: Option<&dyn CommitTemplateLanguageExtension>,
     ) -> Self {
         let mut build_fn_table = CommitTemplateBuildFnTable::builtin();
+        let mut cache_extensions = ExtensionsMap::empty();
+
+        // TODO: Extension methods should be refactored to be plural, to support
+        // multiple extensions in a dynamic load environment
         if let Some(extension) = extension {
             build_fn_table.merge(extension.build_fn_table());
+            extension.build_cache_extensions(&mut cache_extensions);
         }
+
         CommitTemplateLanguage {
             repo,
             workspace_id: workspace_id.clone(),
             id_prefix_context,
             build_fn_table,
             keyword_cache: CommitKeywordCache::default(),
+            cache_extensions,
         }
     }
 }
@@ -156,6 +168,10 @@ impl<'repo> CommitTemplateLanguage<'repo> {
         &self.keyword_cache
     }
 
+    pub fn cache_extension<T: Any>(&self) -> Option<&T> {
+        self.cache_extensions.get::<T>()
+    }
+
     pub fn wrap_commit(
         &self,
         property: impl TemplateProperty<Commit, Output = Commit> + 'repo,

lib/src/extensions_map.rs

@@ -0,0 +1,98 @@
+// Copyright 2024 The Jujutsu Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#![allow(missing_docs)]
+
+use std::any::{Any, TypeId};
+use std::collections::HashMap;
+
+/// Type-safe map that stores objects of arbitrary types.
+///
+/// This allows extensions to store and retrieve their own types unknown to
+/// jj_lib safely.
+#[derive(Default)]
+pub struct ExtensionsMap {
+    values: HashMap<TypeId, Box<dyn Any>>,
+}
+
+impl ExtensionsMap {
+    /// Creates an empty ExtensionsMap.
+    pub fn empty() -> Self {
+        Default::default()
+    }
+
+    /// Returns the specified type if it has already been inserted.
+    pub fn get<V: Any>(&self) -> Option<&V> {
+        self.values
+            .get(&TypeId::of::<V>())
+            .map(|v| v.downcast_ref::<V>().unwrap())
+    }
+
+    /// Inserts a new instance of the specified type.
+    ///
+    /// Requires that this type has not been inserted before.
+    pub fn insert<V: Any>(&mut self, value: V) {
+        assert!(self
+            .values
+            .insert(TypeId::of::<V>(), Box::new(value))
+            .is_none());
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    struct TestTypeA;
+    impl TestTypeA {
+        fn get_a(&self) -> &'static str {
+            "a"
+        }
+    }
+
+    struct TestTypeB;
+    impl TestTypeB {
+        fn get_b(&self) -> &'static str {
+            "b"
+        }
+    }
+
+    #[test]
+    fn test_empty() {
+        let extensions_map = ExtensionsMap::empty();
+        assert!(extensions_map.get::<TestTypeA>().is_none());
+        assert!(extensions_map.get::<TestTypeB>().is_none());
+    }
+
+    #[test]
+    fn test_retrieval() {
+        let mut extensions_map = ExtensionsMap::empty();
+        extensions_map.insert(TestTypeA);
+        extensions_map.insert(TestTypeB);
+        assert_eq!(
+            extensions_map
+                .get::<TestTypeA>()
+                .map(|a| a.get_a())
+                .unwrap_or(""),
+            "a"
+        );
+        assert_eq!(
+            extensions_map
+                .get::<TestTypeB>()
+                .map(|b| b.get_b())
+                .unwrap_or(""),
+            "b"
+        );
+    }
+}

lib/src/lib.rs

@@ -36,6 +36,7 @@ pub mod dag_walk;
 pub mod default_index;
 pub mod default_submodule_store;
 pub mod diff;
+pub mod extensions_map;
 pub mod file_util;
 pub mod files;
 pub mod fmt_util;