generator: Move macro invocation code generation to helper module

Introduced a new helpers module where code generation can be placed.
For this, some generator methods had to become pub(crate) instead of private.

askama_derive/src/generator.rs

@@ -1,5 +1,6 @@
 mod expr;
 mod filter;
+mod helpers;
 mod node;
 
 use std::borrow::Cow;
@@ -85,7 +86,7 @@ struct Generator<'a, 'h> {
     /// Used in blocks to check if we are inside a filter block.
     is_in_filter_block: usize,
     /// Set of called macros we are currently in. Used to prevent (indirect) recursions.
-    seen_callers: Vec<(&'a Call<'a>, &'a Macro<'a>, Option<FileInfo<'a>>)>,
+    seen_callers: Vec<(&'a Macro<'a>, Option<FileInfo<'a>>)>,
 }
 
 impl<'a, 'h> Generator<'a, 'h> {

askama_derive/src/generator/helpers/macro_invocation.rs

@@ -0,0 +1,314 @@
+use core::fmt;
+use std::borrow::Cow;
+use std::collections::HashMap;
+use std::fmt::Write;
+use std::mem;
+
+use parser::node::{Call, Macro, Ws};
+use parser::{Expr, Span, WithSpan};
+use rustc_hash::FxBuildHasher;
+
+use crate::CompileError;
+use crate::generator::node::AstLevel;
+use crate::generator::{Generator, LocalMeta, is_copyable, normalize_identifier};
+use crate::heritage::Context;
+use crate::integration::Buffer;
+
+/// Helper to generate the code for macro invocations
+pub(crate) struct MacroInvocation<'a, 'b> {
+    pub callsite_ctx: &'b Context<'a>,
+    pub callsite_span: Span<'a>,
+    pub callsite_ws: Ws,
+    pub call_args: &'b Vec<WithSpan<'a, Expr<'a>>>,
+    pub call: Option<&'a WithSpan<'a, Call<'a>>>,
+    pub macro_def: &'a Macro<'a>,
+    pub macro_ctx: &'b Context<'a>,
+}
+
+impl<'a, 'b> MacroInvocation<'a, 'b> {
+    pub(crate) fn write<'h>(
+        &self,
+        buf: &'b mut Buffer,
+        generator: &mut Generator<'a, 'h>,
+    ) -> Result<usize, CompileError> {
+        if generator
+            .seen_callers
+            .iter()
+            .any(|(s, _)| std::ptr::eq(*s, self.macro_def))
+        {
+            let mut message = "Found recursion in macro calls:".to_owned();
+            for (m, f) in &generator.seen_callers {
+                if let Some(f) = f {
+                    write!(message, "{f}").unwrap();
+                } else {
+                    write!(message, "\n`{}`", m.name.escape_debug()).unwrap();
+                }
+            }
+            return Err(self
+                .callsite_ctx
+                .generate_error(message, self.callsite_span));
+        } else {
+            generator.seen_callers.push((
+                self.macro_def,
+                self.callsite_ctx.file_info_of(self.callsite_span),
+            ));
+        }
+
+        generator.push_locals(|this| {
+            if let Some(call) = self.call {
+                this.locals.insert(
+                    "caller".into(),
+                    LocalMeta::caller(call, self.callsite_ctx.clone()),
+                );
+            }
+
+            self.ensure_arg_count()?;
+
+            this.flush_ws(self.callsite_ws); // Cannot handle_ws() here: whitespace from macro definition comes first
+            this.write_buf_writable(self.callsite_ctx, buf)?;
+            buf.write('{');
+            this.prepare_ws(self.macro_def.ws1);
+
+            self.write_preamble(buf, this)?;
+
+            let mut size_hint =
+                this.handle(self.macro_ctx, &self.macro_def.nodes, buf, AstLevel::Nested)?;
+
+            this.flush_ws(self.macro_def.ws2);
+            size_hint += this.write_buf_writable(self.callsite_ctx, buf)?;
+            buf.write('}');
+
+            this.prepare_ws(self.callsite_ws);
+            this.seen_callers.pop();
+            Ok(size_hint)
+        })
+    }
+
+    fn write_preamble<'h>(
+        &self,
+        buf: &'b mut Buffer,
+        generator: &mut Generator<'a, 'h>,
+    ) -> Result<(), CompileError> {
+        let mut named_arguments: HashMap<&str, _, FxBuildHasher> = HashMap::default();
+        if let Some(Expr::NamedArgument(_, _)) = self.call_args.last().map(|expr| &**expr) {
+            // First we check that all named arguments actually exist in the called item.
+            for (index, arg) in self.call_args.iter().enumerate().rev() {
+                let Expr::NamedArgument(arg_name, _) = &**arg else {
+                    break;
+                };
+                if !self.macro_def.args.iter().any(|(arg, _)| arg == arg_name) {
+                    return Err(self.callsite_ctx.generate_error(
+                        format_args!(
+                            "no argument named `{arg_name}` in macro {}",
+                            self.macro_def.name
+                        ),
+                        self.callsite_span,
+                    ));
+                }
+                named_arguments.insert(arg_name, (index, arg));
+            }
+        }
+        let mut value = Buffer::new();
+        let mut allow_positional = true;
+        let mut used_named_args = vec![false; self.call_args.len()];
+
+        for (index, (arg, default_value)) in self.macro_def.args.iter().enumerate() {
+            let expr = if let Some((index, expr)) = named_arguments.get(arg) {
+                used_named_args[*index] = true;
+                allow_positional = false;
+                expr
+            } else {
+                match self.call_args.get(index) {
+                    Some(arg_expr) if !matches!(**arg_expr, Expr::NamedArgument(_, _)) => {
+                        // If there is already at least one named argument, then it's not allowed
+                        // to use unnamed ones at this point anymore.
+                        if !allow_positional {
+                            return Err(self.callsite_ctx.generate_error(
+                                format_args!(
+                                    "cannot have unnamed argument (`{arg}`) after named argument \
+                                        in call to macro {}",
+                                    self.macro_def.name
+                                ),
+                                self.callsite_span,
+                            ));
+                        }
+                        arg_expr
+                    }
+                    Some(arg_expr) if used_named_args[index] => {
+                        let Expr::NamedArgument(name, _) = **arg_expr else {
+                            unreachable!()
+                        };
+                        return Err(self.callsite_ctx.generate_error(
+                            format_args!("`{name}` is passed more than once"),
+                            self.callsite_span,
+                        ));
+                    }
+                    _ => {
+                        if let Some(default_value) = default_value {
+                            default_value
+                        } else {
+                            return Err(self.callsite_ctx.generate_error(
+                                format_args!("missing `{arg}` argument"),
+                                self.callsite_span,
+                            ));
+                        }
+                    }
+                }
+            };
+            match &**expr {
+                // If `expr` is already a form of variable then
+                // don't reintroduce a new variable. This is
+                // to avoid moving non-copyable values.
+                Expr::Var(name) if *name != "self" => {
+                    let var = generator.locals.resolve_or_self(name);
+                    generator
+                        .locals
+                        .insert(Cow::Borrowed(arg), LocalMeta::var_with_ref(var));
+                }
+                Expr::AssociatedItem(obj, associated_item) => {
+                    let mut associated_item_buf = Buffer::new();
+                    generator.visit_associated_item(
+                        self.callsite_ctx,
+                        &mut associated_item_buf,
+                        obj,
+                        associated_item,
+                    )?;
+
+                    let associated_item = associated_item_buf.into_string();
+                    let var = generator
+                        .locals
+                        .resolve(&associated_item)
+                        .unwrap_or(associated_item);
+                    generator
+                        .locals
+                        .insert(Cow::Borrowed(arg), LocalMeta::var_with_ref(var));
+                }
+                // Everything else still needs to become variables,
+                // to avoid having the same logic be executed
+                // multiple times, e.g. in the case of macro
+                // parameters being used multiple times.
+                _ => {
+                    value.clear();
+                    let (before, after) = if !is_copyable(expr) {
+                        ("&(", ")")
+                    } else {
+                        ("", "")
+                    };
+                    value.write(generator.visit_expr_root(self.callsite_ctx, expr)?);
+                    // We need to normalize the arg to write it, thus we need to add it to
+                    // locals in the normalized manner
+                    let normalized_arg = normalize_identifier(arg);
+                    buf.write(format_args!(
+                        "let {normalized_arg} = {before}{value}{after};"
+                    ));
+                    generator
+                        .locals
+                        .insert_with_default(Cow::Borrowed(normalized_arg));
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    fn ensure_arg_count(&self) -> Result<(), CompileError> {
+        if self.call_args.len() > self.macro_def.args.len() {
+            return Err(self.callsite_ctx.generate_error(
+                format_args!(
+                    "macro `{}` expected {} argument{}, found {}",
+                    self.macro_def.name,
+                    self.macro_def.args.len(),
+                    if self.macro_def.args.len() > 1 {
+                        "s"
+                    } else {
+                        ""
+                    },
+                    self.call_args.len(),
+                ),
+                self.callsite_span,
+            ));
+        }
+
+        // First we list of arguments position, then we remove every argument with a value.
+        let mut args: Vec<_> = self
+            .macro_def
+            .args
+            .iter()
+            .map(|&(name, _)| Some(name))
+            .collect();
+        for (pos, arg) in self.call_args.iter().enumerate() {
+            let pos = match **arg {
+                Expr::NamedArgument(name, ..) => self
+                    .macro_def
+                    .args
+                    .iter()
+                    .position(|(arg_name, _)| *arg_name == name),
+                _ => Some(pos),
+            };
+            if let Some(pos) = pos
+                && mem::take(&mut args[pos]).is_none()
+            {
+                // This argument was already passed, so error.
+                return Err(self.callsite_ctx.generate_error(
+                    format_args!(
+                        "argument `{}` was passed more than once when calling macro `{}`",
+                        self.macro_def.args[pos].0, self.macro_def.name,
+                    ),
+                    self.callsite_span,
+                ));
+            }
+        }
+
+        // Now we can check off arguments with a default value, too.
+        for (pos, (_, dflt)) in self.macro_def.args.iter().enumerate() {
+            if dflt.is_some() {
+                args[pos] = None;
+            }
+        }
+
+        // Now that we have a needed information, we can print an error message (if needed).
+        struct FmtMissing<'a, I> {
+            count: usize,
+            missing: I,
+            name: &'a str,
+        }
+
+        impl<'a, I: Iterator<Item = &'a str> + Clone> fmt::Display for FmtMissing<'a, I> {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                if self.count == 1 {
+                    let a = self.missing.clone().next().unwrap();
+                    write!(
+                        f,
+                        "missing argument when calling macro `{}`: `{a}`",
+                        self.name
+                    )
+                } else {
+                    write!(f, "missing arguments when calling macro `{}`: ", self.name)?;
+                    for (idx, a) in self.missing.clone().enumerate() {
+                        if idx == self.count - 1 {
+                            write!(f, " and ")?;
+                        } else if idx > 0 {
+                            write!(f, ", ")?;
+                        }
+                        write!(f, "`{a}`")?;
+                    }
+                    Ok(())
+                }
+            }
+        }
+
+        let missing = args.iter().filter_map(Option::as_deref);
+        let fmt_missing = FmtMissing {
+            count: missing.clone().count(),
+            missing,
+            name: self.macro_def.name,
+        };
+        if fmt_missing.count == 0 {
+            Ok(())
+        } else {
+            Err(self
+                .callsite_ctx
+                .generate_error(fmt_missing, self.callsite_span))
+        }
+    }
+}

askama_derive/src/generator/helpers/mod.rs

@@ -0,0 +1,3 @@
+mod macro_invocation;
+
+pub(crate) use macro_invocation::*;

askama_derive/src/generator/node.rs

@@ -1,21 +1,19 @@
 use std::borrow::Cow;
 use std::collections::hash_map::{Entry, HashMap};
-use std::fmt::{self, Debug, Write};
+use std::fmt::Debug;
 use std::mem;
 
 use parser::expr::BinOp;
 use parser::node::{
-    Call, Comment, Cond, CondTest, FilterBlock, If, Include, Let, Lit, Loop, Macro, Match,
-    Whitespace, Ws,
+    Call, Comment, Cond, CondTest, FilterBlock, If, Include, Let, Lit, Loop, Match, Whitespace, Ws,
 };
 use parser::{Expr, Node, Span, Target, WithSpan};
-use rustc_hash::FxBuildHasher;
 
 use super::{
     DisplayWrap, FILTER_SOURCE, Generator, LocalMeta, MapChain, compile_time_escape, is_copyable,
     normalize_identifier,
 };
-use crate::generator::{LocalCallerMeta, Writable};
+use crate::generator::{LocalCallerMeta, Writable, helpers};
 use crate::heritage::{Context, Heritage};
 use crate::integration::Buffer;
 use crate::{CompileError, FileInfo, fmt_left, fmt_right};
@@ -37,7 +35,7 @@ impl<'a> Generator<'a, '_> {
         Ok(size_hint)
     }
 
