parser: add type for `Node::BlockDef`

4 files changed, 58 insertions, 64 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 60ffa27..beb3c49 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -650,8 +650,8 @@ impl<'a> Generator<'a> {
                 Node::Loop(ref loop_block) => {
                     size_hint += self.write_loop(ctx, buf, loop_block)?;
                 }
-                Node::BlockDef(ws1, name, _, ws2) => {
-                    size_hint += self.write_block(buf, Some(name), Ws(ws1.0, ws2.1))?;
+                Node::BlockDef(ref b) => {
+                    size_hint += self.write_block(buf, Some(b.name), Ws(b.ws1.0, b.ws2.1))?;
                 }
                 Node::Include(ws, path) => {
                     size_hint += self.handle_include(ctx, buf, ws, path)?;
@@ -1163,26 +1163,18 @@ impl<'a> Generator<'a> {
         // Get the block definition from the heritage chain
         let heritage = self
             .heritage
-            .as_ref()
             .ok_or_else(|| CompileError::from("no block ancestors available"))?;
-        let (ctx, def) = heritage.blocks[cur.0].get(cur.1).ok_or_else(|| {
+        let (ctx, def) = *heritage.blocks[cur.0].get(cur.1).ok_or_else(|| {
             CompileError::from(match name {
                 None => format!("no super() block found for block '{}'", cur.0),
                 Some(name) => format!("no block found for name '{name}'"),
             })
         })?;
 
-        // Get the nodes and whitespace suppression data from the block definition
-        let (ws1, nodes, ws2) = if let Node::BlockDef(ws1, _, nodes, ws2) = def {
-            (ws1, nodes, ws2)
-        } else {
-            unreachable!()
-        };
-
         // Handle inner whitespace suppression spec and process block nodes
-        self.prepare_ws(*ws1);
+        self.prepare_ws(def.ws1);
         self.locals.push();
-        let size_hint = self.handle(ctx, nodes, buf, AstLevel::Block)?;
+        let size_hint = self.handle(ctx, &def.nodes, buf, AstLevel::Block)?;
 
         if !self.locals.is_current_empty() {
             // Need to flush the buffer before popping the variable stack
@@ -1190,7 +1182,7 @@ impl<'a> Generator<'a> {
         }
 
         self.locals.pop();
-        self.flush_ws(*ws2);
+        self.flush_ws(def.ws2);
 
         // Restore original block context and set whitespace suppression for
         // succeeding whitespace according to the outer WS spec
diff --git a/askama_derive/src/heritage.rs b/askama_derive/src/heritage.rs
index fbbb71f..2a7ef43 100644
--- a/askama_derive/src/heritage.rs
+++ b/askama_derive/src/heritage.rs
@@ -3,7 +3,7 @@ use std::path::{Path, PathBuf};
 
 use crate::config::Config;
 use crate::CompileError;
-use parser::{Loop, Macro, Match, Node};
+use parser::{BlockDef, Loop, Macro, Match, Node};
 
 pub(crate) struct Heritage<'a> {
     pub(crate) root: &'a Context<'a>,
@@ -32,12 +32,12 @@ impl Heritage<'_> {
     }
 }
 
-type BlockAncestry<'a> = HashMap<&'a str, Vec<(&'a Context<'a>, &'a Node<'a>)>>;
+type BlockAncestry<'a> = HashMap<&'a str, Vec<(&'a Context<'a>, &'a BlockDef<'a>)>>;
 
 pub(crate) struct Context<'a> {
     pub(crate) nodes: &'a [Node<'a>],
     pub(crate) extends: Option<PathBuf>,
-    pub(crate) blocks: HashMap<&'a str, &'a Node<'a>>,
+    pub(crate) blocks: HashMap<&'a str, &'a BlockDef<'a>>,
     pub(crate) macros: HashMap<&'a str, &'a Macro<'a>>,
     pub(crate) imports: HashMap<&'a str, PathBuf>,
 }
@@ -49,7 +49,7 @@ impl Context<'_> {
         nodes: &'n [Node<'n>],
     ) -> Result<Context<'n>, CompileError> {
         let mut extends = None;
-        let mut blocks = Vec::new();
+        let mut blocks = HashMap::new();
         let mut macros = HashMap::new();
         let mut imports = HashMap::new();
         let mut nested = vec![nodes];
@@ -76,11 +76,9 @@ impl Context<'_> {
                             "extends, macro or import blocks not allowed below top level".into(),
                         );
                     }
