revset, templater: implement arity-based alias overloading

Still alias function shadows builtin function (of any arity) by name. This
allows to detect argument error as such, but might be a bit inconvenient if
user wants to overload heads() for example. If needed, maybe we can add some
config/revset syntax to import builtin function to alias namespace.

The functions table is keyed by name, not by (name, arity) pair. That's mainly
because std collections require keys to be Borrow, and a pair of borrowed
values is incompatible with owned pair. Another reason is it makes easy to look
up overloads by name.

Alias overloading could also be achieved by adding default parameters, but that
will complicate the implementation a bit more, and can't prevent shadowing of
0-ary immutable_heads().

Closes #2966

CHANGELOG.md

@@ -31,6 +31,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * The `descendants()` revset function now accepts an optional `depth` argument;
   like the `ancestors()` depth argument, it limits the depth of the set.
 
+* Revset/template aliases now support function overloading.
+  [#2966](https://github.com/martinvonz/jj/issues/2966)
+
 ### Fixed bugs
 
 ## [0.18.0] - 2024-06-05

cli/src/revset_util.rs

@@ -134,17 +134,6 @@ pub fn load_revset_aliases(
             }
         }
     }
-
-    // TODO: If we add support for function overloading (#2966), this check can
-    // be removed.
-    let (_, params, _) = aliases_map.get_function(BUILTIN_IMMUTABLE_HEADS).unwrap();
-    if !params.is_empty() {
-        return Err(user_error(format!(
-            "The `revset-aliases.{name}()` function must be declared without arguments",
-            name = BUILTIN_IMMUTABLE_HEADS
-        )));
-    }
-
     Ok(aliases_map)
 }
 
