derive: implement template attribute blocks

Kijewski · Kijewski · commit ecf1210268ca · 2025-02-08T01:50:21.000+01:00
diff --git a/rinja/Cargo.toml b/rinja/Cargo.toml
@@ -44,14 +44,15 @@ maintenance = { status = "actively-developed" }
 
 [features]
 default = ["config", "std", "urlencode"]
-full = ["default", "code-in-doc", "serde_json"]
+full = ["default", "blocks", "code-in-doc", "serde_json"]
 
 alloc = [
     "rinja_derive/alloc",
     "serde?/alloc",
     "serde_json?/alloc",
     "percent-encoding?/alloc"
 ]
+blocks = ["rinja_derive/blocks"]
 code-in-doc = ["rinja_derive/code-in-doc"]
 config = ["rinja_derive/config"]
 serde_json = ["rinja_derive/serde_json", "dep:serde", "dep:serde_json"]
diff --git a/rinja_derive/Cargo.toml b/rinja_derive/Cargo.toml
@@ -39,6 +39,7 @@ syn = { version = "2.0.3", features = ["full"] }
 
 [features]
 alloc = []
+blocks = ["syn/full"]
 code-in-doc = ["dep:pulldown-cmark"]
 config = ["dep:serde", "dep:basic-toml", "parser/config"]
 urlencode = []
diff --git a/rinja_derive/src/generator.rs b/rinja_derive/src/generator.rs
@@ -12,7 +12,10 @@ use parser::node::{Macro, Whitespace};
 use parser::{
     CharLit, Expr, FloatKind, IntKind, MAX_RUST_KEYWORD_LEN, Num, RUST_KEYWORDS, StrLit, WithSpan,
 };
+use proc_macro2::Span;
+use quote::quote_spanned;
 use rustc_hash::FxBuildHasher;
+use syn::{GenericParam, Ident, Lifetime, LifetimeParam, Token};
 
 use crate::heritage::{Context, Heritage};
 use crate::html::write_escaped_str;
