operation_templater: support extensions of the template language

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

@@ -0,0 +1,97 @@
+// 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.
+
+use jj_cli::cli_util::CliRunner;
+use jj_cli::operation_templater::{
+    OperationTemplateBuildFnTable, OperationTemplateLanguageExtension,
+};
+use jj_cli::template_builder::TemplateLanguage;
+use jj_cli::template_parser::{self, TemplateParseError};
+use jj_cli::templater::{TemplateFunction, TemplatePropertyError};
+use jj_lib::extensions_map::ExtensionsMap;
+use jj_lib::object_id::ObjectId;
+use jj_lib::op_store::OperationId;
+use jj_lib::operation::Operation;
+
+struct HexCounter;
+
+fn num_digits_in_id(id: &OperationId) -> i64 {
+    let mut count = 0;
+    for ch in id.hex().chars() {
+        if ch.is_ascii_digit() {
+            count += 1;
+        }
+    }
+    count
+}
+
+fn num_char_in_id(operation: Operation, ch_match: char) -> Result<i64, TemplatePropertyError> {
+    let mut count = 0;
+    for ch in operation.id().hex().chars() {
+        if ch == ch_match {
+            count += 1;
+        }
+    }
+    Ok(count)
+}
+
+impl OperationTemplateLanguageExtension for HexCounter {
+    fn build_fn_table(&self) -> OperationTemplateBuildFnTable {
+        let mut table = OperationTemplateBuildFnTable::empty();
+        table.operation_methods.insert(
+            "num_digits_in_id",
+            |language, _build_context, property, call| {
+                template_parser::expect_no_arguments(call)?;
+                Ok(
+                    language.wrap_integer(TemplateFunction::new(property, |operation| {
+                        Ok(num_digits_in_id(operation.id()))
+                    })),
+                )
+            },
+        );
+        table.operation_methods.insert(
+            "num_char_in_id",
+            |language, _build_context, property, call| {
+                let [string_arg] = template_parser::expect_exact_arguments(call)?;
+                let char_arg =
+                    template_parser::expect_string_literal_with(string_arg, |string, span| {
+                        let chars: Vec<_> = string.chars().collect();
+                        match chars[..] {
+                            [ch] => Ok(ch),
+                            _ => Err(TemplateParseError::unexpected_expression(
+                                "Expected singular character argument",
+                                span,
+                            )),
+                        }
+                    })?;
+
+                Ok(
+                    language.wrap_integer(TemplateFunction::new(property, move |operation| {
+                        num_char_in_id(operation, char_arg)
+                    })),
+                )
+            },
+        );
+
+        table
+    }
+
+    fn build_cache_extensions(&self, _extensions: &mut ExtensionsMap) {}
+}
+
+fn main() -> std::process::ExitCode {
+    CliRunner::init()
+        .set_operation_template_extension(Box::new(HexCounter))
+        .run()
+}

cli/src/cli_util.rs

@@ -86,6 +86,7 @@ use crate::git_util::{
     is_colocated_git_workspace, print_failed_git_export, print_git_import_stats,
 };
 use crate::merge_tools::{DiffEditor, MergeEditor, MergeToolConfigError};
+use crate::operation_templater::OperationTemplateLanguageExtension;
 use crate::template_builder::TemplateLanguage;
 use crate::template_parser::TemplateAliasesMap;
 use crate::templater::Template;
