Allow included templates to `extend`, `import`, and `macro`

Signed-off-by: max <gmx.sht@gmail.com>

3 files changed, 95 insertions, 58 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index c1a8ebe..22f996f 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -1,8 +1,8 @@
 use std::collections::hash_map::{Entry, HashMap};
-use std::path::{Path, PathBuf};
+use std::path::Path;
 use std::{cmp, hash, mem, str};
 
-use crate::config::{get_template_source, WhitespaceHandling};
+use crate::config::WhitespaceHandling;
 use crate::heritage::{Context, Heritage};
 use crate::input::{Source, TemplateInput};
 use crate::CompileError;
@@ -10,7 +10,7 @@ use crate::CompileError;
 use parser::node::{
     Call, Comment, CondTest, If, Include, Let, Lit, Loop, Match, Target, Whitespace, Ws,
 };
-use parser::{Expr, Node, Parsed};
+use parser::{Expr, Node};
 use quote::quote;
 
 pub(crate) struct Generator<'a> {
@@ -20,8 +20,6 @@ pub(crate) struct Generator<'a> {
     contexts: &'a HashMap<&'a Path, Context<'a>>,
     // The heritage contains references to blocks and their ancestry
     heritage: Option<&'a Heritage<'a>>,
-    // Cache ASTs for included templates
-    includes: HashMap<PathBuf, Parsed>,
     // Variables accessible directly from the current scope (not redirected to context)
     locals: MapChain<'a, &'a str, LocalMeta>,
     // Suffix whitespace from the previous literal. Will be flushed to the
@@ -50,7 +48,6 @@ impl<'a> Generator<'a> {
             input,
             contexts,
             heritage,
-            includes: HashMap::default(),
             locals,
             next_ws: None,
             skip_ws: WhitespaceHandling::Preserve,
@@ -846,24 +843,35 @@ impl<'a> Generator<'a> {
             )?;
         }
 
-        // Since nodes must not outlive the Generator, we instantiate a nested `Generator` here to
-        // handle the include's nodes. Unfortunately we can't easily share the `includes` cache.
+        // We clone the context of the child in order to preserve their macros and imports.
+        // But also add all the imports and macros from this template that don't override the
+        // child's ones to preserve this template's context.
+        let child_ctx = &mut self.contexts[path.as_path()].clone();
+        for (name, mac) in &ctx.macros {
+            child_ctx.macros.entry(name).or_insert(mac);
+        }
+        for (name, import) in &ctx.imports {
+            child_ctx
+                .imports
+                .entry(name)
+                .or_insert_with(|| import.clone());
+        }
 
-        let locals = MapChain::with_parent(&self.locals);
-        let mut child = Self::new(self.input, self.contexts, self.heritage, locals);
-
-        let nodes = match self.contexts.get(path.as_path()) {
-            Some(ctx) => ctx.nodes,
-            None => match self.includes.entry(path) {
-                Entry::Occupied(entry) => entry.into_mut().nodes(),
-                Entry::Vacant(entry) => {
-                    let src = get_template_source(entry.key())?;
-                    entry.insert(Parsed::new(src, self.input.syntax)?).nodes()
-                }
-            },
+        // Create a new generator for the child, and call it like in `impl_template` as if it were
+        // a full template, while preserving the context.
+        let heritage = if !child_ctx.blocks.is_empty() || child_ctx.extends.is_some() {
+            Some(Heritage::new(child_ctx, self.contexts))
+        } else {
+            None
         };
 
-        let mut size_hint = child.handle(ctx, nodes, buf, AstLevel::Nested)?;
+        let handle_ctx = match &heritage {
+            Some(heritage) => heritage.root,
+            None => child_ctx,
+        };
+        let locals = MapChain::with_parent(&self.locals);
+        let mut child = Self::new(self.input, self.contexts, heritage.as_ref(), locals);
+        let mut size_hint = child.handle(handle_ctx, handle_ctx.nodes, buf, AstLevel::Top)?;
         size_hint += child.write_buf_writable(buf)?;
         self.prepare_ws(i.ws);
 
diff --git a/askama_derive/src/heritage.rs b/askama_derive/src/heritage.rs
index d75d0a5..bded66b 100644
--- a/askama_derive/src/heritage.rs
+++ b/askama_derive/src/heritage.rs
@@ -35,6 +35,7 @@ impl Heritage<'_> {
 
 type BlockAncestry<'a> = HashMap<&'a str, Vec<(&'a Context<'a>, &'a BlockDef<'a>)>>;
 
+#[derive(Clone)]
 pub(crate) struct Context<'a> {
     pub(crate) nodes: &'a [Node<'a>],
     pub(crate) extends: Option<PathBuf>,