@@ -25,7 +28,7 @@ pub(crate) fn template_to_string(
     input: &TemplateInput<'_>,
     contexts: &HashMap<&Arc<Path>, Context<'_>, FxBuildHasher>,
     heritage: Option<&Heritage<'_, '_>>,
-    tmpl_kind: TmplKind,
+    tmpl_kind: TmplKind<'_>,
 ) -> Result<usize, CompileError> {
     let generator = Generator::new(
         input,
@@ -50,11 +53,13 @@ pub(crate) fn template_to_string(
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub(crate) enum TmplKind {
+pub(crate) enum TmplKind<'a> {
     /// [`rinja::Template`]
     Struct,
     /// [`rinja::helpers::EnumVariantTemplate`]
     Variant,
+    /// Used in `blocks` implementation
+    Block(&'a str),
 }
 
 struct Generator<'a, 'h> {
@@ -113,13 +118,14 @@ impl<'a, 'h> Generator<'a, 'h> {
     fn impl_template(
         mut self,
         buf: &mut Buffer,
-        tmpl_kind: TmplKind,
+        tmpl_kind: TmplKind<'a>,
     ) -> Result<usize, CompileError> {
         let ctx = &self.contexts[&self.input.path];
 
         let target = match tmpl_kind {
             TmplKind::Struct => "rinja::Template",
             TmplKind::Variant => "rinja::helpers::EnumVariantTemplate",
+            TmplKind::Block(trait_name) => trait_name,
         };
         write_header(self.input.ast, buf, target);
         buf.write(
@@ -170,9 +176,147 @@ impl<'a, 'h> Generator<'a, 'h> {
         }
 
         buf.write('}');
+
+        for (block, span) in self.input.blocks {
+            self.impl_block(buf, block, span)?;
+        }
+
         Ok(size_hint)
     }
 
+    fn impl_block(&self, buf: &mut Buffer, block: &str, span: &Span) -> Result<(), CompileError> {
+        // RATIONALE: `*self` must be the input type, implementation details should not leak:
+        // - impl Self { fn as_block(self) } ->
+        // - struct __Rinja__Self__as__block__Wrapper { this: self } ->
+        // - impl Template for __Rinja__Self__as__block__Wrapper { fn render_into_with_values() } ->
+        // - impl __Rinja__Self__as__block for Self { render_into_with_values() }
+
+        let span = *span;
+        buf.write(
+            "\
+            #[allow(missing_docs, non_camel_case_types, non_snake_case, unreachable_pub)]\
+            const _: () = {",
+        );
+
+        let ident = &self.input.ast.ident;
+
+        let doc = format!("A sub-template that renders only the block `{block}` of [`{ident}`].");
+        let method_name = format!("as_{block}");
+        let trait_name = format!("__Rinja__{ident}__as__{block}");
+        let wrapper_name = format!("__Rinja__{ident}__as__{block}__Wrapper");
+        let self_lt_name = format!("'__Rinja__{ident}__as__{block}__self");
+
+        let method_id = Ident::new(&method_name, span);
+        let trait_id = Ident::new(&trait_name, span);
+        let wrapper_id = Ident::new(&wrapper_name, span);
+        let self_lt = Lifetime::new(&self_lt_name, span);
+
+        // generics of the input with an additional lifetime to capture `self`
+        let mut wrapper_generics = self.input.ast.generics.clone();
+        if wrapper_generics.lt_token.is_none() {
+            wrapper_generics.lt_token = Some(Token![<](span));
+            wrapper_generics.gt_token = Some(Token![>](span));
+        }
+        wrapper_generics.params.insert(
+            0,
+            GenericParam::Lifetime(LifetimeParam::new(self_lt.clone())),
+        );
+
+        let (impl_generics, ty_generics, where_clause) = self.input.ast.generics.split_for_impl();
+        let (wrapper_impl_generics, wrapper_ty_generics, wrapper_where_clause) =
+            wrapper_generics.split_for_impl();
+
+        let input = TemplateInput {
+            block: Some((block, span)),
+            blocks: &[],
+            ..self.input.clone()
+        };
+        let size_hint = template_to_string(
+            buf,
+            &input,
+            self.contexts,
+            self.heritage,
+            TmplKind::Block(&trait_name),
+        )?;
+
+        buf.write(quote_spanned! {
+            span =>
+            pub trait #trait_id {
+                fn render_into_with_values<RinjaW>(
+                    &self,
+                    writer: &mut RinjaW,
+                    values: &dyn rinja::Values,
+                ) -> rinja::Result<()>
+                where
+                    RinjaW:
+                        rinja::helpers::core::fmt::Write + ?rinja::helpers::core::marker::Sized;
+            }
+
+            impl #impl_generics #ident #ty_generics #where_clause {
+                #[inline]
+                #[doc = #doc]
+                pub fn #method_id(&self) -> impl rinja::Template + '_ {
+                    #wrapper_id {
+                        this: self,
+                    }
+                }
+            }
+
+            #[rinja::helpers::core::prelude::rust_2021::derive(
+                rinja::helpers::core::prelude::rust_2021::Clone,
+                rinja::helpers::core::prelude::rust_2021::Copy
+            )]
+            pub struct #wrapper_id #wrapper_generics #wrapper_where_clause {
+                this: &#self_lt #ident #ty_generics,
+            }
+
+            impl #wrapper_impl_generics rinja::Template
+            for #wrapper_id #wrapper_ty_generics #wrapper_where_clause {
+                #[inline]
+                fn render_into_with_values<RinjaW>(
+                    &self,
+                    writer: &mut RinjaW,
+                    values: &dyn rinja::Values
+                ) -> rinja::Result<()>
+                where
+                    RinjaW: rinja::helpers::core::fmt::Write + ?rinja::helpers::core::marker::Sized
+                {
+                    <_ as #trait_id>::render_into_with_values(self.this, writer, values)
+                }
+
+                const SIZE_HINT: rinja::helpers::core::primitive::usize = #size_hint;
+            }
+
+            // cannot use `crate::integrations::impl_fast_writable()` w/o cloning the struct
+            impl #wrapper_impl_generics rinja::filters::FastWritable
+            for #wrapper_id #wrapper_ty_generics #wrapper_where_clause {
+                #[inline]
+                fn write_into<RinjaW>(&self, dest: &mut RinjaW) -> rinja::Result<()>
+                where
+                    RinjaW: rinja::helpers::core::fmt::Write + ?rinja::helpers::core::marker::Sized
+                {
+                    <_ as rinja::Template>::render_into(self, dest)
+                }
+            }
+
+            // cannot use `crate::integrations::impl_display()` w/o cloning the struct
+            impl #wrapper_impl_generics rinja::helpers::core::fmt::Display
+            for #wrapper_id #wrapper_ty_generics #wrapper_where_clause {
+                #[inline]
+                fn fmt(
+                    &self,
+                    f: &mut rinja::helpers::core::fmt::Formatter<'_>
+                ) -> rinja::helpers::core::fmt::Result {
+                    <_ as rinja::Template>::render_into(self, f)
+                        .map_err(|_| rinja::helpers::core::fmt::Error)
+                }
+            }
+        });
+
+        buf.write("};");
+        Ok(())
+    }
+
     fn is_var_defined(&self, var_name: &str) -> bool {
         self.locals.get(var_name).is_some() || self.input.fields.iter().any(|f| f == var_name)
     }
