LibJS: Implement async functions as generator functions in BC mode

This applies a simple transformation, and adds a simple wrapper that
translates the generator interface to the async function interface.

Userland/Libraries/LibJS/AST.h

@@ -18,6 +18,7 @@
 #include <AK/Vector.h>
 #include <LibJS/Forward.h>
 #include <LibJS/Runtime/EnvironmentCoordinate.h>
+#include <LibJS/Runtime/FunctionKind.h>
 #include <LibJS/Runtime/PropertyKey.h>
 #include <LibJS/Runtime/Reference.h>
 #include <LibJS/Runtime/Value.h>
@@ -33,12 +34,6 @@ class Identifier;
 class MemberExpression;
 class VariableDeclaration;
 
-enum class FunctionKind {
-    Generator,
-    Regular,
-    Async,
-};
-
 template<class T, class... Args>
 static inline NonnullRefPtr<T>
 create_ast_node(SourceRange range, Args&&... args)
@@ -575,6 +570,7 @@ public:
 
     virtual Value execute(Interpreter&, GlobalObject&) const override;
     virtual void dump(int indent) const override;
+    virtual void generate_bytecode(Bytecode::Generator&) const override;
 
 private:
     NonnullRefPtr<Expression> m_argument;

Userland/Libraries/LibJS/Bytecode/ASTCodegen.cpp

@@ -911,7 +911,7 @@ void ReturnStatement::generate_bytecode(Bytecode::Generator& generator) const
     if (m_argument)
         m_argument->generate_bytecode(generator);
 
-    if (generator.is_in_generator_function())
+    if (generator.is_in_generator_or_async_function())
         generator.emit<Bytecode::Op::Yield>(nullptr);
     else
         generator.emit<Bytecode::Op::Return>();
@@ -1269,4 +1269,16 @@ void ThisExpression::generate_bytecode(Bytecode::Generator& generator) const
     generator.emit<Bytecode::Op::ResolveThisBinding>();
 }
 
+void AwaitExpression::generate_bytecode(Bytecode::Generator& generator) const
+{
+    VERIFY(generator.is_in_async_function());
+
+    // Transform `await expr` to `yield expr`
+    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);
+}
+
 }

Userland/Libraries/LibJS/Bytecode/Generator.cpp

@@ -23,19 +23,19 @@ Generator::~Generator()
 {
 }
 