-                    def @ Node::BlockDef(_, _, _, _) => {
-                        blocks.push(def);
-                        if let Node::BlockDef(_, _, nodes, _) = def {
-                            nested.push(nodes);
-                        }
+                    Node::BlockDef(b) => {
+                        blocks.insert(b.name, b);
+                        nested.push(&b.nodes);
                     }
                     Node::Cond(branches, _) => {
                         for cond in branches {
@@ -106,17 +104,6 @@ impl Context<'_> {
             top = false;
         }
 
-        let blocks: HashMap<_, _> = blocks
-            .iter()
-            .map(|def| {
-                if let Node::BlockDef(_, name, _, _) = def {
-                    (*name, *def)
-                } else {
-                    unreachable!()
-                }
-            })
-            .collect();
-
         Ok(Context {
             nodes,
             extends,
diff --git a/askama_parser/src/lib.rs b/askama_parser/src/lib.rs
index 74004cb..d095f65 100644
--- a/askama_parser/src/lib.rs
+++ b/askama_parser/src/lib.rs
@@ -16,7 +16,7 @@ use nom::{error_position, AsChar, IResult, InputTakeAtPosition};
 
 pub use self::expr::Expr;
 pub use self::node::{
-    Call, Cond, CondTest, Import, Loop, Macro, Match, Node, Target, When, Whitespace, Ws,
+    BlockDef, Call, Cond, CondTest, Import, Loop, Macro, Match, Node, Target, When, Whitespace, Ws,
 };
 
 mod expr;
diff --git a/askama_parser/src/node.rs b/askama_parser/src/node.rs
index 277d60f..a8325ef 100644
--- a/askama_parser/src/node.rs
+++ b/askama_parser/src/node.rs
@@ -26,7 +26,7 @@ pub enum Node<'a> {
     Match(Match<'a>),
     Loop(Loop<'a>),
     Extends(&'a str),
-    BlockDef(Ws, &'a str, Vec<Node<'a>>, Ws),
+    BlockDef(BlockDef<'a>),
     Include(Ws, &'a str),
     Import(Import<'a>),
     Macro(Macro<'a>),
@@ -77,7 +77,7 @@ impl<'a> Node<'a> {
                 Self::extends,
                 Self::include,
                 map(Import::parse, Self::Import),
-                |i| Self::block(i, s),
+                map(|i| BlockDef::parse(i, s), Self::BlockDef),
                 map(|i| Macro::parse(i, s), Self::Macro),
                 |i| Self::raw(i, s),
                 |i| Self::r#break(i, s),
@@ -224,33 +224,6 @@ impl<'a> Node<'a> {
         Ok((i, Self::Include(Ws(pws, nws), name)))
     }
 
-    fn block(i: &'a str, s: &State<'_>) -> IResult<&'a str, Self> {
-        let mut start = tuple((
-            opt(Whitespace::parse),
-            ws(keyword("block")),
-            cut(tuple((ws(identifier), opt(Whitespace::parse), |i| {
-                s.tag_block_end(i)
-            }))),
-        ));
-        let (i, (pws1, _, (name, nws1, _))) = start(i)?;
-
-        let mut end = cut(tuple((
-            |i| Self::many(i, s),
-            cut(tuple((
-                |i| s.tag_block_start(i),
-                opt(Whitespace::parse),
-                ws(keyword("endblock")),
-                cut(tuple((opt(ws(keyword(name))), opt(Whitespace::parse)))),
-            ))),
-        )));
-        let (i, (contents, (_, pws2, _, (_, nws2)))) = end(i)?;
-
-        Ok((
-            i,
-            Self::BlockDef(Ws(pws1, nws1), name, contents, Ws(pws2, nws2)),
-        ))
-    }
-
     fn raw(i: &'a str, s: &State<'_>) -> IResult<&'a str, Self> {
         let endraw = tuple((
             |i| s.tag_block_start(i),
@@ -780,6 +753,48 @@ impl<'a> Match<'a> {
     }
 }
 
+#[derive(Debug, PartialEq)]
+pub struct BlockDef<'a> {
+    pub ws1: Ws,
+    pub name: &'a str,
+    pub nodes: Vec<Node<'a>>,
+    pub ws2: Ws,
+}
+
+impl<'a> BlockDef<'a> {
+    fn parse(i: &'a str, s: &State<'_>) -> IResult<&'a str, Self> {
+        let mut start = tuple((
+            opt(Whitespace::parse),
+            ws(keyword("block")),
+            cut(tuple((ws(identifier), opt(Whitespace::parse), |i| {
+                s.tag_block_end(i)
+            }))),
+        ));
+        let (i, (pws1, _, (name, nws1, _))) = start(i)?;
+
+        let mut end = cut(tuple((
+            |i| Node::many(i, s),
+            cut(tuple((
+                |i| s.tag_block_start(i),
+                opt(Whitespace::parse),
+                ws(keyword("endblock")),
+                cut(tuple((opt(ws(keyword(name))), opt(Whitespace::parse)))),
+            ))),
+        )));
+        let (i, (nodes, (_, pws2, _, (_, nws2)))) = end(i)?;
+
+        Ok((
+            i,
+            BlockDef {
+                ws1: Ws(pws1, nws1),
+                name,
+                nodes,
+                ws2: Ws(pws2, nws2),
+            },
+        ))
+    }
+}
+
 /// First field is "minus/plus sign was used on the left part of the item".
 ///
 /// Second field is "minus/plus sign was used on the right part of the item".