For loop patterns 🦚

src/eval/mod.rs

@@ -275,7 +275,7 @@ impl Eval for Spanned<&ExprBinary> {
 
 impl Spanned<&ExprBinary> {
     /// Apply a basic binary operation.
-    fn apply<F>(&self, ctx: &mut EvalContext, op: F) -> Value
+    fn apply<F>(self, ctx: &mut EvalContext, op: F) -> Value
     where
         F: FnOnce(Value, Value) -> Value,
     {
@@ -311,7 +311,7 @@ impl Spanned<&ExprBinary> {
     }
 
     /// Apply an assignment operation.
-    fn assign<F>(&self, ctx: &mut EvalContext, op: F) -> Value
+    fn assign<F>(self, ctx: &mut EvalContext, op: F) -> Value
     where
         F: FnOnce(Value, Value) -> Value,
     {
@@ -379,30 +379,56 @@ impl Eval for Spanned<&ExprFor> {
 
     fn eval(self, ctx: &mut EvalContext) -> Self::Output {
         let iter = self.v.iter.eval(ctx);
-        if let Value::Array(array) = iter {
-            let mut output = match self.v.body.v {
-                Expr::Template(_) => Value::Template(vec![]),
-                _ => Value::None,
-            };
+        let mut output = if let Expr::Template(_) = self.v.body.v {
+            Value::Template(vec![])
+        } else {
+            Value::None
+        };
 
-            for value in array {
-                ctx.scopes.define(self.v.pat.v.as_str(), value);
-                let value = self.v.body.eval(ctx);
+        macro_rules! iterate {
+            (for ($($binding:ident => $value:ident),*) in $iter:expr) => {
+                #[allow(unused_parens)]
+                for ($($value),*) in $iter {
+                    $(ctx.scopes.define($binding.as_str(), $value);)*
 
-                if let Value::Template(prev) = &mut output {
-                    if let Value::Template(new) = value {
-                        prev.extend(new);
+                    let value = self.v.body.eval(ctx);
+
+                    if let Value::Template(prev) = &mut output {
+                        if let Value::Template(new) = value {
+                            prev.extend(new);
+                        }
                     }
                 }
+
+                return output;
+            };
+        }
+
+        match (self.v.pat.v.clone(), iter) {
+            (ForPattern::Value(v), Value::Str(string)) => {
+                iterate!(for (v => value) in string.chars().map(|c| Value::Str(c.into())));
+            }
+            (ForPattern::Value(v), Value::Array(array)) => {
+                iterate!(for (v => value) in array.into_iter());
+            }
+            (ForPattern::Value(v), Value::Dict(dict)) => {
+                iterate!(for (v => value) in dict.into_iter().map(|p| p.1));
+            }
+            (ForPattern::KeyValue(k, v), Value::Dict(dict)) => {
+                iterate!(for (k => key, v => value) in dict.into_iter());
             }
 
-            return output;
-        } else if iter != Value::Error {
-            ctx.diag(error!(
+            (ForPattern::KeyValue(..), Value::Str(_))
+            | (ForPattern::KeyValue(..), Value::Array(_)) => {
+                ctx.diag(error!(self.v.pat.span, "mismatched pattern",));
+            }
+
+            (_, Value::Error) => {}
+            (_, iter) => ctx.diag(error!(
                 self.v.iter.span,
-                "expected array, found {}",
+                "cannot loop over {}",
                 iter.type_name(),
-            ));
+            )),
         }
 
         Value::Error

src/parse/mod.rs

@@ -423,7 +423,7 @@ fn expr_for(p: &mut Parser) -> Option<Expr> {
     p.assert(Token::For);
 
     let mut expr_for = None;
-    if let Some(pat) = p.span_if(ident) {
+    if let Some(pat) = p.span_if(for_pattern) {
         if p.expect(Token::In) {
             if let Some(iter) = p.span_if(expr) {
                 if let Some(body) = p.span_if(body) {
@@ -440,6 +440,17 @@ fn expr_for(p: &mut Parser) -> Option<Expr> {
     expr_for
 }
 
+/// Parse a for loop pattern.
+fn for_pattern(p: &mut Parser) -> Option<ForPattern> {
+    let first = ident(p)?;
+    if p.eat_if(Token::Comma) {
+        if let Some(second) = ident(p) {
+            return Some(ForPattern::KeyValue(first, second));
+        }
+    }
+    Some(ForPattern::Value(first))
+}
+
 /// Parse an identifier.
 fn ident(p: &mut Parser) -> Option<Ident> {
     match p.peek() {

src/syntax/expr.rs

@@ -528,7 +528,7 @@ impl Pretty for ExprIf {
 #[derive(Debug, Clone, PartialEq)]
 pub struct ExprFor {
     /// The pattern to assign to.
-    pub pat: Spanned<Ident>,
+    pub pat: Spanned<ForPattern>,
     /// The expression to iterate over.
     pub iter: SpanBox<Expr>,
     /// The expression to evaluate for each iteration.
@@ -538,10 +538,32 @@ pub struct ExprFor {
 impl Pretty for ExprFor {
     fn pretty(&self, p: &mut Printer) {
         p.push_str("#for ");
-        p.push_str(&self.pat.v);
+        self.pat.v.pretty(p);
         p.push_str(" #in ");
         self.iter.v.pretty(p);
         p.push_str(" ");
         self.body.v.pretty(p);
     }
 }
+
+/// A pattern in a for loop.
+#[derive(Debug, Clone, PartialEq)]
+pub enum ForPattern {
+    /// A value pattern: `#for v #in array`.
+    Value(Ident),
+    /// A key-value pattern: `#for k, v #in dict`.
+    KeyValue(Ident, Ident),
+}
+
+impl Pretty for ForPattern {
+    fn pretty(&self, p: &mut Printer) {
+        match self {
+            Self::Value(v) => p.push_str(&v),
+            Self::KeyValue(k, v) => {
+                p.push_str(&k);
+                p.push_str(", ");
+                p.push_str(&v);
+            }
+        }
+    }
+}

src/syntax/mod.rs

@@ -127,5 +127,6 @@ mod tests {
         roundtrip("#let x = 1 + 2");
         roundtrip("#if x [y] #else [z]");
         roundtrip("#for x #in y {z}");
+        roundtrip("#for k, x #in y {z}");
     }
 }