-    fn push_locals<T, F>(&mut self, callback: F) -> Result<T, CompileError>
+    pub(crate) fn push_locals<T, F>(&mut self, callback: F) -> Result<T, CompileError>
     where
         F: FnOnce(&mut Self) -> Result<T, CompileError>,
     {
@@ -80,7 +78,7 @@ impl<'a> Generator<'a, '_> {
         res
     }
 
-    fn handle(
+    pub(crate) fn handle(
         &mut self,
         ctx: &Context<'a>,
         nodes: &'a [Node<'_>],
@@ -620,150 +618,19 @@ impl<'a> Generator<'a, '_> {
             (*def, ctx)
         };
 
-        if self
-            .seen_callers
-            .iter()
-            .any(|(_, s, _)| std::ptr::eq(*s, def))
-        {
-            let mut message = "Found recursion in macro calls:".to_owned();
-            for (_, m, f) in &self.seen_callers {
-                if let Some(f) = f {
-                    write!(message, "{f}").unwrap();
-                } else {
-                    write!(message, "\n`{}`", m.name.escape_debug()).unwrap();
-                }
-            }
-            return Err(ctx.generate_error(message, call.span()));
-        } else {
-            self.seen_callers
-                .push((call, def, ctx.file_info_of(call.span())));
+        // whitespaces for the invocation is constructed from
+        // - call-block's outer (start)
+        // - endcall-block's outer (end)
+        helpers::MacroInvocation {
+            callsite_ctx: ctx,
+            callsite_span: call.span(),
+            call: Some(call),
+            callsite_ws: Ws(ws1.0, ws2.1),
+            call_args: args,
+            macro_def: def,
+            macro_ctx: own_ctx,
         }
-        self.flush_ws(ws1); // Cannot handle_ws() here: whitespace from macro definition comes first
-        let size_hint = self.push_locals(|this| {
-            this.locals.insert("caller".into(), LocalMeta::caller(call, ctx.clone()));
-
-            macro_call_ensure_arg_count(call, def, ctx)?;
-
-            this.write_buf_writable(ctx, buf)?;
-            buf.write('{');
-            this.prepare_ws(def.ws1);
-
-            let mut named_arguments: HashMap<&str, _, FxBuildHasher> = HashMap::default();
-            // Since named arguments can only be passed last, we only need to check if the last argument
-            // is a named one.
-            if let Some(Expr::NamedArgument(_, _)) = args.last().map(|expr| &**expr) {
-                // First we check that all named arguments actually exist in the called item.
-                for (index, arg) in args.iter().enumerate().rev() {
-                    let Expr::NamedArgument(arg_name, _) = &**arg else {
-                        break;
-                    };
-                    if !def.args.iter().any(|(arg, _)| arg == arg_name) {
-                        return Err(ctx.generate_error(
-                            format_args!("no argument named `{arg_name}` in macro {name:?}"),
-                            call.span(),
-                        ));
-                    }
-                    named_arguments.insert(arg_name, (index, arg));
-                }
-            }
-
-            let mut value = Buffer::new();
-
-            // Handling both named and unnamed arguments requires to be careful of the named arguments
-            // order. To do so, we iterate through the macro defined arguments and then check if we have
-            // a named argument with this name:
-            //
-            // * If there is one, we add it and move to the next argument.
-            // * If there isn't one, then we pick the next argument (we can do it without checking
-            //   anything since named arguments are always last).
-            let mut allow_positional = true;
-            let mut used_named_args = vec![false; args.len()];
-            for (index, (arg, default_value)) in def.args.iter().enumerate() {
-                let expr = if let Some((index, expr)) = named_arguments.get(arg) {
-                    used_named_args[*index] = true;
-                    allow_positional = false;
-                    expr
-                } else {
-                    match args.get(index) {
-                        Some(arg_expr) if !matches!(**arg_expr, Expr::NamedArgument(_, _)) => {
-                            // If there is already at least one named argument, then it's not allowed
-                            // to use unnamed ones at this point anymore.
-                            if !allow_positional {
-                                return Err(ctx.generate_error(
-                                    format_args!(
-                                        "cannot have unnamed argument (`{arg}`) after named argument \
-                                         in call to macro {name:?}"
-                                    ),
-                                    call.span(),
-                                ));
-                            }
-                            arg_expr
-                        }
-                        Some(arg_expr) if used_named_args[index] => {
-                            let Expr::NamedArgument(name, _) = **arg_expr else { unreachable!() };
-                            return Err(ctx.generate_error(
-                                format_args!("`{name}` is passed more than once"),
-                                call.span(),
-                            ));
-                        }
-                        _ => {
-                            if let Some(default_value) = default_value {
-                                default_value
-                            } else {
-                                return Err(ctx.generate_error(format_args!("missing `{arg}` argument"), call.span()));
-                            }
-                        }
-                    }
-                };
-                match &**expr {
-                    // If `expr` is already a form of variable then
-                    // don't reintroduce a new variable. This is
-                    // to avoid moving non-copyable values.
-                    Expr::Var(name) if *name != "self" => {
-                        let var = this.locals.resolve_or_self(name);
-                        this.locals
-                            .insert(Cow::Borrowed(arg), LocalMeta::var_with_ref(var));
-                    }
-                    Expr::AssociatedItem(obj, associated_item) => {
-                        let mut associated_item_buf = Buffer::new();
-                        this.visit_associated_item(ctx, &mut associated_item_buf, obj, associated_item)?;
-
-                        let associated_item = associated_item_buf.into_string();
-                        let var = this.locals.resolve(&associated_item).unwrap_or(associated_item);
-                        this.locals
-                            .insert(Cow::Borrowed(arg), LocalMeta::var_with_ref(var));
-                    }
-                    // Everything else still needs to become variables,
-                    // to avoid having the same logic be executed
-                    // multiple times, e.g. in the case of macro
-                    // parameters being used multiple times.
-                    _ => {
-                        value.clear();
-                        let (before, after) = if !is_copyable(expr) {
-                            ("&(", ")")
-                        } else {
-                            ("", "")
-                        };
-                        value.write(this.visit_expr_root(ctx, expr)?);
-                        // We need to normalize the arg to write it, thus we need to add it to
-                        // locals in the normalized manner
-                        let normalized_arg = normalize_identifier(arg);
-                        buf.write(format_args!("let {normalized_arg} = {before}{value}{after};"));
-                        this.locals.insert_with_default(Cow::Borrowed(normalized_arg));
-                    }
-                }
-            }
-
-            let mut size_hint = this.handle(own_ctx, &def.nodes, buf, AstLevel::Nested)?;
-
-            this.flush_ws(def.ws2);
-            size_hint += this.write_buf_writable(ctx, buf)?;
-            buf.write('}');
-            Ok(size_hint)
-        })?;
-        self.prepare_ws(ws2);
-        self.seen_callers.pop();
-        Ok(size_hint)
+        .write(buf, self)
     }
 
     fn write_filter_block(
@@ -1252,7 +1119,7 @@ impl<'a> Generator<'a, '_> {
     }
 
     // Write expression buffer and empty
