From b36bbfa91c864d177fb74b0e50bf0c68e67b6952 Mon Sep 17 00:00:00 2001
From: Windel Bouwman <windel@windel.nl>
Date: Sun, 1 Sep 2019 22:31:16 +0200
Subject: [PATCH] Add symboltable scope for comprehensions. Add _ast module
 nodes for comprehensions.

---
 compiler/src/compile.rs     |  4 ++++
 compiler/src/symboltable.rs | 31 ++++++++++++++++++++++++++++++-
 vm/src/stdlib/ast.rs        | 29 ++++++++++++++++++-----------
 3 files changed, 52 insertions(+), 12 deletions(-)

diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
index 7cdfa1a36..7194c8886 100644
--- a/compiler/src/compile.rs
+++ b/compiler/src/compile.rs
@@ -1782,6 +1782,7 @@ impl<O: OutputStream> Compiler<O> {
             line_number,
             name.clone(),
         ));
+        self.enter_scope();
 
         // Create empty object of proper type:
         match kind {
@@ -1891,6 +1892,9 @@ impl<O: OutputStream> Compiler<O> {
         // Fetch code for listcomp function:
         let code = self.pop_code_object();
 
+        // Pop scope
+        self.leave_scope();
+
         // List comprehension code:
         self.emit(Instruction::LoadConst {
             value: bytecode::Constant::Code {
diff --git a/compiler/src/symboltable.rs b/compiler/src/symboltable.rs
index 546e118a2..cb367494f 100644
--- a/compiler/src/symboltable.rs
+++ b/compiler/src/symboltable.rs
@@ -588,6 +588,23 @@ impl SymbolTableBuilder {
                 self.scan_expressions(elements, context)?;
             }
             Comprehension { kind, generators } => {
+                // Comprehensions are compiled as functions, so create a scope for them:
+                let scope_name = match **kind {
+                    ast::ComprehensionKind::GeneratorExpression { .. } => "genexpr",
+                    ast::ComprehensionKind::List { .. } => "listcomp",
+                    ast::ComprehensionKind::Set { .. } => "setcomp",
+                    ast::ComprehensionKind::Dict { .. } => "dictcomp",
+                };
+
+                self.enter_scope(
+                    scope_name,
+                    SymbolTableType::Function,
+                    expression.location.row(),
+                );
+
+                // Register the passed argument to the generator function as the name ".0"
+                self.register_name(".0", SymbolUsage::Parameter)?;
+
                 match **kind {
                     ast::ComprehensionKind::GeneratorExpression { ref element }
                     | ast::ComprehensionKind::List { ref element }
@@ -600,13 +617,25 @@ impl SymbolTableBuilder {
                     }
                 }
 
+                let mut is_first_generator = true;
                 for generator in generators {
                     self.scan_expression(&generator.target, &ExpressionContext::Store)?;
-                    self.scan_expression(&generator.iter, &ExpressionContext::Load)?;
+                    if is_first_generator {
+                        is_first_generator = false;
+                    } else {
+                        self.scan_expression(&generator.iter, &ExpressionContext::Load)?;
+                    }
+
                     for if_expr in &generator.ifs {
                         self.scan_expression(if_expr, &ExpressionContext::Load)?;
                     }
                 }
+
+                self.leave_scope();
+
+                // The first iterable is passed as an argument into the created function:
+                assert!(!generators.is_empty());
+                self.scan_expression(&generators[0].iter, &ExpressionContext::Load)?;
             }
             Call {
                 function,
diff --git a/vm/src/stdlib/ast.rs b/vm/src/stdlib/ast.rs
index 85c8fd8a4..a9d7b44c2 100644
--- a/vm/src/stdlib/ast.rs
+++ b/vm/src/stdlib/ast.rs
@@ -456,18 +456,25 @@ fn expression_to_ast(vm: &VirtualMachine, expression: &ast::Expression) -> PyRes
             let py_generators = map_ast(comprehension_to_ast, vm, generators)?;
 
             match kind.deref() {
-                ast::ComprehensionKind::GeneratorExpression { .. } => {
-                    node!(vm, GeneratorExp, {generators => py_generators})
-                }
-                ast::ComprehensionKind::List { .. } => {
-                    node!(vm, ListComp, {generators => py_generators})
-                }
-                ast::ComprehensionKind::Set { .. } => {
-                    node!(vm, SetComp, {generators => py_generators})
-                }
-                ast::ComprehensionKind::Dict { .. } => {
-                    node!(vm, DictComp, {generators => py_generators})
+                ast::ComprehensionKind::GeneratorExpression { element } => {
+                    node!(vm, GeneratorExp, {
+                        elt => expression_to_ast(vm, element)?,
+                        generators => py_generators
+                    })
                 }
+                ast::ComprehensionKind::List { element } => node!(vm, ListComp, {
+                    elt => expression_to_ast(vm, element)?,
+                    generators => py_generators
+                }),
+                ast::ComprehensionKind::Set { element } => node!(vm, SetComp, {
+                    elt => expression_to_ast(vm, element)?,
+                    generators => py_generators
+                }),
+                ast::ComprehensionKind::Dict { key, value } => node!(vm, DictComp, {
+                    key => expression_to_ast(vm, key)?,
+                    value => expression_to_ast(vm, value)?,
+                    generators => py_generators
+                }),
             }
         }
         Await { value } => {