From 024cf5db042c18053f4e2ed2b6a20b20445e0ba4 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.

---
 src/compile.rs     |  4 ++++
 src/symboltable.rs | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/compile.rs b/src/compile.rs
index 7cdfa1a..7194c88 100644
--- a/src/compile.rs
+++ b/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/src/symboltable.rs b/src/symboltable.rs
index 546e118..cb36749 100644
--- a/src/symboltable.rs
+++ b/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,