From ae825125a5fa22f31a5179c90059bfcecd42545b Mon Sep 17 00:00:00 2001
From: coolreader18 <33094578+coolreader18@users.noreply.github.com>
Date: Wed, 16 Oct 2019 22:28:49 -0500
Subject: [PATCH] Add async with, reorganize how with blocks work

---
 src/compile.rs | 79 +++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 62 insertions(+), 17 deletions(-)

diff --git a/src/compile.rs b/src/compile.rs
index f467cc5..fae828d 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -386,27 +386,72 @@ impl<O: OutputStream> Compiler<O> {
                 body,
             } => {
                 if *is_async {
-                    unimplemented!("async with");
-                } else {
-                    let end_label = self.new_label();
-                    for item in items {
-                        self.compile_expression(&item.context_expr)?;
-                        self.emit(Instruction::SetupWith { end: end_label });
-                        match &item.optional_vars {
-                            Some(var) => {
-                                self.compile_store(var)?;
+                    let end_labels = items
+                        .iter()
+                        .map(|item| {
+                            let end_label = self.new_label();
+                            self.compile_expression(&item.context_expr)?;
+                            self.emit(Instruction::BeforeAsyncWith);
+                            self.emit(Instruction::GetAwaitable);
+                            self.emit(Instruction::LoadConst {
+                                value: bytecode::Constant::None,
+                            });
+                            self.emit(Instruction::YieldFrom);
+                            self.emit(Instruction::SetupAsyncWith { end: end_label });
+                            match &item.optional_vars {
+                                Some(var) => {
+                                    self.compile_store(var)?;
+                                }
+                                None => {
+                                    self.emit(Instruction::Pop);
+                                }
                             }
-                            None => {
-                                self.emit(Instruction::Pop);
-                            }
-                        }
-                    }
+                            Ok(end_label)
+                        })
+                        .collect::<Result<Vec<_>, CompileError>>()?;
 
                     self.compile_statements(body)?;
-                    for _ in 0..items.len() {
-                        self.emit(Instruction::CleanupWith { end: end_label });
+
+                    for end_label in end_labels {
+                        self.emit(Instruction::PopBlock);
+                        self.emit(Instruction::EnterFinally);
+                        self.set_label(end_label);
+                        self.emit(Instruction::WithCleanupStart);
+                        self.emit(Instruction::GetAwaitable);
+                        self.emit(Instruction::LoadConst {
+                            value: bytecode::Constant::None,
+                        });
+                        self.emit(Instruction::YieldFrom);
+                        self.emit(Instruction::WithCleanupFinish);
+                    }
+                } else {
+                    let end_labels = items
+                        .iter()
+                        .map(|item| {
+                            let end_label = self.new_label();
+                            self.compile_expression(&item.context_expr)?;
+                            self.emit(Instruction::SetupWith { end: end_label });
+                            match &item.optional_vars {
+                                Some(var) => {
+                                    self.compile_store(var)?;
+                                }
+                                None => {
+                                    self.emit(Instruction::Pop);
+                                }
+                            }
+                            Ok(end_label)
+                        })
+                        .collect::<Result<Vec<_>, CompileError>>()?;
+
+                    self.compile_statements(body)?;
+
+                    for end_label in end_labels {
+                        self.emit(Instruction::PopBlock);
+                        self.emit(Instruction::EnterFinally);
+                        self.set_label(end_label);
+                        self.emit(Instruction::WithCleanupStart);
+                        self.emit(Instruction::WithCleanupFinish);
                     }
-                    self.set_label(end_label);
                 }
             }
             For {