Change target of for loops into single expression.

windelbouwman · windelbouwman · commit c9aa103cceba · 2018-10-17T13:12:04.000+02:00
diff --git a/parser/src/ast.rs b/parser/src/ast.rs
@@ -78,7 +78,7 @@ pub enum Statement {
         body: Vec<LocatedStatement>,
     },
     For {
-        target: Vec<Expression>,
+        target: Expression,
         iter: Vec<Expression>,
         body: Vec<LocatedStatement>,
         orelse: Option<Vec<LocatedStatement>>,
@@ -200,7 +200,7 @@ pub enum Expression {
 
 #[derive(Debug, PartialEq, Clone)]
 pub struct Comprehension {
-    pub target: Vec<Expression>, // TODO: should this be a non-vector?
+    pub target: Expression,
     pub iter: Expression,
     pub ifs: Vec<Expression>,
 }
diff --git a/parser/src/parser.rs b/parser/src/parser.rs
@@ -337,9 +337,9 @@ mod tests {
                     name: "x".to_string()
                 }),
                 generators: vec![ast::Comprehension {
-                    target: vec![ast::Expression::Identifier {
+                    target: ast::Expression::Identifier {
                         name: "y".to_string()
-                    },],
+                    },
                     iter: ast::Expression::Identifier {
                         name: "z".to_string()
                     },
@@ -348,4 +348,45 @@ mod tests {
             }
         );
     }
+
+    #[test]
+    fn test_parse_double_list_comprehension() {
+        let source = String::from("[x for y, y2 in z for a in b]");
+        let parse_ast = parse_expression(&source).unwrap();
+        assert_eq!(
+            parse_ast,
+            ast::Expression::ListComprehension {
+                element: Box::new(ast::Expression::Identifier {
+                    name: "x".to_string()
+                }),
+                generators: vec![
+                    ast::Comprehension {
+                        target: ast::Expression::Tuple {
+                            elements: vec![
+                                ast::Expression::Identifier {
+                                    name: "y".to_string()
+                                },
+                                ast::Expression::Identifier {
+                                    name: "y2".to_string()
+                                },
+                            ],
+                        },
+                        iter: ast::Expression::Identifier {
+                            name: "z".to_string()
+                        },
+                        ifs: vec![],
+                    },
+                    ast::Comprehension {
+                        target: ast::Expression::Identifier {
+                            name: "a".to_string()
+                        },
+                        iter: ast::Expression::Identifier {
+                            name: "b".to_string()
+                        },
+                        ifs: vec![],
+                    }
+                ],
+            }
+        );
+    }
 }
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
@@ -298,8 +298,11 @@ ForStatement: ast::LocatedStatement = {
       None => None,
     };
     ast::LocatedStatement {
-      location: loc,
-      node: ast::Statement::For { target: e, iter: t, body: s, orelse: or_else },
+          location: loc,
+          node: ast::Statement::For {
+            target: e,
+            iter: t, body: s, orelse: or_else
+          },
     }
   },
 };
@@ -622,8 +625,14 @@ DictEntry: (ast::Expression, ast::Expression) = {
     <e1: Test> ":" <e2: Test> => (e1, e2),
 };
 
-ExpressionList: Vec<ast::Expression> = {
-  <e: Comma<Expression>> => e,
+ExpressionList: ast::Expression = {
+    <e: Comma<Expression>> => {
+        if e.len() == 1 {
+            e.into_iter().next().unwrap()
+        } else {
+            ast::Expression::Tuple { elements: e }
+        }
+    },
 };
 
 #[inline]
@@ -653,12 +662,12 @@ CompFor: Vec<ast::Comprehension> = {
             },
             Some((Some(mut cphs), None)) => {
                 let mut res = vec![];
-                res.append(&mut cphs);
                 res.push(ast::Comprehension {
                     target: e,
                     iter: i,
                     ifs: vec![],
                 });
+                res.append(&mut cphs);
                 res
             },
             _ => {
diff --git a/tests/snippets/comprehensions.py b/tests/snippets/comprehensions.py
@@ -4,3 +4,9 @@
 y = [a*a+1 for a in x]
 assert y == [2, 5, 10]
 
+z = [(b, c) for b in x for c in y]
+# print(z)
+assert z == [
+    (1, 2), (1, 5), (1, 10),
+    (2, 2), (2, 5), (2, 10),
+    (3, 2), (3, 5), (3, 10)]
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
@@ -250,17 +250,7 @@ impl Compiler {
                 self.emit(Instruction::ForIter);
 
                 // Start of loop iteration, set targets:
-                // TODO: can we use compile_store here?
-                for t in target {
-                    match t {
-                        ast::Expression::Identifier { name } => {
-                            self.emit(Instruction::StoreName {
-                                name: name.to_string(),
-                            });
-                        }
-                        _ => panic!("Not impl"),
-                    }
-                }
+                self.compile_store(target);
 
                 // Body of loop:
                 self.compile_statements(body);
@@ -921,17 +911,7 @@ impl Compiler {
                 self.set_label(start_label);
                 self.emit(Instruction::ForIter);
 
-                // TODO: can we use compile_store here?
-                for t in &generator.target {
-                    match t {
-                        ast::Expression::Identifier { name } => {
-                            self.emit(Instruction::StoreName {
-                                name: name.to_string(),
-                            });
-                        }
-                        _ => panic!("Not impl"),
-                    }
-                }
+                self.compile_store(&generator.target);
 
                 // Evaluate element:
                 self.compile_expression(element);
