revset: support defining custom revset functions

cli/src/cli_util.rs

@@ -54,8 +54,9 @@ use jj_lib::repo::{
 };
 use jj_lib::repo_path::{FsPathParseError, RepoPath, RepoPathBuf};
 use jj_lib::revset::{
-    RevsetAliasesMap, RevsetExpression, RevsetExtensions, RevsetFilterPredicate, RevsetIteratorExt,
-    RevsetModifier, RevsetParseContext, RevsetWorkspaceContext, SymbolResolverExtension,
+    RevsetAliasesMap, RevsetExpression, RevsetExtensions, RevsetFilterPredicate, RevsetFunction,
+    RevsetIteratorExt, RevsetModifier, RevsetParseContext, RevsetWorkspaceContext,
+    SymbolResolverExtension,
 };
 use jj_lib::rewrite::restore_tree;
 use jj_lib::settings::{ConfigResultExt as _, UserSettings};
@@ -2747,6 +2748,15 @@ impl CliRunner {
         self
     }
 
+    pub fn add_revset_function_extension(
+        mut self,
+        name: &'static str,
+        func: RevsetFunction,
+    ) -> Self {
+        self.revset_extensions.add_custom_function(name, func);
+        self
+    }
+
     pub fn add_commit_template_extension(
         mut self,
         commit_template_extension: Box<dyn CommitTemplateLanguageExtension>,

cli/src/commit_templater.rs

@@ -489,7 +489,7 @@ fn builtin_commit_methods<'repo>() -> CommitTemplateBuildMethodFnMap<'repo, Comm
     );
     map.insert("mine", |language, _build_ctx, self_property, function| {
         template_parser::expect_no_arguments(function)?;
-        let user_email = language.revset_parse_context.user_email().clone();
+        let user_email = language.revset_parse_context.user_email().to_owned();
         let out_property = self_property.map(move |commit| commit.author().email == user_email);
         Ok(L::wrap_boolean(out_property))
     });

lib/src/revset.rs

@@ -14,7 +14,7 @@
 
 #![allow(missing_docs)]
 
-use std::collections::{HashMap, HashSet};
+use std::collections::{hash_map, HashMap, HashSet};
 use std::convert::Infallible;
 use std::ops::Range;
 use std::path::Path;
@@ -806,6 +806,7 @@ impl fmt::Display for RevsetAliasId<'_> {
 
 #[derive(Clone, Copy, Debug)]
 pub struct ParseState<'a> {
+    function_map: &'a HashMap<&'static str, &'a RevsetFunction>,
     aliases_map: &'a RevsetAliasesMap,
     aliases_expanding: &'a [RevsetAliasId<'a>],
     locals: &'a HashMap<&'a str, Rc<RevsetExpression>>,
@@ -821,6 +822,7 @@ impl<'a> ParseState<'a> {
         locals: &'a HashMap<&str, Rc<RevsetExpression>>,
     ) -> Self {
         ParseState {
+            function_map: &context.function_map,
             aliases_map: context.aliases_map,
             aliases_expanding: &[],
             locals,
@@ -847,6 +849,7 @@ impl<'a> ParseState<'a> {
         let mut aliases_expanding = self.aliases_expanding.to_vec();
         aliases_expanding.push(id);
         let expanding_state = ParseState {
+            function_map: self.function_map,
             aliases_map: self.aliases_map,
             aliases_expanding: &aliases_expanding,
             locals,
@@ -1152,22 +1155,28 @@ fn parse_function_expression(
         state.with_alias_expanding(id, &locals, primary_span, |state| {
             parse_program(defn, state)
         })
-    } else if let Some(func) = BUILTIN_FUNCTION_MAP.get(name) {
+    } else if let Some(func) = state.function_map.get(name) {
         func(name, arguments_pair, state)
     } else {
         Err(RevsetParseError::with_span(
             RevsetParseErrorKind::NoSuchFunction {
                 name: name.to_owned(),
-                candidates: collect_similar(name, all_function_names(state.aliases_map)),
+                candidates: collect_similar(
+                    name,
+                    all_function_names(state.function_map.keys().copied(), state.aliases_map),
+                ),
             },
             name_pair.as_span(),
         ))
     }
 }
 
-fn all_function_names(aliases_map: &RevsetAliasesMap) -> impl Iterator<Item = &str> {
+fn all_function_names<'a>(
+    function_keys: impl Iterator<Item = &'a str>,
+    aliases_map: &'a RevsetAliasesMap,
+) -> impl Iterator<Item = &'a str> {
     itertools::chain(
-        BUILTIN_FUNCTION_MAP.keys().copied(),
+        function_keys,
         aliases_map.function_aliases.keys().map(|n| n.as_ref()),
     )
 }
@@ -2640,7 +2649,7 @@ impl Iterator for ReverseRevsetIterator {
 #[derive(Default)]
 pub struct RevsetExtensions {
     symbol_resolvers: Vec<Box<dyn SymbolResolverExtension>>,
-    // TODO: Add more fields for extending the revset language
+    custom_functions: HashMap<&'static str, RevsetFunction>,
 }
 
 impl RevsetExtensions {
@@ -2651,6 +2660,15 @@ impl RevsetExtensions {
     pub fn add_symbol_resolver(&mut self, symbol_resolver: Box<dyn SymbolResolverExtension>) {
         self.symbol_resolvers.push(symbol_resolver);
     }
+
+    pub fn add_custom_function(&mut self, name: &'static str, func: RevsetFunction) {
+        match self.custom_functions.entry(name) {
+            hash_map::Entry::Occupied(_) => {
+                panic!("Multiple extensions tried to register revset function '{name}'")
+            }
+            hash_map::Entry::Vacant(v) => v.insert(func),
+        };
+    }
 }
 
 /// Information needed to parse revset expression.
@@ -2660,6 +2678,7 @@ pub struct RevsetParseContext<'a> {
     user_email: String,
     extensions: &'a RevsetExtensions,
     workspace: Option<RevsetWorkspaceContext<'a>>,
+    function_map: HashMap<&'static str, &'a RevsetFunction>,
 }
 
 impl<'a> RevsetParseContext<'a> {
@@ -2669,11 +2688,23 @@ impl<'a> RevsetParseContext<'a> {
         extensions: &'a RevsetExtensions,
         workspace: Option<RevsetWorkspaceContext<'a>>,
     ) -> Self {
+        let mut function_map = HashMap::<&'static str, &'a RevsetFunction>::new();
+        for (name, func) in BUILTIN_FUNCTION_MAP.iter() {
+            function_map.insert(name, func);
+        }
+        for (name, func) in &extensions.custom_functions {
+            match function_map.entry(name) {
+                hash_map::Entry::Occupied(_) => panic!("Cannot override builtin function '{name}'"),
+                hash_map::Entry::Vacant(v) => v.insert(func),
+            };
+        }
+
         Self {
             aliases_map,
             user_email,
             extensions,
             workspace,
+            function_map,
         }
     }
 
@@ -2681,7 +2712,7 @@ impl<'a> RevsetParseContext<'a> {
         self.aliases_map
     }
 
-    pub fn user_email(&self) -> &String {
+    pub fn user_email(&self) -> &str {
         &self.user_email
     }