Don't push something on the stack when starting a generator

coolreader18 · coolreader18 · commit 559123c61709 · 2020-12-23T17:40:33.000-06:00
diff --git a/bytecode/src/lib.rs b/bytecode/src/lib.rs
@@ -378,6 +378,7 @@ pub enum Instruction {
     },
     GetAIter,
     GetANext,
+    EndAsyncFor,
 
     /// Reverse order evaluation in MapAdd
     /// required to support named expressions of Python 3.8 in dict comprehension
@@ -902,6 +903,7 @@ impl Instruction {
             SetupAsyncWith { .. } => 0,
             GetAIter => 0,
             GetANext => 1,
+            EndAsyncFor => -1,
         }
     }
 
@@ -1049,6 +1051,7 @@ impl Instruction {
             GetAwaitable => w!(GetAwaitable),
             GetAIter => w!(GetAIter),
             GetANext => w!(GetANext),
+            EndAsyncFor => w!(EndAsyncFor),
             MapAdd { i } => w!(MapAdd, i),
         }
     }
diff --git a/compiler/src/compile.rs b/compiler/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,24 +1401,27 @@ 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));
         self.compile_statements(body)?;
         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 {
             self.compile_statements(orelse)?;
         }
 
         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/compiler/src/ir.rs b/compiler/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];
diff --git a/vm/src/coroutine.rs b/vm/src/coroutine.rs
@@ -38,7 +38,6 @@ pub struct Coro {
     pub closed: AtomicCell<bool>,
     running: AtomicCell<bool>,
     exceptions: PyMutex<Vec<PyBaseExceptionRef>>,
-    started: AtomicCell<bool>,
     variant: Variant,
     name: PyMutex<PyStrRef>,
 }
@@ -50,7 +49,6 @@ impl Coro {
             closed: AtomicCell::new(false),
             running: AtomicCell::new(false),
             exceptions: PyMutex::new(vec![]),
-            started: AtomicCell::new(false),
             variant,
             name: PyMutex::new(name),
         }
@@ -67,34 +65,37 @@ impl Coro {
     where
         F: FnOnce(FrameRef) -> PyResult<ExecutionResult>,
     {
-        self.running.store(true);
-        let curr_exception_stack_len = vm.exceptions.borrow().len();
-        vm.exceptions
-            .borrow_mut()
-            .append(&mut self.exceptions.lock());
+        if self.running.compare_exchange(false, true).is_err() {
+            return Err(vm.new_value_error(format!("{} already executing", self.variant.name())));
+        }
+        let curr_exception_stack_len;
+        {
+            let mut vm_excs = vm.exceptions.borrow_mut();
+            curr_exception_stack_len = vm_excs.len();
+            vm_excs.append(&mut self.exceptions.lock());
+        }
         let result = vm.with_frame(self.frame.clone(), func);
-        std::mem::swap(
-            &mut *self.exceptions.lock(),
-            &mut vm
-                .exceptions
-                .borrow_mut()
-                .split_off(curr_exception_stack_len),
-        );
+        self.exceptions
+            .lock()
+            .extend(vm.exceptions.borrow_mut().drain(curr_exception_stack_len..));
         self.running.store(false);
-        self.started.store(true);
         result
     }
 
     pub fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         if self.closed.load() {
             return Err(vm.new_exception_empty(self.variant.stop_iteration(vm)));
         }
-        if !self.started.load() && !vm.is_none(&value) {
+        let value = if self.frame.lasti() > 0 {
+            Some(value)
+        } else if !vm.is_none(&value) {
             return Err(vm.new_type_error(format!(
                 "can't send non-None value to a just-started {}",
                 self.variant.name()
             )));
-        }
+        } else {
+            None
+        };
         let result = self.run_with_context(vm, |f| f.resume(value, vm));
         self.maybe_close(&result);
         match result {
@@ -155,9 +156,6 @@ impl Coro {
         }
     }
 
-    pub fn started(&self) -> bool {
-        self.started.load()
-    }
     pub fn running(&self) -> bool {
         self.running.load()
     }
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -254,11 +254,13 @@ impl FrameRef {
 
     pub(crate) fn resume(
         &self,
-        value: PyObjectRef,
+        value: Option<PyObjectRef>,
         vm: &VirtualMachine,
     ) -> PyResult<ExecutionResult> {
         self.with_exec(|mut exec| {
-            exec.push_value(value);
+            if let Some(value) = value {
+                exec.push_value(value)
+            }
             exec.run(vm)
         })
     }
@@ -811,6 +813,15 @@ impl ExecutingFrame<'_> {
                 self.push_value(awaitable);
                 Ok(None)
             }
+            bytecode::Instruction::EndAsyncFor => {
+                let exc = self.pop_value();
+                if exc.isinstance(&vm.ctx.exceptions.stop_async_iteration) {
+                    vm.pop_exception().expect("Should have exception in stack");
+                    Ok(None)
+                } else {
+                    Err(exc.downcast().unwrap())
+                }
+            }
             bytecode::Instruction::ForIter { target } => self.execute_for_iter(vm, *target),
             bytecode::Instruction::MakeFunction(flags) => self.execute_make_function(vm, *flags),
             bytecode::Instruction::CallFunctionPositional { nargs } => {

Original file line number	Diff line number	Diff line change
`@@ -378,6 +378,7 @@ pub enum Instruction {`
`378`	`378`	`},`
`379`	`379`	`GetAIter,`
`380`	`380`	`GetANext,`
	`381`	`+ EndAsyncFor,`
`381`	`382`
`382`	`383`	`/// Reverse order evaluation in MapAdd`
`383`	`384`	`/// required to support named expressions of Python 3.8 in dict comprehension`
`@@ -902,6 +903,7 @@ impl Instruction {`
`902`	`903`	`SetupAsyncWith { .. } => 0,`
`903`	`904`	`GetAIter => 0,`
`904`	`905`	`GetANext => 1,`
	`906`	`+ EndAsyncFor => -1,`
`905`	`907`	`}`
`906`	`908`	`}`
`907`	`909`
`@@ -1049,6 +1051,7 @@ impl Instruction {`
`1049`	`1051`	`GetAwaitable => w!(GetAwaitable),`
`1050`	`1052`	`GetAIter => w!(GetAIter),`
`1051`	`1053`	`GetANext => w!(GetANext),`
	`1054`	`+ EndAsyncFor => w!(EndAsyncFor),`
`1052`	`1055`	`MapAdd { i } => w!(MapAdd, i),`
`1053`	`1056`	`}`
`1054`	`1057`	`}`