From 3ecf1c4bcf6e4eedce29bdfcbdb2979b034f2002 Mon Sep 17 00:00:00 2001
From: Haled Odat <8566042+HalidOdat@users.noreply.github.com>
Date: Sat, 23 Sep 2023 22:51:27 +0200
Subject: [PATCH] Make CompletionType::Yield use the return_value

---
 boa_engine/src/vm/mod.rs                       | 13 +++++++++----
 boa_engine/src/vm/opcode/await/mod.rs          | 18 ++++++++++++------
 boa_engine/src/vm/opcode/generator/mod.rs      |  2 +-
 .../src/vm/opcode/generator/yield_stm.rs       | 14 ++++++++------
 boa_engine/src/vm/opcode/mod.rs                |  2 +-
 5 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/boa_engine/src/vm/mod.rs b/boa_engine/src/vm/mod.rs
index aa40c57e09a..5375e27b405 100644
--- a/boa_engine/src/vm/mod.rs
+++ b/boa_engine/src/vm/mod.rs
@@ -206,6 +206,10 @@ impl Vm {
     pub(crate) fn set_return_value(&mut self, value: JsValue) {
         self.return_value = value;
     }
+
+    pub(crate) fn take_return_value(&mut self) -> JsValue {
+        std::mem::take(&mut self.return_value)
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq)]
@@ -355,12 +359,12 @@ impl Context<'_> {
                 Ok(CompletionType::Return) => {
                     self.vm.stack.truncate(self.vm.frame().fp as usize);
 
-                    let execution_result = std::mem::take(&mut self.vm.return_value);
+                    let result = self.vm.take_return_value();
                     if self.vm.frame().exit_early {
-                        return CompletionRecord::Normal(execution_result);
+                        return CompletionRecord::Normal(result);
                     }
 
-                    self.vm.push(execution_result);
+                    self.vm.push(result);
                     self.vm.pop_frame();
                 }
                 Ok(CompletionType::Throw) => {
@@ -401,11 +405,12 @@ impl Context<'_> {
                 }
                 // Early return immediately.
                 Ok(CompletionType::Yield) => {
+                    let result = self.vm.take_return_value();
                     if self.vm.frame().exit_early {
-                        let result = self.vm.pop();
                         return CompletionRecord::Return(result);
                     }
 
+                    self.vm.push(result);
                     self.vm.pop_frame();
                 }
                 Err(err) => {
diff --git a/boa_engine/src/vm/opcode/await/mod.rs b/boa_engine/src/vm/opcode/await/mod.rs
index d865077f256..1123f8f9d30 100644
--- a/boa_engine/src/vm/opcode/await/mod.rs
+++ b/boa_engine/src/vm/opcode/await/mod.rs
@@ -1,7 +1,7 @@
 use boa_gc::{Gc, GcRefCell};
 
 use crate::{
-    builtins::{generator::GeneratorContext, Promise},
+    builtins::{generator::GeneratorContext, promise::PromiseCapability, Promise},
     native_function::NativeFunction,
     object::FunctionObjectBuilder,
     vm::{opcode::Operation, CompletionType, GeneratorResumeKind},
@@ -125,11 +125,17 @@ impl Operation for Await {
             context,
         );
 
-        if let Some(promise_capabality) = context.vm.frame().promise_capability.clone() {
-            context.vm.push(promise_capabality.promise().clone());
-        } else {
-            context.vm.push(JsValue::undefined());
-        }
+        let return_value = context
+            .vm
+            .frame()
+            .promise_capability
+            .as_ref()
+            .map(PromiseCapability::promise)
+            .cloned()
+            .map(JsValue::from)
+            .unwrap_or_default();
+
+        context.vm.set_return_value(return_value);
         Ok(CompletionType::Yield)
     }
 }
diff --git a/boa_engine/src/vm/opcode/generator/mod.rs b/boa_engine/src/vm/opcode/generator/mod.rs
index 6c208a2493f..d74d7103181 100644
--- a/boa_engine/src/vm/opcode/generator/mod.rs
+++ b/boa_engine/src/vm/opcode/generator/mod.rs
@@ -114,7 +114,7 @@ impl Operation for Generator {
                 .async_generator = Some(gen_clone);
         }
 
-        context.vm.push(generator);
+        context.vm.set_return_value(generator.into());
         Ok(CompletionType::Yield)
     }
 }
diff --git a/boa_engine/src/vm/opcode/generator/yield_stm.rs b/boa_engine/src/vm/opcode/generator/yield_stm.rs
index 237fbdca8f4..477a5cf6427 100644
--- a/boa_engine/src/vm/opcode/generator/yield_stm.rs
+++ b/boa_engine/src/vm/opcode/generator/yield_stm.rs
@@ -15,7 +15,9 @@ impl Operation for GeneratorYield {
     const NAME: &'static str = "GeneratorYield";
     const INSTRUCTION: &'static str = "INST - GeneratorYield";
 
-    fn execute(_context: &mut Context<'_>) -> JsResult<CompletionType> {
+    fn execute(context: &mut Context<'_>) -> JsResult<CompletionType> {
+        let value = context.vm.pop();
+        context.vm.set_return_value(value);
         Ok(CompletionType::Yield)
     }
 }
@@ -76,11 +78,11 @@ impl Operation for AsyncGeneratorYield {
 
             context.vm.push(resume_kind);
 
-            Ok(CompletionType::Normal)
-        } else {
-            gen.state = AsyncGeneratorState::SuspendedYield;
-            context.vm.push(JsValue::undefined());
-            GeneratorYield::execute(context)
+            return Ok(CompletionType::Normal);
         }
+
+        gen.state = AsyncGeneratorState::SuspendedYield;
+        context.vm.set_return_value(JsValue::undefined());
+        Ok(CompletionType::Yield)
     }
 }
diff --git a/boa_engine/src/vm/opcode/mod.rs b/boa_engine/src/vm/opcode/mod.rs
index 177e1d20d2a..cd2f7e207eb 100644
--- a/boa_engine/src/vm/opcode/mod.rs
+++ b/boa_engine/src/vm/opcode/mod.rs
@@ -1834,7 +1834,7 @@ generate_opcodes! {
     ///
     /// Operands:
     ///
-    /// Stack: **=>** resume_kind, received
+    /// Stack: value **=>** resume_kind, received
     GeneratorYield,
 
     /// Resumes the current generator function.