diff --git a/rinja_derive/src/input.rs b/rinja_derive/src/input.rs
@@ -12,11 +12,14 @@ use proc_macro2::Span;
 use rustc_hash::FxBuildHasher;
 use syn::punctuated::Punctuated;
 use syn::spanned::Spanned;
-use syn::{Attribute, Expr, ExprLit, ExprPath, Ident, Lit, LitBool, LitStr, Meta, Token};
+use syn::{
+    Attribute, Expr, ExprArray, ExprLit, ExprPath, Ident, Lit, LitBool, LitStr, Meta, Token,
+};
 
 use crate::config::{Config, SyntaxAndCache};
 use crate::{CompileError, FileInfo, MsgValidEscapers, OnceMap};
 
+#[derive(Clone)]
 pub(crate) struct TemplateInput<'a> {
     pub(crate) ast: &'a syn::DeriveInput,
     pub(crate) enum_ast: Option<&'a syn::DeriveInput>,
@@ -25,10 +28,11 @@ pub(crate) struct TemplateInput<'a> {
     pub(crate) source: &'a Source,
     pub(crate) source_span: Option<Span>,
     pub(crate) block: Option<(&'a str, Span)>,
+    pub(crate) blocks: &'a [(String, Span)],
     pub(crate) print: Print,
     pub(crate) escaper: &'a str,
     pub(crate) path: Arc<Path>,
-    pub(crate) fields: Vec<String>,
+    pub(crate) fields: Arc<[String]>,
 }
 
 impl TemplateInput<'_> {
@@ -44,6 +48,7 @@ impl TemplateInput<'_> {
         let TemplateArgs {
             source: (source, source_span),
             block,
+            blocks,
             print,
             escaping,
             ext,
@@ -134,10 +139,11 @@ impl TemplateInput<'_> {
             source,
             source_span: *source_span,
             block: block.as_ref().map(|(block, span)| (block.as_str(), *span)),
+            blocks: blocks.as_slice(),
             print: *print,
             escaper,
             path,
-            fields,
+            fields: fields.into(),
         })
     }
 