@@ -175,14 +164,10 @@ pub fn default_symbol_resolver<'a>(
 pub fn parse_immutable_expression(
     context: &RevsetParseContext,
 ) -> Result<Rc<RevsetExpression>, RevsetParseError> {
-    let (_, params, immutable_heads_str) = context
+    let (_, _, immutable_heads_str) = context
         .aliases_map()
-        .get_function(BUILTIN_IMMUTABLE_HEADS)
+        .get_function(BUILTIN_IMMUTABLE_HEADS, 0)
         .unwrap();
-    assert!(
-        params.is_empty(),
-        "invalid declaration should have been rejected by load_revset_aliases()"
-    );
     // Negated ancestors expression `~::(<heads> | root())` is slightly easier
     // to optimize than negated union `~(::<heads> | root())`.
     let heads = revset::parse(immutable_heads_str, context)?;

cli/src/template_parser.rs

@@ -1028,16 +1028,36 @@ mod tests {
     fn test_parse_alias_decl() {
         let mut aliases_map = TemplateAliasesMap::new();
         aliases_map.insert("sym", r#""is symbol""#).unwrap();
-        aliases_map.insert("func(a)", r#""is function""#).unwrap();
+        aliases_map.insert("func()", r#""is function 0""#).unwrap();
+        aliases_map
+            .insert("func(a, b)", r#""is function 2""#)
+            .unwrap();
+        aliases_map.insert("func(a)", r#""is function a""#).unwrap();
+        aliases_map.insert("func(b)", r#""is function b""#).unwrap();
 
         let (id, defn) = aliases_map.get_symbol("sym").unwrap();
         assert_eq!(id, AliasId::Symbol("sym"));
         assert_eq!(defn, r#""is symbol""#);
 
-        let (id, params, defn) = aliases_map.get_function("func").unwrap();
-        assert_eq!(id, AliasId::Function("func", &["a".to_owned()]));
-        assert_eq!(params, ["a"]);
-        assert_eq!(defn, r#""is function""#);
+        let (id, params, defn) = aliases_map.get_function("func", 0).unwrap();
+        assert_eq!(id, AliasId::Function("func", &[]));
+        assert!(params.is_empty());
+        assert_eq!(defn, r#""is function 0""#);
+
+        let (id, params, defn) = aliases_map.get_function("func", 1).unwrap();
+        assert_eq!(id, AliasId::Function("func", &["b".to_owned()]));
+        assert_eq!(params, ["b"]);
+        assert_eq!(defn, r#""is function b""#);
+
+        let (id, params, defn) = aliases_map.get_function("func", 2).unwrap();
+        assert_eq!(
+            id,
+            AliasId::Function("func", &["a".to_owned(), "b".to_owned()])
+        );
+        assert_eq!(params, ["a", "b"]);
+        assert_eq!(defn, r#""is function 2""#);
+
+        assert!(aliases_map.get_function("func", 3).is_none());
 
         // Formal parameter 'a' can't be redefined
         assert_eq!(
@@ -1149,6 +1169,12 @@ mod tests {
             parse_normalized("a ++ b"),
         );
 
+        // Not recursion because functions are overloaded by arity.
+        assert_eq!(
+            with_aliases([("F(x)", "F(x,b)"), ("F(x,y)", "x ++ y")]).parse_normalized("F(a)"),
+            parse_normalized("a ++ b")
+        );
+
         // Arguments should be resolved in the current scope.
         assert_eq!(
             with_aliases([("F(x,y)", "if(x, y)")]).parse_normalized("F(a ++ y, b ++ x)"),
@@ -1221,5 +1247,12 @@ mod tests {
                 .kind,
             TemplateParseErrorKind::BadAliasExpansion("F(x)".to_owned()),
         );
+        assert_eq!(
+            with_aliases([("F(x)", "F(x,b)"), ("F(x,y)", "F(x|y)")])
+                .parse("F(a)")
+                .unwrap_err()
+                .kind,
+            TemplateParseErrorKind::BadAliasExpansion("F(x)".to_owned())
+        );
     }
 }

cli/tests/test_immutable_commits.rs

@@ -95,17 +95,11 @@ fn test_rewrite_immutable_generic() {
     Parent commit      : kkmpptxz c8d4c7ca main | b
     "###);
 
-    // Error if we redefine immutable_heads() with an argument
+    // immutable_heads() of different arity doesn't shadow the 0-ary one
     test_env.add_config(r#"revset-aliases."immutable_heads(foo)" = "none()""#);
     let stderr = test_env.jj_cmd_failure(&repo_path, &["edit", "root()"]);
     insta::assert_snapshot!(stderr, @r###"
-    Error: The `revset-aliases.immutable_heads()` function must be declared without arguments
-    "###);
-    // ... even if we also update the built-in call sites
-    test_env.add_config(r#"revsets.short-prefixes = "immutable_heads(root())""#);
-    let stderr = test_env.jj_cmd_failure(&repo_path, &["edit", "root()"]);
-    insta::assert_snapshot!(stderr, @r###"
-    Error: The `revset-aliases.immutable_heads()` function must be declared without arguments
+    Error: The root commit 000000000000 is immutable
     "###);
 }
 

docs/revsets.md

@@ -196,11 +196,15 @@ Functions that perform string matching support the following pattern syntax:
 New symbols and functions can be defined in the config file, by using any
 combination of the predefined symbols/functions and other aliases.
 
+Alias functions can be overloaded by the number of parameters. However, builtin
+function will be shadowed by name, and can't co-exist with aliases.
+
 For example:
 
 ```toml
 [revset-aliases]
 'HEAD' = '@-'
+'user()' = 'user("me@example.org")
 'user(x)' = 'author(x) | committer(x)'
 ```
 

docs/templates.md

@@ -273,6 +273,9 @@ can be seen in the `cli/src/config/templates.toml` file in jj's source tree.
 New keywords and functions can be defined as aliases, by using any
 combination of the predefined keywords/functions and other aliases.
 
+Alias functions can be overloaded by the number of parameters. However, builtin
+function will be shadowed by name, and can't co-exist with aliases.
+
 For example:
 
 ```toml

lib/src/dsl_util.rs

@@ -64,6 +64,11 @@ pub struct KeywordArgument<'i, T> {
 }
 
 impl<'i, T> FunctionCallNode<'i, T> {
+    /// Number of arguments assuming named arguments are all unique.
+    pub fn arity(&self) -> usize {
+        self.args.len() + self.keyword_args.len()
+    }
+
     /// Ensures that no arguments passed.
     pub fn expect_no_arguments(&self) -> Result<(), InvalidArguments<'i>> {
         let ([], []) = self.expect_arguments()?;
@@ -218,6 +223,14 @@ impl<'i, T> FunctionCallNode<'i, T> {
         };
         self.invalid_arguments(message, self.args_span)
     }
+
+    fn invalid_arguments_count_with_arities(
+        &self,
+        arities: impl IntoIterator<Item = usize>,
+    ) -> InvalidArguments<'i> {
+        let message = format!("Expected {} arguments", arities.into_iter().join(", "));
+        self.invalid_arguments(message, self.args_span)
+    }
 }
 
 /// Unexpected number of arguments, or invalid combination of arguments.
@@ -350,7 +363,8 @@ impl<R: RuleType> StringLiteralParser<R> {
 #[derive(Clone, Debug, Default)]
 pub struct AliasesMap<P> {
     symbol_aliases: HashMap<String, String>,
-    function_aliases: HashMap<String, (Vec<String>, String)>,
+    // name: [(params, defn)] (sorted by arity)
+    function_aliases: HashMap<String, Vec<(Vec<String>, String)>>,
     // Parser type P helps prevent misuse of AliasesMap of different language.
     parser: P,
 }
@@ -377,7 +391,11 @@ impl<P> AliasesMap<P> {
                 self.symbol_aliases.insert(name, defn.into());
             }
             AliasDeclaration::Function(name, params) => {
-                self.function_aliases.insert(name, (params, defn.into()));
+                let overloads = self.function_aliases.entry(name).or_default();
+                match overloads.binary_search_by_key(&params.len(), |(params, _)| params.len()) {
+                    Ok(i) => overloads[i] = (params, defn.into()),
+                    Err(i) => overloads.insert(i, (params, defn.into())),
+                }
             }
         }
         Ok(())
@@ -400,18 +418,49 @@ impl<P> AliasesMap<P> {
             .map(|(name, defn)| (AliasId::Symbol(name), defn.as_ref()))
     }
 
-    /// Looks up function alias by name. Returns identifier, list of parameter
-    /// names, and definition text.
-    pub fn get_function(&self, name: &str) -> Option<(AliasId<'_>, &[String], &str)> {
-        self.function_aliases
-            .get_key_value(name)
-            .map(|(name, (params, defn))| {
-                (
-                    AliasId::Function(name, params),
-                    params.as_ref(),
-                    defn.as_ref(),
-                )
-            })
+    /// Looks up function alias by name and arity. Returns identifier, list of
+    /// parameter names, and definition text.
+    pub fn get_function(&self, name: &str, arity: usize) -> Option<(AliasId<'_>, &[String], &str)> {
+        let overloads = self.get_function_overloads(name)?;
+        overloads.find_by_arity(arity)
+    }
+
+    /// Looks up function aliases by name.
+    fn get_function_overloads(&self, name: &str) -> Option<AliasFunctionOverloads<'_>> {
+        let (name, overloads) = self.function_aliases.get_key_value(name)?;
+        Some(AliasFunctionOverloads { name, overloads })
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+struct AliasFunctionOverloads<'a> {
+    name: &'a String,
+    overloads: &'a Vec<(Vec<String>, String)>,
+}
+
+impl<'a> AliasFunctionOverloads<'a> {
+    fn arities(self) -> impl DoubleEndedIterator<Item = usize> + ExactSizeIterator + 'a {
+        self.overloads.iter().map(|(params, _)| params.len())
+    }
+
+    fn min_arity(self) -> usize {
+        self.arities().next().unwrap()
+    }
+
+    fn max_arity(self) -> usize {
+        self.arities().next_back().unwrap()
+    }
+
+    fn find_by_arity(self, arity: usize) -> Option<(AliasId<'a>, &'a [String], &'a str)> {
+        let index = self
+            .overloads
+            .binary_search_by_key(&arity, |(params, _)| params.len())
+            .ok()?;
+        let (params, defn) = &self.overloads[index];
+        // Exact parameter names aren't needed to identify a function, but they
+        // provide a better error indication. (e.g. "foo(x, y)" is easier to
+        // follow than "foo/2".)
+        Some((AliasId::Function(self.name, params), params, defn))
     }
 }
 
@@ -564,18 +613,23 @@ where
         function: Box<FunctionCallNode<'i, T>>,
         span: pest::Span<'i>,
     ) -> Result<T, Self::Error> {
-        if let Some((id, params, defn)) = self.aliases_map.get_function(function.name) {
-            // TODO: add support for keyword arguments and arity-based
-            // overloading (#2966)?
-            let arity = params.len();
+        // For better error indication, builtin functions are shadowed by name,
+        // not by (name, arity).
+        if let Some(overloads) = self.aliases_map.get_function_overloads(function.name) {
+            // TODO: add support for keyword arguments
             function
                 .ensure_no_keyword_arguments()
                 .map_err(E::invalid_arguments)?;
-            if function.args.len() != arity {
-                return Err(E::invalid_arguments(
-                    function.invalid_arguments_count(arity, Some(arity)),
-                ));
-            }
+            let Some((id, params, defn)) = overloads.find_by_arity(function.arity()) else {
+                let min = overloads.min_arity();
+                let max = overloads.max_arity();
+                let err = if max - min + 1 == overloads.arities().len() {
+                    function.invalid_arguments_count(min, Some(max))
+                } else {
+                    function.invalid_arguments_count_with_arities(overloads.arities())
+                };
+                return Err(E::invalid_arguments(err));
+            };
             // Resolve arguments in the current scope, and pass them in to the alias
             // expansion scope.
             let args = fold_expression_nodes(self, function.args)?;

lib/src/revset_parser.rs

@@ -1274,6 +1274,37 @@ mod tests {
         assert!(aliases_map.insert("a@b", "none()").is_err());
     }
 
+    #[test]
+    fn test_parse_revset_alias_func_decl() {
+        let mut aliases_map = RevsetAliasesMap::new();
+        aliases_map.insert("func()", r#""is function 0""#).unwrap();
+        aliases_map
+            .insert("func(a, b)", r#""is function 2""#)
+            .unwrap();
+        aliases_map.insert("func(a)", r#""is function a""#).unwrap();
+        aliases_map.insert("func(b)", r#""is function b""#).unwrap();
+
+        let (id, params, defn) = aliases_map.get_function("func", 0).unwrap();
+        assert_eq!(id, AliasId::Function("func", &[]));
+        assert!(params.is_empty());
+        assert_eq!(defn, r#""is function 0""#);
+
+        let (id, params, defn) = aliases_map.get_function("func", 1).unwrap();
+        assert_eq!(id, AliasId::Function("func", &["b".to_owned()]));
+        assert_eq!(params, ["b"]);
+        assert_eq!(defn, r#""is function b""#);
+
+        let (id, params, defn) = aliases_map.get_function("func", 2).unwrap();
+        assert_eq!(
+            id,
+            AliasId::Function("func", &["a".to_owned(), "b".to_owned()])
+        );
+        assert_eq!(params, ["a", "b"]);
+        assert_eq!(defn, r#""is function 2""#);
+
+        assert!(aliases_map.get_function("func", 3).is_none());
+    }
+
     #[test]
     fn test_parse_revset_alias_formal_parameter() {
         let mut aliases_map = RevsetAliasesMap::new();
@@ -1542,6 +1573,12 @@ mod tests {
             parse_normalized("a|b")
         );
 
+        // Not recursion because functions are overloaded by arity.
+        assert_eq!(
+            with_aliases([("F(x)", "F(x,b)"), ("F(x,y)", "x|y")]).parse_normalized("F(a)"),
+            parse_normalized("a|b")
+        );
+
         // Arguments should be resolved in the current scope.
         assert_eq!(
             with_aliases([("F(x,y)", "x|y")]).parse_normalized("F(a::y,b::x)"),
@@ -1613,6 +1650,26 @@ mod tests {
                 message: "Expected 2 arguments".to_owned()
             }
         );
+        assert_eq!(
+            with_aliases([("F(x)", "x"), ("F(x,y)", "x|y")])
+                .parse("F()")
+                .unwrap_err()
+                .kind,
+            RevsetParseErrorKind::InvalidFunctionArguments {
+                name: "F".to_owned(),
+                message: "Expected 1 to 2 arguments".to_owned()
+            }
+        );
+        assert_eq!(
+            with_aliases([("F()", "x"), ("F(x,y)", "x|y")])
+                .parse("F(a)")
+                .unwrap_err()
+                .kind,
+            RevsetParseErrorKind::InvalidFunctionArguments {
+                name: "F".to_owned(),
+                message: "Expected 0, 2 arguments".to_owned()
+            }
+        );
 
         // Keyword argument isn't supported for now.
         assert_eq!(
@@ -1634,5 +1691,12 @@ mod tests {
                 .kind,
             RevsetParseErrorKind::BadAliasExpansion("F(x)".to_owned())
         );
+        assert_eq!(
+            with_aliases([("F(x)", "F(x,b)"), ("F(x,y)", "F(x|y)")])
+                .parse("F(a)")
+                .unwrap_err()
+                .kind,
+            RevsetParseErrorKind::BadAliasExpansion("F(x)".to_owned())
+        );
     }
 }