-Executable Generator::generate(ASTNode const& node, bool is_in_generator_function)
+Executable Generator::generate(ASTNode const& node, FunctionKind enclosing_function_kind)
 {
     Generator generator;
     generator.switch_to_basic_block(generator.make_block());
-    if (is_in_generator_function) {
-        generator.enter_generator_context();
+    generator.m_enclosing_function_kind = enclosing_function_kind;
+    if (generator.is_in_generator_or_async_function()) {
         // Immediately yield with no value.
         auto& start_block = generator.make_block();
         generator.emit<Bytecode::Op::Yield>(Label { start_block });
         generator.switch_to_basic_block(start_block);
     }
     node.generate_bytecode(generator);
-    if (is_in_generator_function) {
+    if (generator.is_in_generator_or_async_function()) {
         // Terminate all unterminated blocks with yield return
         for (auto& block : generator.m_root_basic_blocks) {
             if (block.is_terminated())

Userland/Libraries/LibJS/Bytecode/Generator.h

@@ -17,12 +17,13 @@
 #include <LibJS/Bytecode/Register.h>
 #include <LibJS/Bytecode/StringTable.h>
 #include <LibJS/Forward.h>
+#include <LibJS/Runtime/FunctionKind.h>
 
 namespace JS::Bytecode {
 
 class Generator {
 public:
-    static Executable generate(ASTNode const&, bool is_in_generator_function = false);
+    static Executable generate(ASTNode const&, FunctionKind = FunctionKind::Regular);
 
     Register allocate_register();
 
@@ -111,9 +112,9 @@ public:
         return m_identifier_table->insert(move(string));
     }
 
-    bool is_in_generator_function() const { return m_is_in_generator_function; }
-    void enter_generator_context() { m_is_in_generator_function = true; }
-    void leave_generator_context() { m_is_in_generator_function = false; }
+    bool is_in_generator_or_async_function() const { return m_enclosing_function_kind == FunctionKind::Async || m_enclosing_function_kind == FunctionKind::Generator; }
+    bool is_in_generator_function() const { return m_enclosing_function_kind == FunctionKind::Generator; }
+    bool is_in_async_function() const { return m_enclosing_function_kind == FunctionKind::Async; }
 
 private:
     Generator();
@@ -129,7 +130,7 @@ private:
 
     u32 m_next_register { 2 };
     u32 m_next_block { 1 };
-    bool m_is_in_generator_function { false };
+    FunctionKind m_enclosing_function_kind { FunctionKind::Regular };
     Vector<Label> m_continuable_scopes;
     Vector<Label> m_breakable_scopes;
 };

Userland/Libraries/LibJS/CMakeLists.txt

@@ -39,6 +39,7 @@ set(SOURCES
     Runtime/ArrayIteratorPrototype.cpp
     Runtime/ArrayPrototype.cpp
     Runtime/AsyncFunctionConstructor.cpp
+    Runtime/AsyncFunctionDriverWrapper.cpp
     Runtime/AsyncFunctionPrototype.cpp
     Runtime/AtomicsObject.cpp
     Runtime/BigInt.cpp

Userland/Libraries/LibJS/Runtime/AsyncFunctionDriverWrapper.cpp

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <LibJS/Runtime/AsyncFunctionDriverWrapper.h>
+#include <LibJS/Runtime/GlobalObject.h>
+#include <LibJS/Runtime/NativeFunction.h>
+#include <LibJS/Runtime/PromiseReaction.h>
+#include <LibJS/Runtime/VM.h>
+
+namespace JS {
+
+ThrowCompletionOr<Value> AsyncFunctionDriverWrapper::create(GlobalObject& global_object, GeneratorObject* generator_object)
+{
+    auto wrapper = global_object.heap().allocate<AsyncFunctionDriverWrapper>(global_object, global_object, generator_object);
+    return wrapper->react_to_async_task_completion(global_object.vm(), global_object, js_undefined(), true);
+}
+
+AsyncFunctionDriverWrapper::AsyncFunctionDriverWrapper(GlobalObject& global_object, GeneratorObject* generator_object)
+    : Promise(global_object)
+    , m_generator_object(generator_object)
+    , m_on_fulfillment(NativeFunction::create(global_object, "async.on_fulfillment"sv, [this](VM& vm, GlobalObject& global_object) {
+        return react_to_async_task_completion(vm, global_object, vm.argument(0), true);
+    }))
+    , m_on_rejection(NativeFunction::create(global_object, "async.on_rejection"sv, [this](VM& vm, GlobalObject& global_object) {
+        return react_to_async_task_completion(vm, global_object, vm.argument(0), false);
+    }))
+{
+}
+
+ThrowCompletionOr<Value> AsyncFunctionDriverWrapper::react_to_async_task_completion(VM& vm, GlobalObject& global_object, Value value, bool is_successful)
+{
+    auto generator_result = is_successful
+        ? m_generator_object->next_impl(vm, global_object, value, {})
+        : m_generator_object->next_impl(vm, global_object, {}, value);
+
+    if (generator_result.is_throw_completion()) {
+        VERIFY(generator_result.throw_completion().type() == Completion::Type::Throw);
+        vm.clear_exception();
+        vm.stop_unwind();
+        auto promise = Promise::create(global_object);
+        promise->reject(generator_result.throw_completion().value());
+        return promise;
+    }
+
+    auto result = generator_result.release_value();
+    VERIFY(result.is_object());
+
+    auto promise_value = TRY(result.get(global_object, vm.names.value));
+    if (!promise_value.is_object() || !is<Promise>(promise_value.as_object())) {
+        auto promise = Promise::create(global_object);
+        promise->fulfill(promise_value);
+        return promise;
+    }
+
+    auto* promise = static_cast<Promise*>(&promise_value.as_object());
+    if (TRY(result.get(global_object, vm.names.done)).to_boolean())
+        return promise;
+
+    return promise->perform_then(m_on_fulfillment, m_on_rejection, PromiseCapability { promise, m_on_fulfillment, m_on_rejection });
+}
+
+AsyncFunctionDriverWrapper::~AsyncFunctionDriverWrapper()
+{
+}
+
+void AsyncFunctionDriverWrapper::visit_edges(Cell::Visitor& visitor)
+{
+    Base::visit_edges(visitor);
+    visitor.visit(m_generator_object);
+    visitor.visit(m_on_fulfillment);
+    visitor.visit(m_on_rejection);
+}
+
+}

Userland/Libraries/LibJS/Runtime/AsyncFunctionDriverWrapper.h

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <LibJS/Bytecode/Interpreter.h>
+#include <LibJS/Runtime/ECMAScriptFunctionObject.h>
+#include <LibJS/Runtime/GeneratorObject.h>
+#include <LibJS/Runtime/Object.h>
+#include <LibJS/Runtime/Promise.h>
+
+namespace JS {
+
+class AsyncFunctionDriverWrapper final : public Promise {
+    JS_OBJECT(AsyncFunctionDriverWrapper, Promise);
+
+public:
+    static ThrowCompletionOr<Value> create(GlobalObject&, GeneratorObject*);
+    explicit AsyncFunctionDriverWrapper(GlobalObject&, GeneratorObject*);
+
+    virtual ~AsyncFunctionDriverWrapper() override;
+    void visit_edges(Cell::Visitor&) override;
+
+    ThrowCompletionOr<Value> react_to_async_task_completion(VM&, GlobalObject&, Value, bool is_successful);
+
+private:
+    GeneratorObject* m_generator_object { nullptr };
+    NativeFunction* m_on_fulfillment { nullptr };
+    NativeFunction* m_on_rejection { nullptr };
+};
+
+}

Userland/Libraries/LibJS/Runtime/ECMAScriptFunctionObject.cpp

@@ -13,6 +13,7 @@
 #include <LibJS/Interpreter.h>
 #include <LibJS/Runtime/AbstractOperations.h>
 #include <LibJS/Runtime/Array.h>
+#include <LibJS/Runtime/AsyncFunctionDriverWrapper.h>
 #include <LibJS/Runtime/ECMAScriptFunctionObject.h>
 #include <LibJS/Runtime/Error.h>
 #include <LibJS/Runtime/ExecutionContext.h>
@@ -100,11 +101,10 @@ void ECMAScriptFunctionObject::initialize(GlobalObject& global_object)
             MUST(prototype->define_property_or_throw(vm.names.constructor, { .value = this, .writable = true, .enumerable = false, .configurable = true }));
             break;
         case FunctionKind::Generator:
+        case FunctionKind::Async:
             // prototype is "g1.prototype" in figure-2 (https://tc39.es/ecma262/img/figure-2.png)
             prototype = global_object.generator_object_prototype();
             break;
-        case FunctionKind::Async:
-            break;
         }
         define_direct_property(vm.names.prototype, prototype, Attribute::Writable);
     }
@@ -753,12 +753,10 @@ Completion ECMAScriptFunctionObject::ordinary_call_evaluate_body()
     auto* bytecode_interpreter = Bytecode::Interpreter::current();
 
     if (bytecode_interpreter) {
-        if (m_kind == FunctionKind::Async)
-            return vm.throw_completion<InternalError>(global_object(), ErrorType::NotImplemented, "Async function execution in Bytecode interpreter");
         // FIXME: pass something to evaluate default arguments with
         TRY(function_declaration_instantiation(nullptr));
         if (!m_bytecode_executable.has_value()) {
-            m_bytecode_executable = Bytecode::Generator::generate(m_ecmascript_code, m_kind == FunctionKind::Generator);
+            m_bytecode_executable = Bytecode::Generator::generate(m_ecmascript_code, m_kind);
             m_bytecode_executable->name = m_name;
             auto& passes = JS::Bytecode::Interpreter::optimization_pipeline();
             passes.perform(*m_bytecode_executable);
@@ -778,10 +776,17 @@ Completion ECMAScriptFunctionObject::ordinary_call_evaluate_body()
 
         // NOTE: Running the bytecode should eventually return a completion.
         // Until it does, we assume "return" and include the undefined fallback from the call site.
-        if (m_kind != FunctionKind::Generator)
+        if (m_kind == FunctionKind::Regular)
             return { Completion::Type::Return, result.value_or(js_undefined()), {} };
 
-        auto generator_object = TRY(GeneratorObject::create(global_object(), result, this, vm.running_execution_context().lexical_environment, move(*result_and_frame.frame)));
+        auto generator_object = TRY(GeneratorObject::create(global_object(), result, this, vm.running_execution_context().copy(), move(*result_and_frame.frame)));
+
+        // NOTE: Async functions are entirely transformed to generator functions, and wrapped in a custom driver that returns a promise
+        //       See AwaitExpression::generate_bytecode() for the transformation.
+        if (m_kind == FunctionKind::Async)
+            return { Completion::Type::Return, TRY(AsyncFunctionDriverWrapper::create(global_object(), generator_object)), {} };
+
+        VERIFY(m_kind == FunctionKind::Generator);
         return { Completion::Type::Return, generator_object, {} };
     } else {
         if (m_kind == FunctionKind::Generator)

Userland/Libraries/LibJS/Runtime/FunctionKind.h

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+namespace JS {
+
+enum class FunctionKind {
+    Generator,
+    Regular,
+    Async,
+};
+
+}

Userland/Libraries/LibJS/Runtime/GeneratorFunctionConstructor.cpp

@@ -50,7 +50,7 @@ ThrowCompletionOr<Object*> GeneratorFunctionConstructor::construct(FunctionObjec
     auto* bytecode_interpreter = Bytecode::Interpreter::current();
     VERIFY(bytecode_interpreter);
 
-    auto executable = Bytecode::Generator::generate(function->body(), true);
+    auto executable = Bytecode::Generator::generate(function->body(), FunctionKind::Generator);
     auto& passes = JS::Bytecode::Interpreter::optimization_pipeline();
     passes.perform(executable);
     if constexpr (JS_BYTECODE_DEBUG) {

Userland/Libraries/LibJS/Runtime/GeneratorObject.cpp

@@ -13,21 +13,21 @@
 
 namespace JS {
 
-ThrowCompletionOr<GeneratorObject*> GeneratorObject::create(GlobalObject& global_object, Value initial_value, ECMAScriptFunctionObject* generating_function, Environment* generating_scope, Bytecode::RegisterWindow frame)
+ThrowCompletionOr<GeneratorObject*> GeneratorObject::create(GlobalObject& global_object, Value initial_value, ECMAScriptFunctionObject* generating_function, ExecutionContext execution_context, Bytecode::RegisterWindow frame)
 {
     // This is "g1.prototype" in figure-2 (https://tc39.es/ecma262/img/figure-2.png)
     auto generating_function_prototype = TRY(generating_function->get(global_object.vm().names.prototype));
     auto* generating_function_prototype_object = TRY(generating_function_prototype.to_object(global_object));
-    auto object = global_object.heap().allocate<GeneratorObject>(global_object, global_object, *generating_function_prototype_object);
+    auto object = global_object.heap().allocate<GeneratorObject>(global_object, global_object, *generating_function_prototype_object, move(execution_context));
     object->m_generating_function = generating_function;
-    object->m_environment = generating_scope;
     object->m_frame = move(frame);
     object->m_previous_value = initial_value;
     return object;
 }
 
-GeneratorObject::GeneratorObject(GlobalObject&, Object& prototype)
+GeneratorObject::GeneratorObject(GlobalObject&, Object& prototype, ExecutionContext context)
     : Object(prototype)
+    , m_execution_context(move(context))
 {
 }
 
@@ -42,12 +42,11 @@ GeneratorObject::~GeneratorObject()
 void GeneratorObject::visit_edges(Cell::Visitor& visitor)
 {
     Base::visit_edges(visitor);
-    visitor.visit(m_environment);
     visitor.visit(m_generating_function);
     visitor.visit(m_previous_value);
 }
 
-ThrowCompletionOr<Value> GeneratorObject::next_impl(VM& vm, GlobalObject& global_object, Optional<Value> value_to_throw)
+ThrowCompletionOr<Value> GeneratorObject::next_impl(VM& vm, GlobalObject& global_object, Optional<Value> next_argument, Optional<Value> value_to_throw)
 {
     auto bytecode_interpreter = Bytecode::Interpreter::current();
     VERIFY(bytecode_interpreter);
@@ -92,20 +91,25 @@ ThrowCompletionOr<Value> GeneratorObject::next_impl(VM& vm, GlobalObject& global
     // Restore the snapshot registers
     bytecode_interpreter->enter_frame(m_frame);
 
+    // Temporarily switch to the captured execution context
+    vm.push_execution_context(m_execution_context, global_object);
+
     // Pretend that 'yield' returned the passed value, or threw
     if (value_to_throw.has_value()) {
         vm.throw_exception(global_object, value_to_throw.release_value());
         bytecode_interpreter->accumulator() = js_undefined();
     } else {
-        bytecode_interpreter->accumulator() = vm.argument(0);
+        bytecode_interpreter->accumulator() = next_argument.value_or(js_undefined());
     }
 
-    // Temporarily switch to the captured environment record
-    TemporaryChange change { vm.running_execution_context().lexical_environment, m_environment };
+    // Temporarily switch to the captured execution context
+    vm.push_execution_context(m_execution_context, global_object);
 
     m_previous_value = bytecode_interpreter->run(*m_generating_function->bytecode_executable(), next_block);
 
-    bytecode_interpreter->leave_frame();
+    m_frame = move(*bytecode_interpreter->pop_frame());
+
+    vm.pop_execution_context();
 
     m_done = TRY(generated_continuation(m_previous_value)) == nullptr;
 

Userland/Libraries/LibJS/Runtime/GeneratorObject.h

@@ -16,17 +16,17 @@ class GeneratorObject final : public Object {
     JS_OBJECT(GeneratorObject, Object);
 
 public:
-    static ThrowCompletionOr<GeneratorObject*> create(GlobalObject&, Value, ECMAScriptFunctionObject*, Environment*, Bytecode::RegisterWindow);
-    GeneratorObject(GlobalObject&, Object& prototype);
+    static ThrowCompletionOr<GeneratorObject*> create(GlobalObject&, Value, ECMAScriptFunctionObject*, ExecutionContext, Bytecode::RegisterWindow);
+    GeneratorObject(GlobalObject&, Object& prototype, ExecutionContext);
     virtual void initialize(GlobalObject&) override;
     virtual ~GeneratorObject() override;
     void visit_edges(Cell::Visitor&) override;
 
-    ThrowCompletionOr<Value> next_impl(VM&, GlobalObject&, Optional<Value> value_to_throw);
+    ThrowCompletionOr<Value> next_impl(VM&, GlobalObject&, Optional<Value> next_argument, Optional<Value> value_to_throw);
     void set_done() { m_done = true; }
 
 private:
-    Environment* m_environment { nullptr };
+    ExecutionContext m_execution_context;
     ECMAScriptFunctionObject* m_generating_function { nullptr };
     Value m_previous_value;
     Bytecode::RegisterWindow m_frame;

Userland/Libraries/LibJS/Runtime/GeneratorObjectPrototype.cpp

@@ -35,7 +35,7 @@ GeneratorObjectPrototype::~GeneratorObjectPrototype()
 JS_DEFINE_NATIVE_FUNCTION(GeneratorObjectPrototype::next)
 {
     auto* generator_object = TRY(typed_this_object(global_object));
-    return generator_object->next_impl(vm, global_object, {});
+    return generator_object->next_impl(vm, global_object, vm.argument(0), {});
 }
 
 // 27.5.1.3 Generator.prototype.next ( value ), https://tc39.es/ecma262/#sec-generator.prototype.return
@@ -43,14 +43,14 @@ JS_DEFINE_NATIVE_FUNCTION(GeneratorObjectPrototype::return_)
 {
     auto* generator_object = TRY(typed_this_object(global_object));
     generator_object->set_done();
-    return generator_object->next_impl(vm, global_object, {});
+    return generator_object->next_impl(vm, global_object, {}, {});
 }
 
 // 27.5.1.4 Generator.prototype.next ( value ), https://tc39.es/ecma262/#sec-generator.prototype.throw
 JS_DEFINE_NATIVE_FUNCTION(GeneratorObjectPrototype::throw_)
 {
     auto* generator_object = TRY(typed_this_object(global_object));
-    return generator_object->next_impl(vm, global_object, vm.argument(0));
+    return generator_object->next_impl(vm, global_object, {}, vm.argument(0));
 }
 
 }

Userland/Libraries/LibJS/Runtime/Promise.h

@@ -13,7 +13,7 @@ namespace JS {
 
 ThrowCompletionOr<Object*> promise_resolve(GlobalObject&, Object& constructor, Value);
 
-class Promise final : public Object {
+class Promise : public Object {
     JS_OBJECT(Promise, Object);
 
 public:
@@ -45,9 +45,10 @@ public:
     Value reject(Value reason);
     Value perform_then(Value on_fulfilled, Value on_rejected, Optional<PromiseCapability> result_capability);
 
-private:
+protected:
     virtual void visit_edges(Visitor&) override;
 
+private:
     bool is_settled() const { return m_state == State::Fulfilled || m_state == State::Rejected; }
 
     void trigger_reactions() const;