LibJS: Align codegen AwaitExpressions to YieldExpressions

We use generators in bytecode to approximate async functions, but the
code generated by AwaitExpressions did not have the value processing
paths that Yield requires, eg the `generator.throw()` path, which is
used by AsyncFunctionDriverWrapper to signal Promise rejections.

4 files changed, 47 insertions, 3 deletions

diff --git a/Userland/Libraries/LibJS/Bytecode/ASTCodegen.cpp b/Userland/Libraries/LibJS/Bytecode/ASTCodegen.cpp
index bdbd344963..33d65f41c4 100644
--- a/Userland/Libraries/LibJS/Bytecode/ASTCodegen.cpp
+++ b/Userland/Libraries/LibJS/Bytecode/ASTCodegen.cpp
@@ -2369,12 +2369,49 @@ Bytecode::CodeGenerationErrorOr<void> AwaitExpression::generate_bytecode(Bytecod
 {
     VERIFY(generator.is_in_async_function());
 
-    // Transform `await expr` to `yield expr`
+    // Transform `await expr` to `yield expr`, see AsyncFunctionDriverWrapper
+    // For that we just need to copy most of the code from YieldExpression
+    auto received_completion_register = generator.allocate_register();
+    auto received_completion_type_register = generator.allocate_register();
+    auto received_completion_value_register = generator.allocate_register();
+
+    auto type_identifier = generator.intern_identifier("type");
+    auto value_identifier = generator.intern_identifier("value");
+
     TRY(m_argument->generate_bytecode(generator));
 
     auto& continuation_block = generator.make_block();
     generator.emit<Bytecode::Op::Yield>(Bytecode::Label { continuation_block });
     generator.switch_to_basic_block(continuation_block);
+
+    // The accumulator is set to an object, for example: { "type": 1 (normal), value: 1337 }
+    generator.emit<Bytecode::Op::Store>(received_completion_register);
+
+    generator.emit<Bytecode::Op::GetById>(type_identifier);
+    generator.emit<Bytecode::Op::Store>(received_completion_type_register);
+
+    generator.emit<Bytecode::Op::Load>(received_completion_register);
+    generator.emit<Bytecode::Op::GetById>(value_identifier);
+    generator.emit<Bytecode::Op::Store>(received_completion_value_register);
+
+    auto& normal_completion_continuation_block = generator.make_block();
+    auto& throw_value_block = generator.make_block();
+
+    generator.emit<Bytecode::Op::LoadImmediate>(Value(to_underlying(Completion::Type::Normal)));
+    generator.emit<Bytecode::Op::StrictlyEquals>(received_completion_type_register);
+    generator.emit<Bytecode::Op::JumpConditional>(
+        Bytecode::Label { normal_completion_continuation_block },
+        Bytecode::Label { throw_value_block });
+
+    // Simplification: The only abrupt completion we receive from AsyncFunctionDriverWrapper is Type::Throw
+    //                 So we do not need to account for the Type::Return path
+    generator.switch_to_basic_block(throw_value_block);
+    generator.emit<Bytecode::Op::Load>(received_completion_value_register);
+    generator.perform_needed_unwinds<Bytecode::Op::Throw>();
+    generator.emit<Bytecode::Op::Throw>();
+
+    generator.switch_to_basic_block(normal_completion_continuation_block);
+    generator.emit<Bytecode::Op::Load>(received_completion_value_register);
     return {};
 }
 
diff --git a/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp b/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp
index 2f91b8bf3a..1926ef3ed0 100644
--- a/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp
+++ b/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp
@@ -169,7 +169,7 @@ Interpreter::ValueAndFrame Interpreter::run_and_return_frame(Executable const& e
     if (!m_return_value.is_empty()) {
         return_value = m_return_value;
         m_return_value = {};
-    } else if (!m_saved_return_value.is_null()) {
+    } else if (!m_saved_return_value.is_null() && m_saved_exception.is_null()) {
         return_value = m_saved_return_value.value();
         m_saved_return_value = {};
     }
@@ -192,6 +192,7 @@ Interpreter::ValueAndFrame Interpreter::run_and_return_frame(Executable const& e
     if (!m_saved_exception.is_null()) {
         Value thrown_value = m_saved_exception.value();
         m_saved_exception = {};
+        m_saved_return_value = {};
         if (auto* register_window = frame.get_pointer<NonnullOwnPtr<RegisterWindow>>())
             return { throw_completion(thrown_value), move(*register_window) };
         return { throw_completion(thrown_value), nullptr };
diff --git a/Userland/Libraries/LibJS/Bytecode/Interpreter.h b/Userland/Libraries/LibJS/Bytecode/Interpreter.h
index 63068c090d..45bd02308c 100644
--- a/Userland/Libraries/LibJS/Bytecode/Interpreter.h
+++ b/Userland/Libraries/LibJS/Bytecode/Interpreter.h
@@ -65,7 +65,11 @@ public:
         VERIFY(unwind_contexts().last().finalizer);
         jump(Label { *unwind_contexts().last().finalizer });
     }
-    void do_return(Value return_value) { m_return_value = return_value; }
+    void do_return(Value return_value)
+    {
+        m_return_value = return_value;
+        m_saved_exception = {};
+    }
 
     void enter_unwind_context(Optional<Label> handler_target, Optional<Label> finalizer_target);
     void leave_unwind_context();
diff --git a/Userland/Libraries/LibJS/Bytecode/Op.cpp b/Userland/Libraries/LibJS/Bytecode/Op.cpp
index ee31553b1c..7bbab0cb6b 100644
--- a/Userland/Libraries/LibJS/Bytecode/Op.cpp
+++ b/Userland/Libraries/LibJS/Bytecode/Op.cpp
@@ -836,6 +836,8 @@ ThrowCompletionOr<void> Yield::execute_impl(Bytecode::Interpreter& interpreter)
     auto object = Object::create(interpreter.realm(), nullptr);
     object->define_direct_property("result", yielded_value, JS::default_attributes);
     if (m_continuation_label.has_value())
+        // FIXME: If we get a pointer, which is not accurately representable as a double
+        //        will cause this to explode
         object->define_direct_property("continuation", Value(static_cast<double>(reinterpret_cast<u64>(&m_continuation_label->block()))), JS::default_attributes);
     else
         object->define_direct_property("continuation", Value(0), JS::default_attributes);