implement async for function

HyeockJinKim · HyeockJinKim · commit 4e5fef6d5f67 · 2021-09-02T00:15:00.000+09:00
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:
@@ -2379,10 +2392,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 +2420,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()))
+                    );
                 }
             }
         }
