Implement "if let" statement

2 files changed, 54 insertions, 17 deletions

diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index 24a4a71..fc94977 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -3,7 +3,8 @@ use crate::filters;
 use crate::heritage::{Context, Heritage};
 use crate::input::{Source, TemplateInput};
 use crate::parser::{
-    parse, Cond, Expr, MatchParameter, MatchParameters, MatchVariant, Node, Target, When, Ws,
+    parse, Cond, CondTest, Expr, MatchParameter, MatchParameters, MatchVariant, Node, Target, When,
+    Ws,
 };
 
 use proc_macro2::Span;
@@ -483,23 +484,38 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
                 self.locals.pop();
             }
 
+            self.locals.push();
             let mut arm_size = 0;
-            if let Some(expr) = cond {
+            if let Some(CondTest { target, expr }) = cond {
                 if i == 0 {
                     buf.write("if ");
                 } else {
                     buf.dedent()?;
                     buf.write("} else if ");
                 }
-                // The following syntax `*(&(...) as &bool)` is used to
-                // trigger Rust's automatic dereferencing, to coerce
-                // e.g. `&&&&&bool` to `bool`. First `&(...) as &bool`
-                // coerces e.g. `&&&bool` to `&bool`. Then `*(&bool)`
-                // finally dereferences it to `bool`.
-                buf.write("*(&(");
-                let expr_code = self.visit_expr_root(expr)?;
-                buf.write(&expr_code);
-                buf.write(") as &bool)");
+
+                if let Some((variant, params)) = target {
+                    let mut expr_buf = Buffer::new(0);
+                    self.visit_expr(&mut expr_buf, expr)?;
+                    buf.write("let ");
+                    self.visit_match_variant(buf, variant);
+                    if let Some(params) = params {
+                        self.visit_match_params(buf, params);
+                    }
+                    buf.write(" = &(");
+                    buf.write(&expr_buf.buf);
+                    buf.write(")");
+                } else {
+                    // The following syntax `*(&(...) as &bool)` is used to
+                    // trigger Rust's automatic dereferencing, to coerce
+                    // e.g. `&&&&&bool` to `bool`. First `&(...) as &bool`
+                    // coerces e.g. `&&&bool` to `&bool`. Then `*(&bool)`
+                    // finally dereferences it to `bool`.
+                    buf.write("*(&(");
+                    let expr_code = self.visit_expr_root(expr)?;
+                    buf.write(&expr_code);
+                    buf.write(") as &bool)");
+                }
             } else {
                 buf.dedent()?;
                 buf.write("} else");
@@ -507,7 +523,6 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
             }
 
             buf.writeln(" {")?;
-            self.locals.push();
 
             arm_size += self.handle(ctx, nodes, buf, AstLevel::Nested)?;
             arm_sizes.push(arm_size);
diff --git a/askama_shared/src/parser.rs b/askama_shared/src/parser.rs
index ea13738..3c5a86c 100644
--- a/askama_shared/src/parser.rs
+++ b/askama_shared/src/parser.rs
@@ -18,7 +18,7 @@ pub enum Node<'a> {
     Call(Ws, Option<&'a str>, &'a str, Vec<Expr<'a>>),
     LetDecl(Ws, Target<'a>),
     Let(Ws, Target<'a>, Expr<'a>),
-    Cond(Vec<(Ws, Option<Expr<'a>>, Vec<Node<'a>>)>, Ws),
+    Cond(Vec<Cond<'a>>, Ws),
     Match(Ws, Expr<'a>, Vec<When<'a>>, Ws),
     Loop(Ws, Target<'a>, Expr<'a>, Vec<Node<'a>>, Ws),
     Extends(Expr<'a>),
@@ -144,7 +144,13 @@ pub enum Target<'a> {
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub struct Ws(pub bool, pub bool);
 
-pub type Cond<'a> = (Ws, Option<Expr<'a>>, Vec<Node<'a>>);
+pub type Cond<'a> = (Ws, Option<CondTest<'a>>, Vec<Node<'a>>);
+
+#[derive(Debug, PartialEq)]
+pub struct CondTest<'a> {
+    pub target: Option<(MatchVariant<'a>, Option<MatchParameters<'a>>)>,
+    pub expr: Expr<'a>,
+}
 
 fn ws<F, I, O, E>(mut inner: F) -> impl FnMut(I) -> IResult<I, O, E>
 where
@@ -724,9 +730,25 @@ fn block_call(i: &[u8]) -> IResult<&[u8], Node> {
     ))
 }
 
-fn cond_if(i: &[u8]) -> IResult<&[u8], Expr> {
-    let (i, (_, cond)) = tuple((ws(tag("if")), ws(expr_any)))(i)?;
-    Ok((i, cond))
+fn cond_if(i: &[u8]) -> IResult<&[u8], CondTest> {
+    let mut p = tuple((
+        ws(tag("if")),
+        opt(tuple((
+            ws(alt((tag("let"), tag("set")))),
+            ws(match_variant),
+            opt(alt((match_simple_parameters, match_named_parameters))),
+            ws(tag("=")),
+        ))),
+        ws(expr_any),
+    ));
+    let (i, (_, dest, expr)) = p(i)?;
+    Ok((
+        i,
+        CondTest {
+            target: dest.map(|(_, variant, params, _)| (variant, params)),
+            expr,
+        },
+    ))
 }
 
 fn cond_block<'a>(i: &'a [u8], s: &'a Syntax<'a>) -> IResult<&'a [u8], Cond<'a>> {