diff --git a/askama_derive/src/input.rs b/askama_derive/src/input.rs
index 57fbc04..54facdc 100644
--- a/askama_derive/src/input.rs
+++ b/askama_derive/src/input.rs
@@ -113,46 +113,74 @@ impl TemplateInput<'_> {
         let mut check = vec![(self.path.clone(), source)];
         while let Some((path, source)) = check.pop() {
             let parsed = Parsed::new(source, self.syntax)?;
-            for n in parsed.nodes() {
-                use Node::*;
-                match n {
-                    Extends(extends) => {
-                        let extends = self.config.find_template(extends.path, Some(&path))?;
-                        let dependency_path = (path.clone(), extends.clone());
-                        if dependency_graph.contains(&dependency_path) {
-                            return Err(format!(
-                                "cyclic dependency in graph {:#?}",
-                                dependency_graph
-                                    .iter()
-                                    .map(|e| format!("{:#?} --> {:#?}", e.0, e.1))
-                                    .collect::<Vec<String>>()
-                            )
-                            .into());
+
+            let mut top = true;
+            let mut nested = vec![parsed.nodes()];
+            while let Some(nodes) = nested.pop() {
+                for n in nodes {
+                    use Node::*;
+                    match n {
+                        Extends(extends) if top => {
+                            let extends = self.config.find_template(extends.path, Some(&path))?;
+                            let dependency_path = (path.clone(), extends.clone());
+                            if dependency_graph.contains(&dependency_path) {
+                                return Err(format!(
+                                    "cyclic dependency in graph {:#?}",
+                                    dependency_graph
+                                        .iter()
+                                        .map(|e| format!("{:#?} --> {:#?}", e.0, e.1))
+                                        .collect::<Vec<String>>()
+                                )
+                                .into());
+                            }
+                            dependency_graph.push(dependency_path);
+                            let source = get_template_source(&extends)?;
+                            check.push((extends, source));
                         }
-                        dependency_graph.push(dependency_path);
-                        let source = get_template_source(&extends)?;
-                        check.push((extends, source));
-                    }
-                    Import(import) => {
-                        let import = self.config.find_template(import.path, Some(&path))?;
-                        let source = get_template_source(&import)?;
-                        check.push((import, source));
+                        Macro(m) if top => {
+                            nested.push(&m.nodes);
+                        }
+                        Import(import) if top => {
+                            let import = self.config.find_template(import.path, Some(&path))?;
+                            let source = get_template_source(&import)?;
+                            check.push((import, source));
+                        }
+                        Include(include) => {
+                            let include = self.config.find_template(include.path, Some(&path))?;
+                            let source = get_template_source(&include)?;
+                            check.push((include, source));
+                        }
+                        BlockDef(b) => {
+                            nested.push(&b.nodes);
+                        }
+                        If(i) => {
+                            for cond in &i.branches {
+                                nested.push(&cond.nodes);
+                            }
+                        }
+                        Loop(l) => {
+                            nested.push(&l.body);
+                            nested.push(&l.else_nodes);
+                        }
+                        Match(m) => {
+                            for arm in &m.arms {
+                                nested.push(&arm.nodes);
+                            }
+                        }
+                        Lit(_)
+                        | Comment(_)
+                        | Expr(_, _)
+                        | Call(_)
+                        | Extends(_)
+                        | Let(_)
+                        | Import(_)
+                        | Macro(_)
+                        | Raw(_)
+                        | Continue(_)
+                        | Break(_) => {}
                     }
-                    If(_)
-                    | Loop(_)
-                    | Match(_)
-                    | BlockDef(_)
-                    | Include(_)
-                    | Lit(_)
-                    | Comment(_)
-                    | Expr(_, _)
-                    | Call(_)
-                    | Let(_)
-                    | Macro(_)
-                    | Raw(_)
-                    | Continue(_)
-                    | Break(_) => {}
                 }
+                top = false;
             }
             map.insert(path, parsed);
         }