@@ -347,6 +353,7 @@ impl AnyTemplateArgs {
 pub(crate) struct TemplateArgs {
     pub(crate) source: (Source, Option<Span>),
     block: Option<(String, Span)>,
+    blocks: Vec<(String, Span)>,
     print: Print,
     escaping: Option<String>,
     ext: Option<String>,
@@ -396,6 +403,12 @@ impl TemplateArgs {
                 }
             },
             block: args.block.map(|value| (value.value(), value.span())),
+            blocks: args
+                .blocks
+                .unwrap_or_default()
+                .into_iter()
+                .map(|value| (value.value(), value.span()))
+                .collect(),
             print: args.print.unwrap_or_default(),
             escaping: args.escape.map(|value| value.value()),
             ext: args.ext.as_ref().map(|value| value.value()),
@@ -413,6 +426,7 @@ impl TemplateArgs {
         Self {
             source: (Source::Source("".into()), None),
             block: None,
+            blocks: vec![],
             print: Print::default(),
             escaping: None,
             ext: Some("txt".to_string()),
@@ -692,6 +706,7 @@ pub(crate) struct PartialTemplateArgs {
     pub(crate) config: Option<LitStr>,
     pub(crate) whitespace: Option<Whitespace>,
     pub(crate) crate_name: Option<ExprPath>,
+    pub(crate) blocks: Option<Vec<LitStr>>,
 }
 
 #[derive(Clone)]
@@ -754,6 +769,7 @@ const _: () = {
             config: None,
             whitespace: None,
             crate_name: None,
+            blocks: None,
         };
         let mut has_data = false;
 
@@ -806,6 +822,23 @@ const _: () = {
                     ensure_only_once(ident, &mut this.crate_name)?;
                     this.crate_name = Some(get_exprpath(ident, pair.value)?);
                     continue;
+                } else if ident == "blocks" {
+                    if is_enum_variant {
+                        return Err(CompileError::no_file_info(
+                            "template attribute `blocks` can only be used on the `enum`, \
+                            not its variants",
+                            Some(ident.span()),
+                        ));
+                    }
+                    ensure_only_once(ident, &mut this.blocks)?;
+                    this.blocks = Some(
+                        get_exprarray(ident, pair.value)?
+                            .elems
+                            .into_iter()
+                            .map(|value| get_strlit(ident, get_lit(ident, value)?))
+                            .collect::<Result<_, _>>()?,
+                    );
+                    continue;
                 }
 
                 let value = get_lit(ident, pair.value)?;
@@ -965,6 +998,21 @@ const _: () = {
         }
     }
 
+    fn get_exprarray(name: &Ident, mut expr: Expr) -> Result<ExprArray, CompileError> {
+        loop {
+            match expr {
+                Expr::Array(array) => return Ok(array),
+                Expr::Group(group) => expr = *group.expr,
+                v => {
+                    return Err(CompileError::no_file_info(
+                        format_args!("template attribute `{name}` expects an array"),
+                        Some(v.span()),
+                    ));
+                }
+            }
+        }
+    }
+
     fn ensure_source_only_once(
         name: &Ident,
         source: &Option<PartialTemplateArgsSource>,
diff --git a/rinja_derive/src/lib.rs b/rinja_derive/src/lib.rs
@@ -262,7 +262,7 @@ fn build_template_item(
     ast: &syn::DeriveInput,
     enum_ast: Option<&syn::DeriveInput>,
     template_args: &TemplateArgs,
-    tmpl_kind: TmplKind,
+    tmpl_kind: TmplKind<'_>,
 ) -> Result<usize, CompileError> {
     let config_path = template_args.config_path();
     let s = read_config_file(config_path, template_args.config_span)?;
diff --git a/rinja_derive_standalone/Cargo.toml b/rinja_derive_standalone/Cargo.toml
@@ -46,6 +46,7 @@ syn = { version = "2.0.3", features = ["full"] }
 default = ["__standalone"]
 __standalone = []
 
+blocks = ["syn/full"]
 code-in-doc = ["dep:pulldown-cmark"]
 config = ["dep:serde", "dep:basic-toml", "parser/config"]
 urlencode = []
diff --git a/testing/tests/blocks.rs b/testing/tests/blocks.rs
