[WIP] Implement enum variants

.github/workflows/rust.yml

@@ -48,7 +48,8 @@ jobs:
         with:
           tool: cargo-nextest
       - uses: Swatinem/rust-cache@v2
-      - run: cd ${{ matrix.package }} && cargo nextest run --no-tests=warn
+      - run: cd ${{ matrix.package }} && cargo build --all-targets
+      - run: cd ${{ matrix.package }} && cargo nextest run --all-targets --no-fail-fast --no-tests=warn
       - run: cd ${{ matrix.package }} && cargo clippy --all-targets -- -D warnings
 
   Docs:

rinja/src/helpers.rs

@@ -268,3 +268,7 @@ impl<L: FastWritable, R: FastWritable> FastWritable for Concat<L, R> {
         self.1.write_into(dest)
     }
 }
+
+pub trait EnumVariantTemplate {
+    fn render_into<W: fmt::Write + ?Sized>(&self, writer: &mut W) -> crate::Result<()>;
+}

rinja_derive/src/generator.rs

@@ -25,9 +25,12 @@ pub(crate) fn template_to_string(
     input: &TemplateInput<'_>,
     contexts: &HashMap<&Arc<Path>, Context<'_>, FxBuildHasher>,
     heritage: Option<&Heritage<'_, '_>>,
-    target: Option<&str>,
+    tmpl_kind: TmplKind,
 ) -> Result<usize, CompileError> {
-    let ctx = &contexts[&input.path];
+    if tmpl_kind == TmplKind::Struct {
+        buf.write("const _: () = { extern crate rinja as rinja;");
+    }
+
     let generator = Generator::new(
         input,
         contexts,
@@ -36,13 +39,27 @@ pub(crate) fn template_to_string(
         input.block.is_some(),
         0,
     );
-    let mut result = generator.build(ctx, buf, target);
-    if let Err(err) = &mut result {
-        if err.span.is_none() {
+    let size_hint = match generator.impl_template(buf, tmpl_kind) {
+        Err(mut err) if err.span.is_none() => {
             err.span = input.source_span;
+            Err(err)
         }
+        result => result,
+    }?;
+
+    if tmpl_kind == TmplKind::Struct {
+        impl_everything(input.ast, buf);
+        buf.write("};");
     }
-    result
+    Ok(size_hint)
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum TmplKind {
+    /// [`rinja::Template`]
+    Struct,
+    /// [`rinja::helpers::EnumVariantTemplate`]
+    Variant,
 }
 
 struct Generator<'a, 'h> {
@@ -97,31 +114,18 @@ impl<'a, 'h> Generator<'a, 'h> {
         }
     }
 
-    // Takes a Context and generates the relevant implementations.
-    fn build(
-        mut self,
-        ctx: &Context<'a>,
-        buf: &mut Buffer,
-        target: Option<&str>,
-    ) -> Result<usize, CompileError> {
-        if target.is_none() {
-            buf.write("const _: () = { extern crate rinja as rinja;");
-        }
-        let size_hint = self.impl_template(ctx, buf, target.unwrap_or("rinja::Template"))?;
-        if target.is_none() {
-            impl_everything(self.input.ast, buf);
-            buf.write("};");
-        }
-        Ok(size_hint)
-    }
-
     // Implement `Template` for the given context struct.
     fn impl_template(
-        &mut self,
-        ctx: &Context<'a>,
+        mut self,
         buf: &mut Buffer,
-        target: &str,
+        tmpl_kind: TmplKind,
     ) -> Result<usize, CompileError> {
+        let ctx = &self.contexts[&self.input.path];
+
+        let target = match tmpl_kind {
+            TmplKind::Struct => "rinja::Template",
+            TmplKind::Variant => "rinja::helpers::EnumVariantTemplate",
+        };
         write_header(self.input.ast, buf, target);
         buf.write(
             "fn render_into<RinjaW>(&self, __rinja_writer: &mut RinjaW) -> rinja::Result<()>\
@@ -157,12 +161,12 @@ impl<'a, 'h> Generator<'a, 'h> {
 
         let size_hint = self.impl_template_inner(ctx, buf)?;
 
-        buf.write(format_args!(
-            "\
-                rinja::Result::Ok(())\
-            }}\
-            const SIZE_HINT: rinja::helpers::core::primitive::usize = {size_hint}usize;",
-        ));
+        buf.write("rinja::Result::Ok(()) }");
+        if tmpl_kind == TmplKind::Struct {
+            buf.write(format_args!(
+                "const SIZE_HINT: rinja::helpers::core::primitive::usize = {size_hint}usize;",
+            ));
+        }
 
         buf.write('}');
         Ok(size_hint)

rinja_derive/src/input.rs

@@ -271,6 +271,63 @@ impl TemplateInput<'_> {
     }
 }
 
+pub(crate) enum AnyTemplateArgs {
+    Struct(TemplateArgs),
+    Enum {
+        enum_args: Option<PartialTemplateArgs>,
+        vars_args: Vec<Option<PartialTemplateArgs>>,
+        has_default_impl: bool,
+    },
+}
+
+impl AnyTemplateArgs {
+    pub(crate) fn new(ast: &syn::DeriveInput) -> Result<Self, CompileError> {
+        let syn::Data::Enum(enum_data) = &ast.data else {
+            return Ok(Self::Struct(TemplateArgs::new(ast)?));
+        };
+
+        let enum_args = PartialTemplateArgs::new(ast, &ast.attrs)?;
+        let vars_args = enum_data
+            .variants
+            .iter()
+            .map(|variant| PartialTemplateArgs::new(ast, &variant.attrs))
+            .collect::<Result<Vec<_>, _>>()?;
+        if vars_args.is_empty() {
+            return Ok(Self::Struct(TemplateArgs::from_partial(ast, enum_args)?));
+        }
+
+        let mut needs_default_impl = vars_args.len();
+        let enum_source = enum_args.as_ref().and_then(|v| v.source.as_ref());
+        for (variant, var_args) in enum_data.variants.iter().zip(&vars_args) {
+            if var_args
+                .as_ref()
+                .and_then(|v| v.source.as_ref())
+                .or(enum_source)
+                .is_none()
+            {
+                return Err(CompileError::new_with_span(
+                    #[cfg(not(feature = "code-in-doc"))]
+                    "either all `enum` variants need a `path` or `source` argument, \
+                    or the `enum` itself needs a default implementation",
+                    #[cfg(feature = "code-in-doc")]
+                    "either all `enum` variants need a `path`, `source` or `in_doc` argument, \
+                    or the `enum` itself needs a default implementation",
+                    None,
+                    Some(variant.ident.span()),
+                ));
+            } else if !var_args.is_none() {
+                needs_default_impl -= 1;
+            }
+        }
+
+        Ok(Self::Enum {
+            enum_args,
+            vars_args,
+            has_default_impl: needs_default_impl > 0,
+        })
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct TemplateArgs {
     pub(crate) source: (Source, Option<Span>),
@@ -626,6 +683,17 @@ pub(crate) enum PartialTemplateArgsSource {
     InDoc(Span, Source),
 }
 
+impl PartialTemplateArgsSource {
+    pub(crate) fn span(&self) -> Span {
+        match self {
+            Self::Path(s) => s.span(),
+            Self::Source(s) => s.span(),
+            #[cfg(feature = "code-in-doc")]
+            Self::InDoc(s, _) => s.span(),
+        }
+    }
+}
+
 // implement PartialTemplateArgs::new()
 const _: () = {
     impl PartialTemplateArgs {