From 1de475b4042f9ab8ff1ae9f54f927cfc2a5fe5c0 Mon Sep 17 00:00:00 2001
From: Hendiadyoin1 <leon.a@serenityos.org>
Date: Sat, 18 May 2024 17:25:43 +0200
Subject: [PATCH] LibJS: Prepare yield object before re-routing it through
 finally

---
 Userland/Libraries/LibJS/Bytecode/Generator.h |  6 ++-
 .../Libraries/LibJS/Bytecode/Instruction.h    |  1 +
 .../Libraries/LibJS/Bytecode/Interpreter.cpp  | 43 +++++++++++++------
 .../Libraries/LibJS/Bytecode/Interpreter.h    |  1 +
 Userland/Libraries/LibJS/Bytecode/Op.h        | 25 +++++++++++
 .../LibJS/Tests/try-return-finally.js         | 40 +++++++++++++++++
 6 files changed, 102 insertions(+), 14 deletions(-)

diff --git a/Userland/Libraries/LibJS/Bytecode/Generator.h b/Userland/Libraries/LibJS/Bytecode/Generator.h
index 8d4fb234c92..e481d66686d 100644
--- a/Userland/Libraries/LibJS/Bytecode/Generator.h
+++ b/Userland/Libraries/LibJS/Bytecode/Generator.h
@@ -283,11 +283,13 @@ public:
             // *  Interpreter::run_bytecode::handle_ContinuePendingUnwind
             // *  Return::execute_impl
             // *  Yield::execute_impl
-            emit<Bytecode::Op::Mov>(Operand(Register::saved_return_value()), value);
+            if constexpr (IsSame<OpType, Op::Yield>)
+                emit<Bytecode::Op::PrepareYield>(Operand(Register::saved_return_value()), value);
+            else
+                emit<Bytecode::Op::Mov>(Operand(Register::saved_return_value()), value);
             emit<Bytecode::Op::Mov>(Operand(Register::exception()), add_constant(Value {}));
             // FIXME: Do we really need to clear the return value register here?
             emit<Bytecode::Op::Mov>(Operand(Register::return_value()), add_constant(Value {}));
-
             emit<Bytecode::Op::Jump>(Label { *m_current_basic_block->finalizer() });
             return;
         }
diff --git a/Userland/Libraries/LibJS/Bytecode/Instruction.h b/Userland/Libraries/LibJS/Bytecode/Instruction.h
index ec026b9c896..f30793758f4 100644
--- a/Userland/Libraries/LibJS/Bytecode/Instruction.h
+++ b/Userland/Libraries/LibJS/Bytecode/Instruction.h
@@ -110,6 +110,7 @@
     O(NewRegExp)                       \
     O(NewTypeError)                    \
     O(Not)                             \
+    O(PrepareYield)                    \
     O(PostfixDecrement)                \
     O(PostfixIncrement)                \
     O(PutById)                         \
diff --git a/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp b/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp
index 76ebf165af6..0b206058e8d 100644
--- a/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp
+++ b/Userland/Libraries/LibJS/Bytecode/Interpreter.cpp
@@ -161,6 +161,22 @@ ALWAYS_INLINE void Interpreter::set(Operand op, Value value)
     m_registers_and_constants_and_locals.data()[op.index()] = value;
 }
 
+ALWAYS_INLINE Value Interpreter::do_yield(Value value, Optional<Label> continuation)
+{
+    auto object = Object::create(realm(), nullptr);
+    object->define_direct_property("result", value, JS::default_attributes);
+
+    if (continuation.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(continuation->address()), JS::default_attributes);
+    else
+        object->define_direct_property("continuation", js_null(), JS::default_attributes);
+
+    object->define_direct_property("isAwait", Value(false), JS::default_attributes);
+    return object;
+}
+
 // 16.1.6 ScriptEvaluation ( scriptRecord ), https://tc39.es/ecma262/#sec-runtime-semantics-scriptevaluation
 ThrowCompletionOr<Value> Interpreter::run(Script& script_record, JS::GCPtr<Environment> lexical_environment_override)
 {
@@ -619,6 +635,7 @@ FLATTEN_ON_CLANG void Interpreter::run_bytecode(size_t entry_point)
             HANDLE_INSTRUCTION_WITHOUT_EXCEPTION_CHECK(NewRegExp);
             HANDLE_INSTRUCTION_WITHOUT_EXCEPTION_CHECK(NewTypeError);
             HANDLE_INSTRUCTION(Not);
+            HANDLE_INSTRUCTION_WITHOUT_EXCEPTION_CHECK(PrepareYield);
             HANDLE_INSTRUCTION(PostfixDecrement);
             HANDLE_INSTRUCTION(PostfixIncrement);
             HANDLE_INSTRUCTION(PutById);
@@ -1762,19 +1779,14 @@ void LeaveUnwindContext::execute_impl(Bytecode::Interpreter& interpreter) const
 void Yield::execute_impl(Bytecode::Interpreter& interpreter) const
 {
     auto yielded_value = interpreter.get(m_value).value_or(js_undefined());
+    interpreter.do_return(
+        interpreter.do_yield(yielded_value, m_continuation_label));
+}
 
-    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(m_continuation_label->address()), JS::default_attributes);
-    else
-        object->define_direct_property("continuation", js_null(), JS::default_attributes);
-
-    object->define_direct_property("isAwait", Value(false), JS::default_attributes);
-    interpreter.do_return(object);
+void PrepareYield::execute_impl(Bytecode::Interpreter& interpreter) const
+{
+    auto value = interpreter.get(m_value).value_or(js_undefined());
+    interpreter.set(m_dest, interpreter.do_yield(value, {}));
 }
 
 void Await::execute_impl(Bytecode::Interpreter& interpreter) const
