implement async for function in compiler

HyeockJinKim · youknowone · commit 4e9cc31e9f2b · 2022-08-09T22:20:26.000+09:00
diff --git a/bytecode/src/lib.rs b/bytecode/src/lib.rs
@@ -1037,7 +1037,7 @@ impl Instruction {
             }
             GetAIter => 0,
             GetANext => 1,
-            EndAsyncFor => -2,
+            EndAsyncFor => -1,
         }
     }
 
diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
@@ -2460,37 +2460,48 @@ 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();
 
             // Setup for loop:
-            self.emit(Instruction::SetupLoop {
-                break_target: after_block,
-            });
+            if !generator.is_async {
+                self.emit(Instruction::SetupLoop {
+                    break_target: after_block,
+                });
+            }
 
             if loop_labels.is_empty() {
                 // Load iterator onto stack (passed as first argument):
                 self.emit(Instruction::LoadFast(arg0));
             } else {
                 // Evaluate iterated item:
                 self.compile_expression(&generator.iter)?;
-
-                // Get iterator / turn item into an iterator
-                self.emit(Instruction::GetIter);
+                if generator.is_async {
+                    self.emit(Instruction::GetAIter);
+                } else {
+                    // Get iterator / turn item into an iterator
+                    self.emit(Instruction::GetIter);
+                }
             }
 
             loop_labels.push((loop_block, after_block));
-
             self.switch_to_block(loop_block);
-            self.emit(Instruction::ForIter {
-                target: after_block,
-            });
-
+            if generator.is_async {
+                is_async = true;
+                self.emit(Instruction::SetupFinally {
+                    handler: after_block,
+                });
+                self.emit(Instruction::GetANext);
+                self.emit_constant(ConstantData::None);
+                self.emit(Instruction::YieldFrom);
+                self.emit(Instruction::PopBlock);
+            } else {
+                self.emit(Instruction::ForIter {
+                    target: after_block,
+                });
+            }
             self.compile_store(&generator.target)?;
 
             // Now evaluate the ifs:
@@ -2507,7 +2518,11 @@ impl Compiler {
 
             // End of for loop:
             self.switch_to_block(after_block);
-            self.emit(Instruction::PopBlock);
+            if is_async {
+                self.emit(Instruction::EndAsyncFor);
+            } else {
+                self.emit(Instruction::PopBlock);
+            }
         }
 
         if return_none {
@@ -2544,10 +2559,19 @@ impl Compiler {
         self.compile_expression(&generators[0].iter)?;
 
         // Get iterator / turn item into an iterator
-        self.emit(Instruction::GetIter);
+        if is_async {
+            self.emit(Instruction::GetAIter);
+        } else {
+            self.emit(Instruction::GetIter);
+        }
 
         // Call just created <listcomp> function:
         self.emit(Instruction::CallFunctionPositional { nargs: 1 });
+        if is_async {
+            self.emit(Instruction::GetAwaitable);
+            self.emit_constant(ConstantData::None);
+            self.emit(Instruction::YieldFrom);
+        }
         Ok(())
     }
 
@@ -2563,7 +2587,9 @@ impl Compiler {
                 // "generator_stop" => {}
                 "annotations" => self.future_annotations = true,
                 other => {
-                    return Err(self.error(CompileErrorType::InvalidFutureFeature(other.to_owned())))
+                    return Err(
+                        self.error(CompileErrorType::InvalidFutureFeature(other.to_owned()))
+                    );
                 }
             }
         }
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -893,7 +893,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
`@@ -1037,7 +1037,7 @@ impl Instruction {`
`1037`	`1037`	`}`
`1038`	`1038`	`GetAIter => 0,`
`1039`	`1039`	`GetANext => 1,`
`1040`		`- EndAsyncFor => -2,`
	`1040`	`+ EndAsyncFor => -1,`
`1041`	`1041`	`}`
`1042`	`1042`	`}`
`1043`	`1043`
Original file line number	Diff line number	Diff line change
`@@ -893,7 +893,6 @@ impl ExecutingFrame<'_> {`
`893`	`893`	`}`
`894`	`894`	`bytecode::Instruction::EndAsyncFor => {`
`895`	`895`	`let exc = self.pop_value();`
`896`		`- self.pop_value(); // async iterator we were calling __anext__ on`
`897`	`896`	`if exc.fast_isinstance(vm.ctx.exceptions.stop_async_iteration) {`
`898`	`897`	`vm.take_exception().expect("Should have exception in stack");`
`899`	`898`	`Ok(None)`