From 64fa9ccba96f63083cba56343e6b0f88d590a7b2 Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Mon, 21 Dec 2020 19:03:49 -0600
Subject: [PATCH] Don't push something on the stack when starting a generator

---
 src/compile.rs | 36 ++++++++++++------------------------
 src/ir.rs      |  5 +----
 2 files changed, 13 insertions(+), 28 deletions(-)

diff --git a/src/compile.rs b/src/compile.rs
index 3de5fc2..d9a1c3a 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -1376,9 +1376,8 @@ impl Compiler {
         // The thing iterated:
         self.compile_expression(iter)?;
 
-        if is_async {
+        let check_asynciter_block = if is_async {
             let check_asynciter_block = self.new_block();
-            let body_block = self.new_block();
 
             self.emit(Instruction::GetAIter);
 
@@ -1391,21 +1390,8 @@ impl Compiler {
             self.emit(Instruction::YieldFrom);
             self.compile_store(target)?;
             self.emit(Instruction::PopBlock);
-            self.emit(Instruction::Jump { target: body_block });
 
-            self.switch_to_block(check_asynciter_block);
-            self.emit(Instruction::Duplicate);
-            let stopasynciter = self.name("StopAsyncIteration");
-            self.emit(Instruction::LoadGlobal(stopasynciter));
-            self.emit(Instruction::CompareOperation {
-                op: bytecode::ComparisonOperator::ExceptionMatch,
-            });
-            self.emit(Instruction::JumpIfTrue { target: else_block });
-            self.emit(Instruction::Raise {
-                kind: bytecode::RaiseKind::Reraise,
-            });
-
-            self.switch_to_block(body_block);
+            Some(check_asynciter_block)
         } else {
             // Retrieve Iterator
             self.emit(Instruction::GetIter);
@@ -1415,7 +1401,8 @@ impl Compiler {
 
             // Start of loop iteration, set targets:
             self.compile_store(target)?;
-        }
+            None
+        };
 
         let was_in_loop = self.ctx.loop_data;
         self.ctx.loop_data = Some((for_block, after_block));
@@ -1423,6 +1410,11 @@ impl Compiler {
         self.ctx.loop_data = was_in_loop;
         self.emit(Instruction::Jump { target: for_block });
 
+        if let Some(check_asynciter_block) = check_asynciter_block {
+            self.switch_to_block(check_asynciter_block);
+            self.emit(Instruction::EndAsyncFor);
+        }
+
         self.switch_to_block(else_block);
         self.emit(Instruction::PopBlock);
         if let Some(orelse) = orelse {
@@ -1430,9 +1422,6 @@ impl Compiler {
         }
 
         self.switch_to_block(after_block);
-        if is_async {
-            self.emit(Instruction::Pop);
-        }
 
         Ok(())
     }
@@ -2169,14 +2158,13 @@ impl Compiler {
         );
         let arg0 = self.varname(".0");
 
+        // if this is a generator expression, we need to insert a LoadConst(None) before we return;
+        // the other kinds load their collection types below
         let is_genexpr = matches!(kind, ast::ComprehensionKind::GeneratorExpression { .. });
 
         // Create empty object of proper type:
         match kind {
-            ast::ComprehensionKind::GeneratorExpression { .. } => {
-                // pop the None that was passed to kickstart the generator
-                self.emit(Instruction::Pop);
-            }
+            ast::ComprehensionKind::GeneratorExpression { .. } => {}
             ast::ComprehensionKind::List { .. } => {
                 self.emit(Instruction::BuildList {
                     size: 0,
diff --git a/src/ir.rs b/src/ir.rs
index dde8962..177df64 100644
--- a/src/ir.rs
+++ b/src/ir.rs
@@ -165,10 +165,7 @@ impl CodeInfo {
         let mut maxdepth = 0u32;
         let mut stack = Vec::with_capacity(self.blocks.len());
         let mut startdepths = vec![u32::MAX; self.blocks.len()];
-        // if it's a generator or a coroutine, it starts with something already on the stack
-        startdepths[0] =
-            self.flags
-                .intersects(CodeFlags::IS_GENERATOR | CodeFlags::IS_COROUTINE) as u32;
+        startdepths[0] = 0;
         stack.push(Label(0));
         'process_blocks: while let Some(block) = stack.pop() {
             let mut depth = startdepths[block.0 as usize];