@@ -2466,6 +2478,13 @@ ByteString Yield::to_byte_string_impl(Bytecode::Executable const& executable) co
         format_operand("value"sv, m_value, executable));
 }
 
+ByteString PrepareYield::to_byte_string_impl(Bytecode::Executable const& executable) const
+{
+    return ByteString::formatted("PrepareYield {}, {}",
+        format_operand("dst"sv, m_dest, executable),
+        format_operand("value"sv, m_value, executable));
+}
+
 ByteString Await::to_byte_string_impl(Bytecode::Executable const& executable) const
 {
     return ByteString::formatted("Await {}, continuation:{}",
diff --git a/Userland/Libraries/LibJS/Bytecode/Interpreter.h b/Userland/Libraries/LibJS/Bytecode/Interpreter.h
index ee1e5c47326..116619f190c 100644
--- a/Userland/Libraries/LibJS/Bytecode/Interpreter.h
+++ b/Userland/Libraries/LibJS/Bytecode/Interpreter.h
@@ -59,6 +59,7 @@ public:
     [[nodiscard]] Value get(Operand) const;
     void set(Operand, Value);
 
+    Value do_yield(Value value, Optional<Label> continuation);
     void do_return(Value value)
     {
         reg(Register::return_value()) = value;
diff --git a/Userland/Libraries/LibJS/Bytecode/Op.h b/Userland/Libraries/LibJS/Bytecode/Op.h
index 4080c9d8eb3..d2a080ee516 100644
--- a/Userland/Libraries/LibJS/Bytecode/Op.h
+++ b/Userland/Libraries/LibJS/Bytecode/Op.h
@@ -2214,6 +2214,31 @@ private:
     Operand m_value;
 };
 
+class PrepareYield final : public Instruction {
+public:
+    explicit PrepareYield(Operand dest, Operand value)
+        : Instruction(Type::PrepareYield)
+        , m_dest(dest)
+        , m_value(value)
+    {
+    }
+
+    void execute_impl(Bytecode::Interpreter&) const;
+    ByteString to_byte_string_impl(Bytecode::Executable const&) const;
+    void visit_operands_impl(Function<void(Operand&)> visitor)
+    {
+        visitor(m_dest);
+        visitor(m_value);
+    }
+
+    Operand destination() const { return m_dest; }
+    Operand value() const { return m_value; }
+
+private:
+    Operand m_dest;
+    Operand m_value;
+};
+
 class Await final : public Instruction {
 public:
     constexpr static bool IsTerminator = true;
diff --git a/Userland/Libraries/LibJS/Tests/try-return-finally.js b/Userland/Libraries/LibJS/Tests/try-return-finally.js
index 7eec43d0eb7..79b78cc5fce 100644
--- a/Userland/Libraries/LibJS/Tests/try-return-finally.js
+++ b/Userland/Libraries/LibJS/Tests/try-return-finally.js
@@ -38,3 +38,43 @@ test("restore exception after generator yield in finally", () => {
     expect(() => generator.next()).toThrowWithMessage(Error, "foo");
     expect(generator.next().done).toBe(true);
 });
+
+test("yield, then return from finally", () => {
+    let test = [];
+    let generator = (function* () {
+        try {
+            yield 1;
+            test.push(1);
+        } finally {
+            test.push(2);
+            return 2;
+        }
+        expect.fail("unreachable");
+    })();
+
+    let result = generator.next();
+    expect(result.value).toBe(1);
+    expect(result.done).toBe(false);
+    result = generator.next();
+    expect(result.value).toBe(2);
+    expect(result.done).toBe(true);
+    expect(test).toEqual([1, 2]);
+});
+
+test("return from async through finally", () => {
+    let test = 0;
+    let result = (async function () {
+        try {
+            return { y: 5 };
+        } finally {
+            test = 42;
+        }
+        expect.fail("unreachable");
+    })();
+
+    expect(result).toBeInstanceOf(Promise);
+    expect(test).toBe(42);
+    result.then(value => {
+        expect(value).toEqual({ y: 5 });
+    });
+});