-    fn write_buf_writable(
+    pub(crate) fn write_buf_writable(
         &mut self,
         ctx: &Context<'_>,
         buf: &mut Buffer,
@@ -1384,7 +1251,7 @@ impl<'a> Generator<'a, '_> {
 
     // Combines `flush_ws()` and `prepare_ws()` to handle both trailing whitespace from the
     // preceding literal and leading whitespace from the succeeding literal.
-    fn handle_ws(&mut self, ws: Ws) {
+    pub(crate) fn handle_ws(&mut self, ws: Ws) {
         self.flush_ws(ws);
         self.prepare_ws(ws);
     }
@@ -1396,7 +1263,7 @@ impl<'a> Generator<'a, '_> {
     // If the previous literal left some trailing whitespace in `next_ws` and the
     // prefix whitespace suppressor from the given argument, flush that whitespace.
     // In either case, `next_ws` is reset to `None` (no trailing whitespace).
-    fn flush_ws(&mut self, ws: Ws) {
+    pub(crate) fn flush_ws(&mut self, ws: Ws) {
         if self.next_ws.is_none() {
             return;
         }
@@ -1429,7 +1296,7 @@ impl<'a> Generator<'a, '_> {
     // Sets `skip_ws` to match the suffix whitespace suppressor from the given
     // argument, to determine whether to suppress leading whitespace from the
     // next literal.
-    fn prepare_ws(&mut self, ws: Ws) {
+    pub(crate) fn prepare_ws(&mut self, ws: Ws) {
         self.skip_ws = self.should_trim_ws(ws.1);
     }
 }
@@ -1578,100 +1445,8 @@ fn median(sizes: &mut [usize]) -> usize {
     }
 }
 
-fn macro_call_ensure_arg_count(
-    call: &WithSpan<'_, Call<'_>>,
-    def: &Macro<'_>,
-    ctx: &Context<'_>,
-) -> Result<(), CompileError> {
-    if call.args.len() > def.args.len() {
-        return Err(ctx.generate_error(
-            format_args!(
-                "macro `{}` expected {} argument{}, found {}",
-                def.name,
-                def.args.len(),
-                if def.args.len() > 1 { "s" } else { "" },
-                call.args.len(),
-            ),
-            call.span(),
-        ));
-    }
-
-    // First we list of arguments position, then we remove every argument with a value.
-    let mut args: Vec<_> = def.args.iter().map(|&(name, _)| Some(name)).collect();
-    for (pos, arg) in call.args.iter().enumerate() {
-        let pos = match **arg {
-            Expr::NamedArgument(name, ..) => {
-                def.args.iter().position(|(arg_name, _)| *arg_name == name)
-            }
-            _ => Some(pos),
-        };
-        if let Some(pos) = pos
-            && mem::take(&mut args[pos]).is_none()
-        {
-            // This argument was already passed, so error.
-            return Err(ctx.generate_error(
-                format_args!(
-                    "argument `{}` was passed more than once when calling macro `{}`",
-                    def.args[pos].0, def.name,
-                ),
-                call.span(),
-            ));
-        }
-    }
-
-    // Now we can check off arguments with a default value, too.
-    for (pos, (_, dflt)) in def.args.iter().enumerate() {
-        if dflt.is_some() {
-            args[pos] = None;
-        }
-    }
-
-    // Now that we have a needed information, we can print an error message (if needed).
-    struct FmtMissing<'a, I> {
-        count: usize,
-        missing: I,
-        name: &'a str,
-    }
-
-    impl<'a, I: Iterator<Item = &'a str> + Clone> fmt::Display for FmtMissing<'a, I> {
-        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-            if self.count == 1 {
-                let a = self.missing.clone().next().unwrap();
-                write!(
-                    f,
-                    "missing argument when calling macro `{}`: `{a}`",
-                    self.name
-                )
-            } else {
-                write!(f, "missing arguments when calling macro `{}`: ", self.name)?;
-                for (idx, a) in self.missing.clone().enumerate() {
-                    if idx == self.count - 1 {
-                        write!(f, " and ")?;
-                    } else if idx > 0 {
-                        write!(f, ", ")?;
-                    }
-                    write!(f, "`{a}`")?;
-                }
-                Ok(())
-            }
-        }
-    }
-
-    let missing = args.iter().filter_map(Option::as_deref);
-    let fmt_missing = FmtMissing {
-        count: missing.clone().count(),
-        missing,
-        name: def.name,
-    };
-    if fmt_missing.count == 0 {
-        Ok(())
-    } else {
-        Err(ctx.generate_error(fmt_missing, call.span()))
-    }
-}
-
 #[derive(Clone, Copy, PartialEq)]
-enum AstLevel {
+pub(crate) enum AstLevel {
     Top,
     Block,
     Nested,