From ec1431819fda2269485c3b8f1dd33ef93c303119 Mon Sep 17 00:00:00 2001
From: HyeockJinKim <kherootz@gmail.com>
Date: Wed, 25 Aug 2021 18:58:48 +0900
Subject: [PATCH] implement async for function in compiler

---
 compiler/codegen/src/compile.rs | 49 ++++++++++++++++++++++++---------
 compiler/core/src/bytecode.rs   |  2 +-
 vm/src/frame.rs                 |  1 -
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
index bf0fbe1dec..34e61339ef 100644
--- 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
index 8ada60a9e6..f468b0f75e 100644
--- 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
index 49307541fb..f818d35d50 100644
--- 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)