@@ -188,6 +189,7 @@ pub struct CommandHelper {
     settings: UserSettings,
     layered_configs: LayeredConfigs,
     commit_template_extension: Option<Arc<dyn CommitTemplateLanguageExtension>>,
+    operation_template_extension: Option<Arc<dyn OperationTemplateLanguageExtension>>,
     maybe_workspace_loader: Result<WorkspaceLoader, CommandError>,
     store_factories: StoreFactories,
     working_copy_factories: HashMap<String, Box<dyn WorkingCopyFactory>>,
@@ -252,6 +254,10 @@ impl CommandHelper {
         Ok(template_builder::parse(language, template_text, &aliases)?)
     }
 
+    pub fn operation_template_extension(&self) -> Option<&dyn OperationTemplateLanguageExtension> {
+        self.operation_template_extension.as_deref()
+    }
+
     pub fn workspace_loader(&self) -> Result<&WorkspaceLoader, CommandError> {
         self.maybe_workspace_loader.as_ref().map_err(Clone::clone)
     }
@@ -2314,6 +2320,7 @@ pub struct CliRunner {
     store_factories: Option<StoreFactories>,
     working_copy_factories: Option<HashMap<String, Box<dyn WorkingCopyFactory>>>,
     commit_template_extension: Option<Arc<dyn CommitTemplateLanguageExtension>>,
+    operation_template_extension: Option<Arc<dyn OperationTemplateLanguageExtension>>,
     dispatch_fn: CliDispatchFn,
     start_hook_fns: Vec<CliDispatchFn>,
     process_global_args_fns: Vec<ProcessGlobalArgsFn>,
@@ -2336,6 +2343,7 @@ impl CliRunner {
             store_factories: None,
             working_copy_factories: None,
             commit_template_extension: None,
+            operation_template_extension: None,
             dispatch_fn: Box::new(crate::commands::run_command),
             start_hook_fns: vec![],
             process_global_args_fns: vec![],
@@ -2377,6 +2385,14 @@ impl CliRunner {
         self
     }
 
+    pub fn set_operation_template_extension(
+        mut self,
+        operation_template_extension: Box<dyn OperationTemplateLanguageExtension>,
+    ) -> Self {
+        self.operation_template_extension = Some(operation_template_extension.into());
+        self
+    }
+
     pub fn add_start_hook(mut self, start_hook_fn: CliDispatchFn) -> Self {
         self.start_hook_fns.push(start_hook_fn);
         self
@@ -2497,6 +2513,7 @@ impl CliRunner {
             settings,
             layered_configs,
             commit_template_extension: self.commit_template_extension,
+            operation_template_extension: self.operation_template_extension,
             maybe_workspace_loader,
             store_factories: self.store_factories.unwrap_or_default(),
             working_copy_factories,

cli/src/commands/operation.rs

@@ -162,6 +162,7 @@ fn cmd_op_log(
         let language = OperationTemplateLanguage::new(
             repo_loader.op_store().root_operation_id(),
             current_op_id,
+            command.operation_template_extension(),
         );
         let text = match &args.template {
             Some(value) => value.to_owned(),

cli/src/operation_templater.rs

@@ -12,38 +12,61 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::any::Any;
+use std::collections::HashMap;
 use std::io;
 
 use itertools::Itertools as _;
+use jj_lib::extensions_map::ExtensionsMap;
 use jj_lib::object_id::ObjectId;
 use jj_lib::op_store::OperationId;
 use jj_lib::operation::Operation;
 
 use crate::formatter::Formatter;
 use crate::template_builder::{
-    self, BuildContext, CoreTemplateBuildFnTable, CoreTemplatePropertyKind, IntoTemplateProperty,
-    TemplateBuildMethodFnMap, TemplateLanguage,
+    self, merge_fn_map, BuildContext, CoreTemplateBuildFnTable, CoreTemplatePropertyKind,
+    IntoTemplateProperty, TemplateBuildMethodFnMap, TemplateLanguage,
 };
 use crate::template_parser::{self, FunctionCallNode, TemplateParseResult};
 use crate::templater::{
     IntoTemplate, PlainTextFormattedProperty, Template, TemplateFunction, TemplateProperty,
     TemplatePropertyFn, TimestampRange,
 };
 
+pub trait OperationTemplateLanguageExtension {
+    fn build_fn_table(&self) -> OperationTemplateBuildFnTable;
+
+    fn build_cache_extensions(&self, extensions: &mut ExtensionsMap);
+}
+
 pub struct OperationTemplateLanguage {
     root_op_id: OperationId,
     current_op_id: Option<OperationId>,
     build_fn_table: OperationTemplateBuildFnTable,
+    cache_extensions: ExtensionsMap,
 }
 
 impl OperationTemplateLanguage {
     /// Sets up environment where operation template will be transformed to
     /// evaluation tree.
-    pub fn new(root_op_id: &OperationId, current_op_id: Option<&OperationId>) -> Self {
+    pub fn new(
+        root_op_id: &OperationId,
+        current_op_id: Option<&OperationId>,
+        extension: Option<&dyn OperationTemplateLanguageExtension>,
+    ) -> Self {
+        let mut build_fn_table = OperationTemplateBuildFnTable::builtin();
+        let mut cache_extensions = ExtensionsMap::empty();
+        if let Some(extension) = extension {
+            let ext_table = extension.build_fn_table();
+            build_fn_table.merge(ext_table);
+            extension.build_cache_extensions(&mut cache_extensions);
+        }
+
         OperationTemplateLanguage {
             root_op_id: root_op_id.clone(),
             current_op_id: current_op_id.cloned(),
-            build_fn_table: OperationTemplateBuildFnTable::builtin(),
+            build_fn_table,
+            cache_extensions,
         }
     }
 }
@@ -94,6 +117,10 @@ impl TemplateLanguage<'static> for OperationTemplateLanguage {
 }
 
 impl OperationTemplateLanguage {
+    pub fn cache_extension<T: Any>(&self) -> Option<&T> {
+        self.cache_extensions.get::<T>()
+    }
+
     pub fn wrap_operation(
         &self,
         property: impl TemplateProperty<Operation, Output = Operation> + 'static,
@@ -155,7 +182,7 @@ pub type OperationTemplateBuildMethodFnMap<T> =
     TemplateBuildMethodFnMap<'static, OperationTemplateLanguage, T>;
 
 /// Symbol table of methods available in the operation template.
-struct OperationTemplateBuildFnTable {
+pub struct OperationTemplateBuildFnTable {
     pub core: CoreTemplateBuildFnTable<'static, OperationTemplateLanguage>,
     pub operation_methods: OperationTemplateBuildMethodFnMap<Operation>,
     pub operation_id_methods: OperationTemplateBuildMethodFnMap<OperationId>,
@@ -170,6 +197,26 @@ impl OperationTemplateBuildFnTable {
             operation_id_methods: builtin_operation_id_methods(),
         }
     }
+
+    pub fn empty() -> Self {
+        OperationTemplateBuildFnTable {
+            core: CoreTemplateBuildFnTable::empty(),
+            operation_methods: HashMap::new(),
+            operation_id_methods: HashMap::new(),
+        }
+    }
+
+    fn merge(&mut self, other: OperationTemplateBuildFnTable) {
+        let OperationTemplateBuildFnTable {
+            core,
+            operation_methods,
+            operation_id_methods,
+        } = other;
+
+        self.core.merge(core);
+        merge_fn_map(&mut self.operation_methods, operation_methods);
+        merge_fn_map(&mut self.operation_id_methods, operation_id_methods);
+    }
 }
 
 fn builtin_operation_methods() -> OperationTemplateBuildMethodFnMap<Operation> {