Skip to main content

jj_lib/
dsl_util.rs

1// Copyright 2020-2024 The Jujutsu Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Domain-specific language helpers.
16
17use std::collections::HashMap;
18use std::fmt;
19use std::slice;
20
21use itertools::Itertools as _;
22use pest::RuleType;
23use pest::iterators::Pair;
24use pest::iterators::Pairs;
25
26pub use crate::symbol_util::escape_string;
27pub use crate::symbol_util::format_string;
28use crate::symbol_util::unescape_char;
29
30/// Manages diagnostic messages emitted during parsing.
31///
32/// `T` is usually a parse error type of the language, which contains a message
33/// and source span of 'static lifetime.
34#[derive(Debug)]
35pub struct Diagnostics<T> {
36    // This might be extended to [{ kind: Warning|Error, message: T }, ..].
37    diagnostics: Vec<T>,
38}
39
40impl<T> Diagnostics<T> {
41    /// Creates new empty diagnostics collector.
42    pub fn new() -> Self {
43        Self {
44            diagnostics: Vec::new(),
45        }
46    }
47
48    /// Returns `true` if there are no diagnostic messages.
49    pub fn is_empty(&self) -> bool {
50        self.diagnostics.is_empty()
51    }
52
53    /// Returns the number of diagnostic messages.
54    pub fn len(&self) -> usize {
55        self.diagnostics.len()
56    }
57
58    /// Returns iterator over diagnostic messages.
59    pub fn iter(&self) -> slice::Iter<'_, T> {
60        self.diagnostics.iter()
61    }
62
63    /// Adds a diagnostic message of warning level.
64    pub fn add_warning(&mut self, diag: T) {
65        self.diagnostics.push(diag);
66    }
67
68    /// Moves diagnostic messages of different type (such as fileset warnings
69    /// emitted within `file()` revset.)
70    pub fn extend_with<U>(&mut self, diagnostics: Diagnostics<U>, mut f: impl FnMut(U) -> T) {
71        self.diagnostics
72            .extend(diagnostics.diagnostics.into_iter().map(&mut f));
73    }
74}
75
76impl<T> Default for Diagnostics<T> {
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82impl<'a, T> IntoIterator for &'a Diagnostics<T> {
83    type Item = &'a T;
84    type IntoIter = slice::Iter<'a, T>;
85
86    fn into_iter(self) -> Self::IntoIter {
87        self.iter()
88    }
89}
90
91/// AST node without type or name checking.
92#[derive(Clone, Debug, Eq, PartialEq)]
93pub struct ExpressionNode<'i, T> {
94    /// Expression item such as identifier, literal, function call, etc.
95    pub kind: T,
96    /// Span of the node.
97    pub span: pest::Span<'i>,
98}
99
100impl<'i, T> ExpressionNode<'i, T> {
101    /// Wraps the given expression and span.
102    pub fn new(kind: T, span: pest::Span<'i>) -> Self {
103        Self { kind, span }
104    }
105}
106
107/// `<name>:<value>` expression in AST.
108#[derive(Clone, Debug, Eq, PartialEq)]
109pub struct PatternNode<'i, T> {
110    /// Pattern name or type (such as `glob`.)
111    pub name: &'i str,
112    /// Span of the pattern name.
113    pub name_span: pest::Span<'i>,
114    /// Value expression.
115    pub value: ExpressionNode<'i, T>,
116}
117
118/// Function call in AST.
119#[derive(Clone, Debug, Eq, PartialEq)]
120pub struct FunctionCallNode<'i, T> {
121    /// Function name.
122    pub name: &'i str,
123    /// Span of the function name.
124    pub name_span: pest::Span<'i>,
125    /// List of positional arguments.
126    pub args: Vec<ExpressionNode<'i, T>>,
127    /// List of keyword arguments.
128    pub keyword_args: Vec<KeywordArgument<'i, T>>,
129    /// Span of the arguments list.
130    pub args_span: pest::Span<'i>,
131}
132
133/// Keyword argument pair in AST.
134#[derive(Clone, Debug, Eq, PartialEq)]
135pub struct KeywordArgument<'i, T> {
136    /// Parameter name.
137    pub name: &'i str,
138    /// Span of the parameter name.
139    pub name_span: pest::Span<'i>,
140    /// Value expression.
141    pub value: ExpressionNode<'i, T>,
142}
143
144impl<'i, T> FunctionCallNode<'i, T> {
145    /// Number of arguments assuming named arguments are all unique.
146    pub fn arity(&self) -> usize {
147        self.args.len() + self.keyword_args.len()
148    }
149
150    /// Ensures that no arguments passed.
151    pub fn expect_no_arguments(&self) -> Result<(), InvalidArguments<'i>> {
152        let ([], []) = self.expect_arguments()?;
153        Ok(())
154    }
155
156    /// Extracts exactly N required arguments.
157    pub fn expect_exact_arguments<const N: usize>(
158        &self,
159    ) -> Result<&[ExpressionNode<'i, T>; N], InvalidArguments<'i>> {
160        let (args, []) = self.expect_arguments()?;
161        Ok(args)
162    }
163
164    /// Extracts N required arguments and remainders.
165    ///
166    /// This can be used to get all the positional arguments without requiring
167    /// any (N = 0):
168    /// ```ignore
169    /// let ([], content_nodes) = function.expect_some_arguments()?;
170    /// ```
171    /// Avoid accessing `function.args` directly, as that may allow keyword
172    /// arguments to be silently ignored.
173    #[expect(clippy::type_complexity)]
174    pub fn expect_some_arguments<const N: usize>(
175        &self,
176    ) -> Result<(&[ExpressionNode<'i, T>; N], &[ExpressionNode<'i, T>]), InvalidArguments<'i>> {
177        self.ensure_no_keyword_arguments()?;
178        if self.args.len() >= N {
179            let (required, rest) = self.args.split_at(N);
180            Ok((required.try_into().unwrap(), rest))
181        } else {
182            Err(self.invalid_arguments_count(N, None))
183        }
184    }
185
186    /// Extracts N required arguments and M optional arguments.
187    #[expect(clippy::type_complexity)]
188    pub fn expect_arguments<const N: usize, const M: usize>(
189        &self,
190    ) -> Result<
191        (
192            &[ExpressionNode<'i, T>; N],
193            [Option<&ExpressionNode<'i, T>>; M],
194        ),
195        InvalidArguments<'i>,
196    > {
197        self.ensure_no_keyword_arguments()?;
198        let count_range = N..=(N + M);
199        if count_range.contains(&self.args.len()) {
200            let (required, rest) = self.args.split_at(N);
201            let mut optional = rest.iter().map(Some).collect_vec();
202            optional.resize(M, None);
203            Ok((
204                required.try_into().unwrap(),
205                optional.try_into().ok().unwrap(),
206            ))
207        } else {
208            let (min, max) = count_range.into_inner();
209            Err(self.invalid_arguments_count(min, Some(max)))
210        }
211    }
212
213    /// Extracts N required arguments and M optional arguments. Some of them can
214    /// be specified as keyword arguments.
215    ///
216    /// `names` is a list of parameter names. Unnamed positional arguments
217    /// should be padded with `""`.
218    #[expect(clippy::type_complexity)]
219    pub fn expect_named_arguments<const N: usize, const M: usize>(
220        &self,
221        names: &[&str],
222    ) -> Result<
223        (
224            [&ExpressionNode<'i, T>; N],
225            [Option<&ExpressionNode<'i, T>>; M],
226        ),
227        InvalidArguments<'i>,
228    > {
229        if self.keyword_args.is_empty() {
230            let (required, optional) = self.expect_arguments::<N, M>()?;
231            Ok((required.each_ref(), optional))
232        } else {
233            let (required, optional) = self.expect_named_arguments_vec(names, N, N + M)?;
234            Ok((
235                required.try_into().ok().unwrap(),
236                optional.try_into().ok().unwrap(),
237            ))
238        }
239    }
240
241    #[expect(clippy::type_complexity)]
242    fn expect_named_arguments_vec(
243        &self,
244        names: &[&str],
245        min: usize,
246        max: usize,
247    ) -> Result<
248        (
249            Vec<&ExpressionNode<'i, T>>,
250            Vec<Option<&ExpressionNode<'i, T>>>,
251        ),
252        InvalidArguments<'i>,
253    > {
254        assert!(names.len() <= max);
255
256        if self.args.len() > max {
257            return Err(self.invalid_arguments_count(min, Some(max)));
258        }
259        let mut extracted = Vec::with_capacity(max);
260        extracted.extend(self.args.iter().map(Some));
261        extracted.resize(max, None);
262
263        for arg in &self.keyword_args {
264            let name = arg.name;
265            let span = arg.name_span.start_pos().span(&arg.value.span.end_pos());
266            let pos = names.iter().position(|&n| n == name).ok_or_else(|| {
267                self.invalid_arguments(format!(r#"Unexpected keyword argument "{name}""#), span)
268            })?;
269            if extracted[pos].is_some() {
270                return Err(self.invalid_arguments(
271                    format!(r#"Got multiple values for keyword "{name}""#),
272                    span,
273                ));
274            }
275            extracted[pos] = Some(&arg.value);
276        }
277
278        let optional = extracted.split_off(min);
279        let required = extracted.into_iter().flatten().collect_vec();
280        if required.len() != min {
281            return Err(self.invalid_arguments_count(min, Some(max)));
282        }
283        Ok((required, optional))
284    }
285
286    fn ensure_no_keyword_arguments(&self) -> Result<(), InvalidArguments<'i>> {
287        if let (Some(first), Some(last)) = (self.keyword_args.first(), self.keyword_args.last()) {
288            let span = first.name_span.start_pos().span(&last.value.span.end_pos());
289            Err(self.invalid_arguments("Unexpected keyword arguments".to_owned(), span))
290        } else {
291            Ok(())
292        }
293    }
294
295    fn invalid_arguments(&self, message: String, span: pest::Span<'i>) -> InvalidArguments<'i> {
296        InvalidArguments {
297            name: self.name,
298            message,
299            span,
300        }
301    }
302
303    fn invalid_arguments_count(&self, min: usize, max: Option<usize>) -> InvalidArguments<'i> {
304        let message = match (min, max) {
305            (min, Some(max)) if min == max => format!("Expected {min} arguments"),
306            (min, Some(max)) => format!("Expected {min} to {max} arguments"),
307            (min, None) => format!("Expected at least {min} arguments"),
308        };
309        self.invalid_arguments(message, self.args_span)
310    }
311
312    fn invalid_arguments_count_with_arities(
313        &self,
314        arities: impl IntoIterator<Item = usize>,
315    ) -> InvalidArguments<'i> {
316        let message = format!("Expected {} arguments", arities.into_iter().join(", "));
317        self.invalid_arguments(message, self.args_span)
318    }
319}
320
321/// Unexpected number of arguments, or invalid combination of arguments.
322///
323/// This error is supposed to be converted to language-specific parse error
324/// type, where lifetime `'i` will be eliminated.
325#[derive(Clone, Debug)]
326pub struct InvalidArguments<'i> {
327    /// Function name.
328    pub name: &'i str,
329    /// Error message.
330    pub message: String,
331    /// Span of the bad arguments.
332    pub span: pest::Span<'i>,
333}
334
335/// Expression item that can be transformed recursively by using `folder: F`.
336pub trait FoldableExpression<'i>: Sized {
337    /// Transforms `self` by applying the `folder` to inner items.
338    fn fold<F>(self, folder: &mut F, span: pest::Span<'i>) -> Result<Self, F::Error>
339    where
340        F: ExpressionFolder<'i, Self> + ?Sized;
341}
342
343/// Visitor-like interface to transform AST nodes recursively.
344pub trait ExpressionFolder<'i, T: FoldableExpression<'i>> {
345    /// Transform error.
346    type Error;
347
348    /// Transforms the expression `node`. By default, inner items are
349    /// transformed recursively.
350    fn fold_expression(
351        &mut self,
352        node: ExpressionNode<'i, T>,
353    ) -> Result<ExpressionNode<'i, T>, Self::Error> {
354        let ExpressionNode { kind, span } = node;
355        let kind = kind.fold(self, span)?;
356        Ok(ExpressionNode { kind, span })
357    }
358
359    /// Transforms identifier.
360    fn fold_identifier(&mut self, name: &'i str, span: pest::Span<'i>) -> Result<T, Self::Error>;
361
362    /// Transforms pattern.
363    fn fold_pattern(
364        &mut self,
365        pattern: Box<PatternNode<'i, T>>,
366        span: pest::Span<'i>,
367    ) -> Result<T, Self::Error>;
368
369    /// Transforms function call.
370    fn fold_function_call(
371        &mut self,
372        function: Box<FunctionCallNode<'i, T>>,
373        span: pest::Span<'i>,
374    ) -> Result<T, Self::Error>;
375}
376
377/// Transforms list of `nodes` by using `folder`.
378pub fn fold_expression_nodes<'i, F, T>(
379    folder: &mut F,
380    nodes: Vec<ExpressionNode<'i, T>>,
381) -> Result<Vec<ExpressionNode<'i, T>>, F::Error>
382where
383    F: ExpressionFolder<'i, T> + ?Sized,
384    T: FoldableExpression<'i>,
385{
386    nodes
387        .into_iter()
388        .map(|node| folder.fold_expression(node))
389        .try_collect()
390}
391
392/// Transforms pattern value by using `folder`.
393pub fn fold_pattern_value<'i, F, T>(
394    folder: &mut F,
395    pattern: PatternNode<'i, T>,
396) -> Result<PatternNode<'i, T>, F::Error>
397where
398    F: ExpressionFolder<'i, T> + ?Sized,
399    T: FoldableExpression<'i>,
400{
401    Ok(PatternNode {
402        name: pattern.name,
403        name_span: pattern.name_span,
404        value: folder.fold_expression(pattern.value)?,
405    })
406}
407
408/// Transforms function call arguments by using `folder`.
409pub fn fold_function_call_args<'i, F, T>(
410    folder: &mut F,
411    function: FunctionCallNode<'i, T>,
412) -> Result<FunctionCallNode<'i, T>, F::Error>
413where
414    F: ExpressionFolder<'i, T> + ?Sized,
415    T: FoldableExpression<'i>,
416{
417    Ok(FunctionCallNode {
418        name: function.name,
419        name_span: function.name_span,
420        args: fold_expression_nodes(folder, function.args)?,
421        keyword_args: function
422            .keyword_args
423            .into_iter()
424            .map(|arg| {
425                Ok(KeywordArgument {
426                    name: arg.name,
427                    name_span: arg.name_span,
428                    value: folder.fold_expression(arg.value)?,
429                })
430            })
431            .try_collect()?,
432        args_span: function.args_span,
433    })
434}
435
436/// Helper to parse string literal.
437#[derive(Debug)]
438pub struct StringLiteralParser<R> {
439    /// String content part.
440    pub content_rule: R,
441    /// Escape sequence part including backslash character.
442    pub escape_rule: R,
443}
444
445impl<R: RuleType> StringLiteralParser<R> {
446    /// Parses the given string literal `pairs` into string.
447    pub fn parse(&self, pairs: Pairs<R>) -> String {
448        let mut result = String::new();
449        for part in pairs {
450            if part.as_rule() == self.content_rule {
451                result.push_str(part.as_str());
452            } else if part.as_rule() == self.escape_rule {
453                result.push(unescape_char(part.as_str()));
454            } else {
455                panic!("unexpected part of string: {part:?}");
456            }
457        }
458        result
459    }
460}
461
462/// Helper to parse function call.
463#[derive(Debug)]
464pub struct FunctionCallParser<R> {
465    /// Function name.
466    pub function_name_rule: R,
467    /// List of positional and keyword arguments.
468    pub function_arguments_rule: R,
469    /// Pair of parameter name and value.
470    pub keyword_argument_rule: R,
471    /// Parameter name.
472    pub argument_name_rule: R,
473    /// Value expression.
474    pub argument_value_rule: R,
475}
476
477impl<R: RuleType> FunctionCallParser<R> {
478    /// Parses the given `pair` as function call.
479    pub fn parse<'i, T, E: From<InvalidArguments<'i>>>(
480        &self,
481        pair: Pair<'i, R>,
482        // parse_name can be defined for any Pair<'_, R>, but parse_value should
483        // be allowed to construct T by capturing Pair<'i, R>.
484        parse_name: impl Fn(Pair<'i, R>) -> Result<&'i str, E>,
485        parse_value: impl Fn(Pair<'i, R>) -> Result<ExpressionNode<'i, T>, E>,
486    ) -> Result<FunctionCallNode<'i, T>, E> {
487        let [name_pair, args_pair] = pair.into_inner().collect_array().unwrap();
488        assert_eq!(name_pair.as_rule(), self.function_name_rule);
489        assert_eq!(args_pair.as_rule(), self.function_arguments_rule);
490        let name_span = name_pair.as_span();
491        let args_span = args_pair.as_span();
492        let function_name = parse_name(name_pair)?;
493        let mut args = Vec::new();
494        let mut keyword_args = Vec::new();
495        for pair in args_pair.into_inner() {
496            let span = pair.as_span();
497            if pair.as_rule() == self.argument_value_rule {
498                if !keyword_args.is_empty() {
499                    return Err(InvalidArguments {
500                        name: function_name,
501                        message: "Positional argument follows keyword argument".to_owned(),
502                        span,
503                    }
504                    .into());
505                }
506                args.push(parse_value(pair)?);
507            } else if pair.as_rule() == self.keyword_argument_rule {
508                let [name_pair, value_pair] = pair.into_inner().collect_array().unwrap();
509                assert_eq!(name_pair.as_rule(), self.argument_name_rule);
510                assert_eq!(value_pair.as_rule(), self.argument_value_rule);
511                let name_span = name_pair.as_span();
512                let arg = KeywordArgument {
513                    name: parse_name(name_pair)?,
514                    name_span,
515                    value: parse_value(value_pair)?,
516                };
517                keyword_args.push(arg);
518            } else {
519                panic!("unexpected argument rule {pair:?}");
520            }
521        }
522        Ok(FunctionCallNode {
523            name: function_name,
524            name_span,
525            args,
526            keyword_args,
527            args_span,
528        })
529    }
530}
531
532/// A function alias containing `(params, definition, description)`.
533type FunctionAlias<V> = (Vec<String>, V, Option<String>);
534
535/// Map of symbol, pattern, and function aliases.
536#[derive(Clone, Debug, Default)]
537pub struct AliasesMap<P, V> {
538    symbol_aliases: HashMap<String, (V, Option<String>)>,
539    // name: (param, defn)
540    pattern_aliases: HashMap<String, (String, V, Option<String>)>,
541    // name: [(params, defn)] (sorted by arity)
542    function_aliases: HashMap<String, Vec<FunctionAlias<V>>>,
543    // Parser type P helps prevent misuse of AliasesMap of different language.
544    parser: P,
545}
546
547impl<P, V> AliasesMap<P, V> {
548    /// Creates an empty aliases map with default-constructed parser.
549    pub fn new() -> Self
550    where
551        P: Default,
552    {
553        Self {
554            symbol_aliases: Default::default(),
555            pattern_aliases: Default::default(),
556            function_aliases: Default::default(),
557            parser: Default::default(),
558        }
559    }
560
561    /// Adds new substitution rule `decl = defn`.
562    ///
563    /// Returns error if `decl` is invalid. The `defn` part isn't checked. A bad
564    /// `defn` will be reported when the alias is substituted.
565    pub fn insert(
566        &mut self,
567        decl: impl AsRef<str>,
568        defn: impl Into<V>,
569        doc: Option<String>,
570    ) -> Result<(), P::Error>
571    where
572        P: AliasDeclarationParser,
573    {
574        match self.parser.parse_declaration(decl.as_ref())? {
575            AliasDeclaration::Symbol(name) => {
576                self.symbol_aliases.insert(name, (defn.into(), doc));
577            }
578            AliasDeclaration::Pattern(name, param) => {
579                self.pattern_aliases.insert(name, (param, defn.into(), doc));
580            }
581            AliasDeclaration::Function(name, params) => {
582                let overloads = self.function_aliases.entry(name).or_default();
583                match overloads.binary_search_by_key(&params.len(), |(params, _, _)| params.len()) {
584                    Ok(i) => overloads[i] = (params, defn.into(), doc),
585                    Err(i) => overloads.insert(i, (params, defn.into(), doc)),
586                }
587            }
588        }
589        Ok(())
590    }
591
592    /// Iterates symbol names in arbitrary order.
593    pub fn symbol_names(&self) -> impl Iterator<Item = &str> {
594        self.symbol_aliases.keys().map(|n| n.as_ref())
595    }
596
597    /// Iterates pattern names in arbitrary order.
598    pub fn pattern_names(&self) -> impl Iterator<Item = &str> {
599        self.pattern_aliases.keys().map(|n| n.as_ref())
600    }
601
602    /// Iterates function names in arbitrary order.
603    pub fn function_names(&self) -> impl Iterator<Item = &str> {
604        self.function_aliases.keys().map(|n| n.as_ref())
605    }
606
607    /// Looks up symbol alias by name. Returns identifier, definition text, and
608    /// optional description.
609    pub fn get_symbol(&self, name: &str) -> Option<(AliasId<'_>, &V, Option<&str>)> {
610        self.symbol_aliases
611            .get_key_value(name)
612            .map(|(name, (defn, doc))| (AliasId::Symbol(name), defn, doc.as_deref()))
613    }
614
615    /// Looks up pattern alias by name. Returns identifier, parameter name,
616    /// definition text, and optional description.
617    pub fn get_pattern(&self, name: &str) -> Option<(AliasId<'_>, &str, &V, Option<&str>)> {
618        self.pattern_aliases
619            .get_key_value(name)
620            .map(|(name, (param, defn, doc))| {
621                (
622                    AliasId::Pattern(name, param),
623                    param.as_ref(),
624                    defn,
625                    doc.as_deref(),
626                )
627            })
628    }
629
630    /// Looks up function alias by name and arity. Returns identifier, list of
631    /// parameter names, definition text, and optional description.
632    pub fn get_function(
633        &self,
634        name: &str,
635        arity: usize,
636    ) -> Option<(AliasId<'_>, &[String], &V, Option<&str>)> {
637        let overloads = self.get_function_overloads(name)?;
638        overloads.find_by_arity(arity)
639    }
640
641    /// Looks up function aliases by name.
642    fn get_function_overloads(&self, name: &str) -> Option<AliasFunctionOverloads<'_, V>> {
643        let (name, overloads) = self.function_aliases.get_key_value(name)?;
644        Some(AliasFunctionOverloads { name, overloads })
645    }
646}
647
648#[derive(Clone, Debug)]
649struct AliasFunctionOverloads<'a, V> {
650    name: &'a String,
651    overloads: &'a Vec<(Vec<String>, V, Option<String>)>,
652}
653
654impl<'a, V> AliasFunctionOverloads<'a, V> {
655    fn arities(&self) -> impl DoubleEndedIterator<Item = usize> + ExactSizeIterator {
656        self.overloads.iter().map(|(params, _, _)| params.len())
657    }
658
659    fn min_arity(&self) -> usize {
660        self.arities().next().unwrap()
661    }
662
663    fn max_arity(&self) -> usize {
664        self.arities().next_back().unwrap()
665    }
666
667    fn find_by_arity(
668        &self,
669        arity: usize,
670    ) -> Option<(AliasId<'a>, &'a [String], &'a V, Option<&'a str>)> {
671        let index = self
672            .overloads
673            .binary_search_by_key(&arity, |(params, _, _)| params.len())
674            .ok()?;
675        let (params, defn, doc) = &self.overloads[index];
676        // Exact parameter names aren't needed to identify a function, but they
677        // provide a better error indication. (e.g. "foo(x, y)" is easier to
678        // follow than "foo/2".)
679        Some((
680            AliasId::Function(self.name, params),
681            params,
682            defn,
683            doc.as_deref(),
684        ))
685    }
686}
687
688/// Borrowed reference to identify alias expression.
689#[derive(Clone, Copy, Debug, Eq, PartialEq)]
690pub enum AliasId<'a> {
691    /// Symbol name.
692    Symbol(&'a str),
693    /// Pattern name and parameter name.
694    Pattern(&'a str, &'a str),
695    /// Function name and parameter names.
696    Function(&'a str, &'a [String]),
697    /// Function parameter name.
698    Parameter(&'a str),
699}
700
701impl fmt::Display for AliasId<'_> {
702    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
703        match self {
704            Self::Symbol(name) => write!(f, "{name}"),
705            Self::Pattern(name, param) => write!(f, "{name}:{param}"),
706            Self::Function(name, params) => {
707                write!(f, "{name}({params})", params = params.join(", "))
708            }
709            Self::Parameter(name) => write!(f, "{name}"),
710        }
711    }
712}
713
714/// Parsed declaration part of alias rule.
715#[derive(Clone, Debug)]
716pub enum AliasDeclaration {
717    /// Symbol name.
718    Symbol(String),
719    /// Pattern name and parameter.
720    Pattern(String, String),
721    /// Function name and parameters.
722    Function(String, Vec<String>),
723}
724
725// AliasDeclarationParser and AliasDefinitionParser can be merged into a single
726// trait, but it's unclear whether doing that would simplify the abstraction.
727
728/// Parser for symbol and function alias declaration.
729pub trait AliasDeclarationParser {
730    /// Parse error type.
731    type Error;
732
733    /// Parses symbol or function name and parameters.
734    fn parse_declaration(&self, source: &str) -> Result<AliasDeclaration, Self::Error>;
735}
736
737/// Parser for symbol and function alias definition.
738pub trait AliasDefinitionParser {
739    /// Expression item type.
740    type Output<'i>;
741    /// Parse error type.
742    type Error;
743
744    /// Parses alias body.
745    fn parse_definition<'i>(
746        &self,
747        source: &'i str,
748    ) -> Result<ExpressionNode<'i, Self::Output<'i>>, Self::Error>;
749}
750
751/// Expression item that supports alias substitution.
752pub trait AliasExpandableExpression<'i>: FoldableExpression<'i> {
753    /// Wraps identifier.
754    fn identifier(name: &'i str) -> Self;
755    /// Wraps pattern.
756    fn pattern(pattern: Box<PatternNode<'i, Self>>) -> Self;
757    /// Wraps function call.
758    fn function_call(function: Box<FunctionCallNode<'i, Self>>) -> Self;
759    /// Wraps substituted expression.
760    fn alias_expanded(id: AliasId<'i>, subst: Box<ExpressionNode<'i, Self>>) -> Self;
761}
762
763/// Error that may occur during alias substitution.
764pub trait AliasExpandError: Sized {
765    /// Unexpected number of arguments, or invalid combination of arguments.
766    fn invalid_arguments(err: InvalidArguments<'_>) -> Self;
767    /// Recursion detected during alias substitution.
768    fn recursive_expansion(id: AliasId<'_>, span: pest::Span<'_>) -> Self;
769    /// Attaches alias trace to the current error.
770    fn within_alias_expansion(self, id: AliasId<'_>, span: pest::Span<'_>) -> Self;
771}
772
773/// Expands aliases recursively in tree of `T`.
774#[derive(Debug)]
775struct AliasExpander<'i, 'a, T, P> {
776    /// Alias symbols and functions that are globally available.
777    aliases_map: &'i AliasesMap<P, String>,
778    /// Local variables set in the outermost scope.
779    locals: &'a HashMap<&'i str, ExpressionNode<'i, T>>,
780    /// Stack of aliases and local parameters currently expanding.
781    states: Vec<AliasExpandingState<'i, T>>,
782}
783
784#[derive(Debug)]
785struct AliasExpandingState<'i, T> {
786    id: AliasId<'i>,
787    locals: HashMap<&'i str, ExpressionNode<'i, T>>,
788}
789
790impl<'i, T, P, E> AliasExpander<'i, '_, T, P>
791where
792    T: AliasExpandableExpression<'i> + Clone,
793    P: AliasDefinitionParser<Output<'i> = T, Error = E>,
794    E: AliasExpandError,
795{
796    /// Local variables available to the current scope.
797    fn current_locals(&self) -> &HashMap<&'i str, ExpressionNode<'i, T>> {
798        self.states.last().map_or(self.locals, |s| &s.locals)
799    }
800
801    fn expand_defn(
802        &mut self,
803        id: AliasId<'i>,
804        defn: &'i str,
805        locals: HashMap<&'i str, ExpressionNode<'i, T>>,
806        span: pest::Span<'i>,
807    ) -> Result<T, E> {
808        // The stack should be short, so let's simply do linear search.
809        if self.states.iter().any(|s| s.id == id) {
810            return Err(E::recursive_expansion(id, span));
811        }
812        self.states.push(AliasExpandingState { id, locals });
813        // Parsed defn could be cached if needed.
814        let result = self
815            .aliases_map
816            .parser
817            .parse_definition(defn)
818            .and_then(|node| self.fold_expression(node))
819            .map(|node| T::alias_expanded(id, Box::new(node)))
820            .map_err(|e| e.within_alias_expansion(id, span));
821        self.states.pop();
822        result
823    }
824}
825
826impl<'i, T, P, E> ExpressionFolder<'i, T> for AliasExpander<'i, '_, T, P>
827where
828    T: AliasExpandableExpression<'i> + Clone,
829    P: AliasDefinitionParser<Output<'i> = T, Error = E>,
830    E: AliasExpandError,
831{
832    type Error = E;
833
834    fn fold_identifier(&mut self, name: &'i str, span: pest::Span<'i>) -> Result<T, Self::Error> {
835        if let Some(subst) = self.current_locals().get(name) {
836            let id = AliasId::Parameter(name);
837            Ok(T::alias_expanded(id, Box::new(subst.clone())))
838        } else if let Some((id, defn, _doc)) = self.aliases_map.get_symbol(name) {
839            let locals = HashMap::new(); // Don't spill out the current scope
840            self.expand_defn(id, defn, locals, span)
841        } else {
842            Ok(T::identifier(name))
843        }
844    }
845
846    fn fold_pattern(
847        &mut self,
848        pattern: Box<PatternNode<'i, T>>,
849        span: pest::Span<'i>,
850    ) -> Result<T, Self::Error> {
851        if let Some((id, param, defn, _doc)) = self.aliases_map.get_pattern(pattern.name) {
852            // Resolve argument in the current scope, and pass it in to the
853            // alias expansion scope.
854            let arg = self.fold_expression(pattern.value)?;
855            let locals = HashMap::from([(param, arg)]);
856            self.expand_defn(id, defn, locals, span)
857        } else {
858            let pattern = Box::new(fold_pattern_value(self, *pattern)?);
859            Ok(T::pattern(pattern))
860        }
861    }
862
863    fn fold_function_call(
864        &mut self,
865        function: Box<FunctionCallNode<'i, T>>,
866        span: pest::Span<'i>,
867    ) -> Result<T, Self::Error> {
868        // For better error indication, builtin functions are shadowed by name,
869        // not by (name, arity).
870        if let Some(overloads) = self.aliases_map.get_function_overloads(function.name) {
871            // TODO: add support for keyword arguments
872            function
873                .ensure_no_keyword_arguments()
874                .map_err(E::invalid_arguments)?;
875            let Some((id, params, defn, _doc)) = overloads.find_by_arity(function.arity()) else {
876                let min = overloads.min_arity();
877                let max = overloads.max_arity();
878                let err = if max - min + 1 == overloads.arities().len() {
879                    function.invalid_arguments_count(min, Some(max))
880                } else {
881                    function.invalid_arguments_count_with_arities(overloads.arities())
882                };
883                return Err(E::invalid_arguments(err));
884            };
885            // Resolve arguments in the current scope, and pass them in to the alias
886            // expansion scope.
887            let args = fold_expression_nodes(self, function.args)?;
888            let locals = params.iter().map(|s| s.as_str()).zip(args).collect();
889            self.expand_defn(id, defn, locals, span)
890        } else {
891            let function = Box::new(fold_function_call_args(self, *function)?);
892            Ok(T::function_call(function))
893        }
894    }
895}
896
897/// Expands aliases recursively.
898pub fn expand_aliases<'i, T, P>(
899    node: ExpressionNode<'i, T>,
900    aliases_map: &'i AliasesMap<P, String>,
901) -> Result<ExpressionNode<'i, T>, P::Error>
902where
903    T: AliasExpandableExpression<'i> + Clone,
904    P: AliasDefinitionParser<Output<'i> = T>,
905    P::Error: AliasExpandError,
906{
907    expand_aliases_with_locals(node, aliases_map, &HashMap::new())
908}
909
910/// Expands aliases recursively with the outermost local variables.
911///
912/// Local variables are similar to alias symbols, but are scoped. Alias symbols
913/// are globally accessible from alias expressions, but local variables aren't.
914pub fn expand_aliases_with_locals<'i, T, P>(
915    node: ExpressionNode<'i, T>,
916    aliases_map: &'i AliasesMap<P, String>,
917    locals: &HashMap<&'i str, ExpressionNode<'i, T>>,
918) -> Result<ExpressionNode<'i, T>, P::Error>
919where
920    T: AliasExpandableExpression<'i> + Clone,
921    P: AliasDefinitionParser<Output<'i> = T>,
922    P::Error: AliasExpandError,
923{
924    let mut expander = AliasExpander {
925        aliases_map,
926        locals,
927        states: Vec::new(),
928    };
929    expander.fold_expression(node)
930}
931
932/// Collects similar names from the `candidates` list.
933pub fn collect_similar<I>(name: &str, candidates: I) -> Vec<String>
934where
935    I: IntoIterator,
936    I::Item: AsRef<str>,
937{
938    candidates
939        .into_iter()
940        .filter(|cand| {
941            // The parameter is borrowed from clap f5540d26
942            strsim::jaro(name, cand.as_ref()) > 0.7
943        })
944        .map(|s| s.as_ref().to_owned())
945        .sorted_unstable()
946        .collect()
947}
948
949#[cfg(test)]
950mod tests {
951    use super::*;
952
953    #[test]
954    fn test_expect_arguments() {
955        fn empty_span() -> pest::Span<'static> {
956            pest::Span::new("", 0, 0).unwrap()
957        }
958
959        fn function(
960            name: &'static str,
961            args: impl Into<Vec<ExpressionNode<'static, u32>>>,
962            keyword_args: impl Into<Vec<KeywordArgument<'static, u32>>>,
963        ) -> FunctionCallNode<'static, u32> {
964            FunctionCallNode {
965                name,
966                name_span: empty_span(),
967                args: args.into(),
968                keyword_args: keyword_args.into(),
969                args_span: empty_span(),
970            }
971        }
972
973        fn value(v: u32) -> ExpressionNode<'static, u32> {
974            ExpressionNode::new(v, empty_span())
975        }
976
977        fn keyword(name: &'static str, v: u32) -> KeywordArgument<'static, u32> {
978            KeywordArgument {
979                name,
980                name_span: empty_span(),
981                value: value(v),
982            }
983        }
984
985        let f = function("foo", [], []);
986        assert!(f.expect_no_arguments().is_ok());
987        assert!(f.expect_some_arguments::<0>().is_ok());
988        assert!(f.expect_arguments::<0, 0>().is_ok());
989        assert!(f.expect_named_arguments::<0, 0>(&[]).is_ok());
990
991        let f = function("foo", [value(0)], []);
992        assert!(f.expect_no_arguments().is_err());
993        assert_eq!(
994            f.expect_some_arguments::<0>().unwrap(),
995            (&[], [value(0)].as_slice())
996        );
997        assert_eq!(
998            f.expect_some_arguments::<1>().unwrap(),
999            (&[value(0)], [].as_slice())
1000        );
1001        assert!(f.expect_arguments::<0, 0>().is_err());
1002        assert_eq!(
1003            f.expect_arguments::<0, 1>().unwrap(),
1004            (&[], [Some(&value(0))])
1005        );
1006        assert_eq!(f.expect_arguments::<1, 1>().unwrap(), (&[value(0)], [None]));
1007        assert!(f.expect_named_arguments::<0, 0>(&[]).is_err());
1008        assert_eq!(
1009            f.expect_named_arguments::<0, 1>(&["a"]).unwrap(),
1010            ([], [Some(&value(0))])
1011        );
1012        assert_eq!(
1013            f.expect_named_arguments::<1, 0>(&["a"]).unwrap(),
1014            ([&value(0)], [])
1015        );
1016
1017        let f = function("foo", [], [keyword("a", 0)]);
1018        assert!(f.expect_no_arguments().is_err());
1019        assert!(f.expect_some_arguments::<1>().is_err());
1020        assert!(f.expect_arguments::<0, 1>().is_err());
1021        assert!(f.expect_arguments::<1, 0>().is_err());
1022        assert!(f.expect_named_arguments::<0, 0>(&[]).is_err());
1023        assert!(f.expect_named_arguments::<0, 1>(&[]).is_err());
1024        assert!(f.expect_named_arguments::<1, 0>(&[]).is_err());
1025        assert_eq!(
1026            f.expect_named_arguments::<1, 0>(&["a"]).unwrap(),
1027            ([&value(0)], [])
1028        );
1029        assert_eq!(
1030            f.expect_named_arguments::<1, 1>(&["a", "b"]).unwrap(),
1031            ([&value(0)], [None])
1032        );
1033        assert!(f.expect_named_arguments::<1, 1>(&["b", "a"]).is_err());
1034
1035        let f = function("foo", [value(0)], [keyword("a", 1), keyword("b", 2)]);
1036        assert!(f.expect_named_arguments::<0, 0>(&[]).is_err());
1037        assert!(f.expect_named_arguments::<1, 1>(&["a", "b"]).is_err());
1038        assert_eq!(
1039            f.expect_named_arguments::<1, 2>(&["c", "a", "b"]).unwrap(),
1040            ([&value(0)], [Some(&value(1)), Some(&value(2))])
1041        );
1042        assert_eq!(
1043            f.expect_named_arguments::<2, 1>(&["c", "b", "a"]).unwrap(),
1044            ([&value(0), &value(2)], [Some(&value(1))])
1045        );
1046        assert_eq!(
1047            f.expect_named_arguments::<0, 3>(&["c", "b", "a"]).unwrap(),
1048            ([], [Some(&value(0)), Some(&value(2)), Some(&value(1))])
1049        );
1050
1051        let f = function("foo", [], [keyword("a", 0), keyword("a", 1)]);
1052        assert!(f.expect_named_arguments::<1, 1>(&["", "a"]).is_err());
1053    }
1054}