Fixed whitespace issue when generating match (#399)

* Fixed #397

* Updated parser to ignore whitespace between match and when

* Updated test cases

* Updated Python script to generate match ws tests

* Added match ws tests

* Resolved rustfmt lint

3 files changed, 28 insertions, 23 deletions

diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index 6a185c1..320e55e 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -444,8 +444,8 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
                 Node::Cond(ref conds, ws) => {
                     self.write_cond(ctx, buf, conds, ws)?;
                 }
-                Node::Match(ws1, ref expr, inter, ref arms, ws2) => {
-                    self.write_match(ctx, buf, ws1, expr, inter, arms, ws2)?;
+                Node::Match(ws1, ref expr, ref arms, ws2) => {
+                    self.write_match(ctx, buf, ws1, expr, arms, ws2)?;
                 }
                 Node::Loop(ws1, ref var, ref iter, ref body, ws2) => {
                     self.write_loop(ctx, buf, ws1, var, iter, body, ws2)?;
@@ -561,23 +561,28 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         buf: &mut Buffer,
         ws1: WS,
         expr: &Expr,
-        inter: Option<&'a str>,
         arms: &'a [When],
         ws2: WS,
     ) -> Result<usize, CompileError> {
         self.flush_ws(ws1);
         let flushed = self.write_buf_writable(buf)?;
         let mut arm_sizes = Vec::new();
-        if let Some(inter) = inter {
-            if !inter.is_empty() {
-                self.next_ws = Some(inter);
-            }
-        }
 
         let expr_code = self.visit_expr_root(expr)?;
         buf.writeln(&format!("match &{} {{", expr_code))?;
-        for arm in arms {
+
+        let mut arm_size = 0;
+        for (i, arm) in arms.iter().enumerate() {
             let &(ws, ref variant, ref params, ref body) = arm;
+            self.handle_ws(ws);
+
+            if i > 0 {
+                arm_sizes.push(arm_size + self.write_buf_writable(buf)?);
+
+                buf.writeln("}")?;
+                self.locals.pop();
+            }
+
             self.locals.push();
             match *variant {
                 Some(ref param) => {
@@ -624,15 +629,17 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
                 }
             }
             buf.writeln(" => {")?;
-            self.handle_ws(ws);
-            let arm_size = self.handle(ctx, body, buf, AstLevel::Nested)?;
-            arm_sizes.push(arm_size + self.write_buf_writable(buf)?);
-            buf.writeln("}")?;
-            self.locals.pop();
+
+            arm_size = self.handle(ctx, body, buf, AstLevel::Nested)?;
         }
 
-        buf.writeln("}")?;
         self.handle_ws(ws2);
+        arm_sizes.push(arm_size + self.write_buf_writable(buf)?);
+        buf.writeln("}")?;
+        self.locals.pop();
+
+        buf.writeln("}")?;
+
         Ok(flushed + median(&mut arm_sizes))
     }
 
diff --git a/askama_shared/src/heritage.rs b/askama_shared/src/heritage.rs
index bc9e701..87f23e0 100644
--- a/askama_shared/src/heritage.rs
+++ b/askama_shared/src/heritage.rs
@@ -89,7 +89,7 @@ impl<'a> Context<'a> {
                     Node::Loop(_, _, _, nodes, _) => {
                         nested.push(nodes);
                     }
-                    Node::Match(_, _, _, arms, _) => {
+                    Node::Match(_, _, arms, _) => {
                         for (_, _, _, arm) in arms {
                             nested.push(arm);
                         }
diff --git a/askama_shared/src/parser.rs b/askama_shared/src/parser.rs
index 8210ed0..a40e967 100644
--- a/askama_shared/src/parser.rs
+++ b/askama_shared/src/parser.rs
@@ -19,7 +19,7 @@ pub enum Node<'a> {
     LetDecl(WS, Target<'a>),
     Let(WS, Target<'a>, Expr<'a>),
     Cond(Vec<(WS, Option<Expr<'a>>, Vec<Node<'a>>)>, WS),
-    Match(WS, Expr<'a>, Option<&'a str>, Vec<When<'a>>, WS),
+    Match(WS, Expr<'a>, Vec<When<'a>>, WS),
     Loop(WS, Target<'a>, Expr<'a>, Vec<Node<'a>>, WS),
     Extends(Expr<'a>),
     BlockDef(WS, &'a str, Vec<Node<'a>>, WS),
@@ -771,8 +771,8 @@ fn block_match<'a>(i: &'a [u8], s: &'a Syntax<'a>) -> IResult<&'a [u8], Node<'a>
         arms.push(arm);
     }
 
-    let inter = match inter {
-        Some(Node::Lit(lws, val, rws)) => {
+    match inter {
+        Some(Node::Lit(_, val, rws)) => {
             assert!(
                 val.is_empty(),
                 "only whitespace allowed between match and first when, found {}",
@@ -783,18 +783,16 @@ fn block_match<'a>(i: &'a [u8], s: &'a Syntax<'a>) -> IResult<&'a [u8], Node<'a>
                 "only whitespace allowed between match and first when, found {}",
                 rws
             );
-            Some(lws)
         }
-        None => None,
+        None => {}
         _ => panic!("only literals allowed between match and first when"),
-    };
+    }
 
     Ok((
         i,
         Node::Match(
             WS(pws1.is_some(), nws1.is_some()),
             expr,
-            inter,
             arms,
             WS(pws2.is_some(), nws2.is_some()),
         ),