Handle SlotArray loads in jitted generators

Author: rhuanjl

lib/Runtime/ByteCode/ByteCodeEmitter.cpp

@@ -1792,7 +1792,7 @@ void ByteCodeGenerator::FinalizeRegisters(FuncInfo* funcInfo, Js::FunctionBody*
         }
     }
 
-    // NOTE: The FB expects the yield reg to be the final non-temp.
+    // NOTE: The FunctionBody expects the yield reg to be the final non-temp.
     if (byteCodeFunction->IsCoroutine())
     {
         if (funcInfo->root->IsAsync())

lib/Runtime/Library/JavascriptGenerator.cpp

@@ -1,5 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft. All rights reserved.
+// Copyright (c) 2021 ChakraCore Project Contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 #include "RuntimeLibraryPch.h"
@@ -171,6 +172,18 @@ Var JavascriptGenerator::CallGenerator(Var data, ResumeYieldKind resumeKind)
     JavascriptLibrary* library = scriptContext->GetLibrary();
     Var result = nullptr;
 
+    if (this->frame)
+    {
+        // if the function already has a state it may be going to resume in the jit
+        // if so copy any innerScopes into registers jit can access
+        uint32 innerScopeCount = this->scriptFunction->GetFunctionBody()->GetInnerScopeCount();
+        for (uint32 i = 0; i < innerScopeCount; ++i)
+        {
+            Js::RegSlot reg = this->scriptFunction->GetFunctionBody()->GetFirstInnerScopeRegister() + i;
+            this->frame->SetNonVarReg(reg, this->frame->InnerScopeFromIndex(i));
+        }
+    }
+
     SetResumeYieldProperties(data, resumeKind);
 
     {

test/es6/generator-jit-bugs.js

@@ -77,7 +77,7 @@ check(gen3.next().value, 1);
 check(gen3.next().value, 2);
 
 // Test 4 - yield* iterator fails to be restored after Bail on No Profile
-title("Bail on no profile losing yield* iterator")
+title("Bail on no profile losing yield* iterator");
 function* gf4() {
     yield 0;
     yield* [1,2,3];
@@ -90,4 +90,51 @@ check(gen4.next().value, 1);
 check(gen4.next().value, 2);
 check(gen4.next().value, 3);
 
+title("Load Scope Slots in presence of for-in");
+function* gf5(v1) {
+    for(v0 in v1) {
+        yield undefined;
+        let v2 = {}
+        function v3() { v2;}
+    }
+}
+
+const gen5 = gf5([0, 1]);
+
+check(gen5.next().value, undefined);
+check(gen5.next().value, undefined);
+check(gen5.next().value, undefined);
+check(gen5.next().value, undefined);
+
+title("Load Scope Slots used in loop control");
+function* gf6 () {
+    for (let v1 = 0; v1 < 1000; ++v1) {
+        function foo() {v1;}
+        yield v1;
+    }
+}
+
+const gen6 = gf6();
+
+check(gen6.next().value, 0);
+check(gen6.next().value, 1);
+check(gen6.next().value, 2);
+check(gen6.next().value, 3);
+
+title("Load Scope Slots used in loop control and captured indirectly");
+function* gf7(v1) {
+    for (const v2 in v1) {
+        yield v2;
+        const v4 = [v2];
+        function foo() { v4; }
+    }
+}
+
+const gen7 = gf7([0, 1, 2]);
+check(gen7.next().value, 0);
+check(gen7.next().value, 1);
+check(gen7.next().value, 2);
+check(gen7.next().value, undefined);
+
+
 print("pass");