From f3501f44cb490b112d20f7d0c8dca152136a5fd1 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Thu, 25 Apr 2024 14:10:00 +0900
Subject: [PATCH] Basic async for comprehension support

---
 compiler/codegen/src/compile.rs | 39 ++++++++++++++++++++-------------
 vm/src/frame.rs                 |  2 ++
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
index 0baa08b6b..b04821b80 100644
--- a/compiler/codegen/src/compile.rs
+++ b/compiler/codegen/src/compile.rs
@@ -2629,24 +2629,30 @@ impl Compiler {
         compile_element: &dyn Fn(&mut Self) -> CompileResult<()>,
     ) -> CompileResult<()> {
         let prev_ctx = self.ctx;
+        let is_async = generators.iter().any(|g| g.is_async);
 
         self.ctx = CompileContext {
             loop_data: None,
             in_class: prev_ctx.in_class,
-            func: FunctionContext::Function,
+            func: if is_async {
+                FunctionContext::AsyncFunction
+            } else {
+                FunctionContext::Function
+            },
         };
 
         // We must have at least one generator:
         assert!(!generators.is_empty());
 
+        let flags = bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED;
+        let flags = if is_async {
+            flags | bytecode::CodeFlags::IS_COROUTINE
+        } else {
+            flags
+        };
+
         // Create magnificent function <listcomp>:
-        self.push_output(
-            bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
-            1,
-            1,
-            0,
-            name.to_owned(),
-        );
+        self.push_output(flags, 1, 1, 0, name.to_owned());
         let arg0 = self.varname(".0")?;
 
         let return_none = init_collection.is_none();
@@ -2656,11 +2662,12 @@ impl Compiler {
         }
 
         let mut loop_labels = vec![];
-        let mut is_async = false;
         for generator in generators {
             let loop_block = self.new_block();
             let after_block = self.new_block();
 
+            // emit!(self, Instruction::SetupLoop);
+
             if loop_labels.is_empty() {
                 // Load iterator onto stack (passed as first argument):
                 emit!(self, Instruction::LoadFast(arg0));
@@ -2679,13 +2686,16 @@ impl Compiler {
             loop_labels.push((loop_block, after_block));
             self.switch_to_block(loop_block);
             if generator.is_async {
-                is_async = true;
-                emit!(self, Instruction::SetupExcept {
-                    handler: after_block,
-                });
+                emit!(
+                    self,
+                    Instruction::SetupExcept {
+                        handler: after_block,
+                    }
+                );
                 emit!(self, Instruction::GetANext);
                 self.emit_constant(ConstantData::None);
                 emit!(self, Instruction::YieldFrom);
+                self.compile_store(&generator.target)?;
                 emit!(self, Instruction::PopBlock);
             } else {
                 emit!(
@@ -2694,10 +2704,9 @@ impl Compiler {
                         target: after_block,
                     }
                 );
+                self.compile_store(&generator.target)?;
             }
 
-            self.compile_store(&generator.target)?;
-
             // Now evaluate the ifs:
             for if_condition in &generator.ifs {
                 self.compile_jump_if(if_condition, false, loop_block)?
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
index 03809d45d..f0d433c06 100644
--- a/vm/src/frame.rs
+++ b/vm/src/frame.rs
@@ -1944,6 +1944,7 @@ impl ExecutingFrame<'_> {
     #[track_caller]
     fn pop_block(&mut self) -> Block {
         let block = self.state.blocks.pop().expect("No more blocks to pop!");
+        // eprintln!("popped block: {:?} stack: {} truncate to {}", block.typ, self.state.stack.len(), block.level);
         #[cfg(debug_assertions)]
         if self.state.stack.len() < block.level {
             dbg!(&self);
@@ -2002,6 +2003,7 @@ impl ExecutingFrame<'_> {
     }
 
     #[inline]
+    #[track_caller]
     fn nth_value(&self, depth: u32) -> &PyObject {
         let stack = &self.state.stack;
         &stack[stack.len() - depth as usize - 1]