implement async for function in compiler

HyeockJinKim · HyeockJinKim · commit 78d80db957f6 · 2021-09-05T17:30:07.000+09:00
diff --git a/bytecode/src/lib.rs b/bytecode/src/lib.rs
@@ -1061,7 +1061,7 @@ impl Instruction {
             }
             GetAIter => 0,
             GetANext => 1,
-            EndAsyncFor => -2,
+            EndAsyncFor => -1,
         }
     }
 
diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
@@ -33,6 +33,7 @@ enum CallType {
     Keyword { nargs: u32 },
     Ex { has_kwargs: bool },
 }
+
 impl CallType {
     fn normal_call(self) -> Instruction {
         match self {
@@ -69,6 +70,7 @@ pub struct CompileOpts {
     /// not emit assert statements
     pub optimize: u8,
 }
+
 impl Default for CompileOpts {
     fn default() -> Self {
         CompileOpts { optimize: 0 }
@@ -1608,7 +1610,7 @@ impl Compiler {
                         "starred assignment target must be in a list or tuple".to_owned(),
                     ),
                     _ => CompileErrorType::Assign(target.node.name()),
-                }))
+                }));
             }
         }
 
@@ -1907,10 +1909,10 @@ impl Compiler {
             YieldFrom { value } => {
                 match self.ctx.func {
                     FunctionContext::NoFunction => {
-                        return Err(self.error(CompileErrorType::InvalidYieldFrom))
+                        return Err(self.error(CompileErrorType::InvalidYieldFrom));
                     }
                     FunctionContext::AsyncFunction => {
-                        return Err(self.error(CompileErrorType::AsyncYieldFrom))
+                        return Err(self.error(CompileErrorType::AsyncYieldFrom));
                     }
                     FunctionContext::Function => {}
                 }
@@ -2295,37 +2297,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:
@@ -2342,7 +2355,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 {
@@ -2379,10 +2396,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(())
     }
 
@@ -2398,7 +2424,9 @@ impl Compiler {
                 // "generator_stop" => {}
                 // "annotations" => {}
                 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
@@ -902,7 +902,6 @@ impl ExecutingFrame<'_> {
             }
             bytecode::Instruction::EndAsyncFor => {
                 let exc = self.pop_value();
-                self.pop_value(); // async iterator we were calling __anext__ on
                 if exc.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
`@@ -1061,7 +1061,7 @@ impl Instruction {`
`1061`	`1061`	`}`
`1062`	`1062`	`GetAIter => 0,`
`1063`	`1063`	`GetANext => 1,`
`1064`		`- EndAsyncFor => -2,`
	`1064`	`+ EndAsyncFor => -1,`
`1065`	`1065`	`}`
`1066`	`1066`	`}`
`1067`	`1067`
Original file line number	Diff line number	Diff line change
`@@ -902,7 +902,6 @@ impl ExecutingFrame<'_> {`
`902`	`902`	`}`
`903`	`903`	`bytecode::Instruction::EndAsyncFor => {`
`904`	`904`	`let exc = self.pop_value();`
`905`		`- self.pop_value(); // async iterator we were calling __anext__ on`
`906`	`905`	`if exc.isinstance(&vm.ctx.exceptions.stop_async_iteration) {`
`907`	`906`	`vm.take_exception().expect("Should have exception in stack");`
`908`	`907`	`Ok(None)`