implement async for function in compiler

HyeockJinKim · youknowone · commit ec1431819fda · 2024-04-23T18:25:47.000+09:00
diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
@@ -2654,11 +2654,8 @@ impl Compiler {
         }
 
         let mut loop_labels = vec![];
+        let mut is_async = false;
         for generator in generators {
-            if generator.is_async {
-                unimplemented!("async for comprehensions");
-            }
-
             let loop_block = self.new_block();
             let after_block = self.new_block();
 
@@ -2670,18 +2667,32 @@ impl Compiler {
                 self.compile_expression(&generator.iter)?;
 
                 // Get iterator / turn item into an iterator
-                emit!(self, Instruction::GetIter);
+                if generator.is_async {
+                    emit!(self, Instruction::GetAIter);
+                } else {
+                    emit!(self, Instruction::GetIter);
+                }
             }
 
             loop_labels.push((loop_block, after_block));
-
             self.switch_to_block(loop_block);
-            emit!(
-                self,
-                Instruction::ForIter {
-                    target: after_block,
-                }
-            );
+            if generator.is_async {
+                is_async = true;
+                emit!(self, Instruction::SetupFinally {
+                    handler: after_block,
+                });
+                emit!(self, Instruction::GetANext);
+                self.emit_constant(ConstantData::None);
+                emit!(self, Instruction::YieldFrom);
+                emit!(self, Instruction::PopBlock);
+            } else {
+                emit!(
+                    self,
+                    Instruction::ForIter {
+                        target: after_block,
+                    }
+                );
+            }
 
             self.compile_store(&generator.target)?;
 
@@ -2699,6 +2710,9 @@ impl Compiler {
 
             // End of for loop:
             self.switch_to_block(after_block);
+            if is_async {
+                emit!(self, Instruction::EndAsyncFor);
+            }
         }
 
         if return_none {
@@ -2735,10 +2749,19 @@ impl Compiler {
         self.compile_expression(&generators[0].iter)?;
 
         // Get iterator / turn item into an iterator
-        emit!(self, Instruction::GetIter);
+        if is_async {
+            emit!(self, Instruction::GetAIter);
+        } else {
+            emit!(self, Instruction::GetIter);
+        };
 
         // Call just created <listcomp> function:
         emit!(self, Instruction::CallFunctionPositional { nargs: 1 });
+        if is_async {
+            emit!(self, Instruction::GetAwaitable);
+            self.emit_constant(ConstantData::None);
+            emit!(self, Instruction::YieldFrom);
+        }
         Ok(())
     }
 
diff --git a/compiler/core/src/bytecode.rs b/compiler/core/src/bytecode.rs
@@ -1268,7 +1268,7 @@ impl Instruction {
             }
             GetAIter => 0,
             GetANext => 1,
-            EndAsyncFor => -2,
+            EndAsyncFor => -1,
             ExtendedArg => 0,
         }
     }
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -1000,7 +1000,6 @@ impl ExecutingFrame<'_> {
             }
             bytecode::Instruction::EndAsyncFor => {
                 let exc = self.pop_value();
-                self.pop_value(); // async iterator we were calling __anext__ on
                 if exc.fast_isinstance(vm.ctx.exceptions.stop_async_iteration) {
                     vm.take_exception().expect("Should have exception in stack");
                     Ok(None)

Original file line number	Diff line number	Diff line change
`@@ -1268,7 +1268,7 @@ impl Instruction {`
`1268`	`1268`	`}`
`1269`	`1269`	`GetAIter => 0,`
`1270`	`1270`	`GetANext => 1,`
`1271`		`- EndAsyncFor => -2,`
	`1271`	`+ EndAsyncFor => -1,`
`1272`	`1272`	`ExtendedArg => 0,`
`1273`	`1273`	`}`
`1274`	`1274`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1000,7 +1000,6 @@ impl ExecutingFrame<'_> {`
`1000`	`1000`	`}`
`1001`	`1001`	`bytecode::Instruction::EndAsyncFor => {`
`1002`	`1002`	`let exc = self.pop_value();`
`1003`		`- self.pop_value(); // async iterator we were calling __anext__ on`
`1004`	`1003`	`if exc.fast_isinstance(vm.ctx.exceptions.stop_async_iteration) {`
`1005`	`1004`	`vm.take_exception().expect("Should have exception in stack");`
`1006`	`1005`	`Ok(None)`