Generator/Async JIT Fix nested For-In loops

- For in loops use an enumerator object
- for optimisation reasons this is cached in a known location
- in nested for-in loops containing yield cache could fail to restore
- remove a step from the process and bring load nearer to usage

Author: rhuanjl

lib/Backend/Func.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 "Backend.h"
@@ -149,7 +150,7 @@ Func::Func(JitArenaAllocator *alloc, JITTimeWorkItem * workItem,
     , m_forInEnumeratorArrayOffset(-1)
     , argInsCount(0)
     , m_globalObjTypeSpecFldInfoArray(nullptr)
-    , m_forInEnumeratorForGeneratorSym(nullptr)
+    , m_generatorFrameSym(nullptr)
 #if LOWER_SPLIT_INT64
     , m_int64SymPairMap(nullptr)
 #endif

lib/Backend/Func.h

@@ -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.
 //-------------------------------------------------------------------------------------------------------
 #pragma once
@@ -1063,21 +1064,21 @@ static const unsigned __int64 c_debugFillPattern8 = 0xcececececececece;
     StackSym* m_loopParamSym;
     StackSym* m_bailoutReturnValueSym;
     StackSym* m_hasBailedOutSym;
-    StackSym* m_forInEnumeratorForGeneratorSym;
+    StackSym* m_generatorFrameSym;
 
 public:
     StackSym* tempSymDouble;
     StackSym* tempSymBool;
 
-    void SetForInEnumeratorSymForGeneratorSym(StackSym* sym)
+    void SetGeneratorFrameSym(StackSym* sym)
     {
-        Assert(this->m_forInEnumeratorForGeneratorSym == nullptr);
-        this->m_forInEnumeratorForGeneratorSym = sym;
+        Assert(this->m_generatorFrameSym == nullptr);
+        this->m_generatorFrameSym = sym;
     }
 
-    StackSym* GetForInEnumeratorSymForGeneratorSym() const
+    StackSym* GetGeneratorFrameSym() const
     {
-        return this->m_forInEnumeratorForGeneratorSym;
+        return this->m_generatorFrameSym;
     }
 
     // StackSyms' corresponding getters/setters

lib/Backend/IRBuilder.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 "Backend.h"
@@ -1255,7 +1256,7 @@ IRBuilder::EnsureLoopBodyForInEnumeratorArrayOpnd()
 }
 
 IR::Opnd *
-IRBuilder::BuildForInEnumeratorOpnd(uint forInLoopLevel)
+IRBuilder::BuildForInEnumeratorOpnd(uint forInLoopLevel, uint32 offset)
 {
     Assert(forInLoopLevel < this->m_func->GetJITFunctionBody()->GetForInLoopDepth());
     if (this->IsLoopBody())
@@ -1270,7 +1271,7 @@ IRBuilder::BuildForInEnumeratorOpnd(uint forInLoopLevel)
     else if (this->m_func->GetJITFunctionBody()->IsCoroutine())
     {
         return IR::IndirOpnd::New(
-            this->m_generatorJumpTable.EnsureForInEnumeratorArrayOpnd(),
+            this->m_generatorJumpTable.BuildForInEnumeratorArrayOpnd(offset),
             forInLoopLevel * sizeof(Js::ForInObjectEnumerator),
             TyMachPtr,
             this->m_func
@@ -2949,7 +2950,7 @@ IRBuilder::BuildProfiledReg1Unsigned1(Js::OpCode newOpcode, uint32 offset, Js::R
     if (newOpcode == Js::OpCode::InitForInEnumerator)
     {
         IR::RegOpnd * src1Opnd = this->BuildSrcOpnd(R0);
-        IR::Opnd * src2Opnd = this->BuildForInEnumeratorOpnd(C1);
+        IR::Opnd * src2Opnd = this->BuildForInEnumeratorOpnd(C1, offset);
         IR::Instr *instr = IR::ProfiledInstr::New(Js::OpCode::InitForInEnumerator, nullptr, src1Opnd, src2Opnd, m_func);
         instr->AsProfiledInstr()->u.profileId = profileId;
         this->AddInstr(instr, offset);
@@ -3084,7 +3085,7 @@ IRBuilder::BuildReg1Unsigned1(Js::OpCode newOpcode, uint offset, Js::RegSlot R0,
         {
             IR::Instr *instr = IR::Instr::New(Js::OpCode::InitForInEnumerator, m_func);
             instr->SetSrc1(this->BuildSrcOpnd(R0));
-            instr->SetSrc2(this->BuildForInEnumeratorOpnd(C1));
+            instr->SetSrc2(this->BuildForInEnumeratorOpnd(C1, offset));
             this->AddInstr(instr, offset);
             return;
         }
@@ -6935,7 +6936,7 @@ IRBuilder::BuildBrReg1Unsigned1(Js::OpCode newOpcode, uint32 offset)
 void
 IRBuilder::BuildBrBReturn(Js::OpCode newOpcode, uint32 offset, Js::RegSlot DestRegSlot, uint32 forInLoopLevel, uint32 targetOffset)
 {
-    IR::Opnd *srcOpnd = this->BuildForInEnumeratorOpnd(forInLoopLevel);
+    IR::Opnd *srcOpnd = this->BuildForInEnumeratorOpnd(forInLoopLevel, offset);
     IR::RegOpnd *     destOpnd = this->BuildDstOpnd(DestRegSlot);
     IR::BranchInstr * branchInstr = IR::BranchInstr::New(newOpcode, destOpnd, nullptr, srcOpnd, m_func);
     this->AddBranchInstr(branchInstr, offset, targetOffset);
@@ -7922,14 +7923,12 @@ IRBuilder::GeneratorJumpTable::BuildJumpTable()
     //
     // s1 = Ld_A prm1
     // s2 = Ld_A s1[offset of JavascriptGenerator::frame]
-    //      BrNotAddr_A s2 !nullptr $initializationCode
+    //      BrNotAddr_A s2 !nullptr $jumpTable
     //
     // $createInterpreterStackFrame:
     // call helper
     //
-    // $initializationCode:
-    // load for-in enumerator address from interpreter stack frame
-    //
+    // Br $startOfFunc
     // 
     // $jumpTable:
     //
@@ -7963,23 +7962,25 @@ IRBuilder::GeneratorJumpTable::BuildJumpTable()
     IR::LabelInstr* functionBegin = IR::LabelInstr::New(Js::OpCode::Label, this->m_func);
     LABELNAMESET(functionBegin, "GeneratorFunctionBegin");
 
-    IR::LabelInstr* initCode = IR::LabelInstr::New(Js::OpCode::Label, this->m_func);
-    LABELNAMESET(initCode, "GeneratorInitializationAndJumpTable");
+    IR::LabelInstr* jumpTable = IR::LabelInstr::New(Js::OpCode::Label, this->m_func);
+    LABELNAMESET(jumpTable, "GeneratorJumpTable");
 
-    // BrNotAddr_A s2 nullptr $initializationCode
-    IR::BranchInstr* skipCreateInterpreterFrame = IR::BranchInstr::New(Js::OpCode::BrNotAddr_A, initCode, genFrameOpnd, IR::AddrOpnd::NewNull(this->m_func), this->m_func);
+    // If there is already a stack frame, generator function has previously begun execution - don't recreate, skip down to jump table
+    // BrNotAddr_A s2 nullptr $jumpTable
+    IR::BranchInstr* skipCreateInterpreterFrame = IR::BranchInstr::New(Js::OpCode::BrNotAddr_A, jumpTable, genFrameOpnd, IR::AddrOpnd::NewNull(this->m_func), this->m_func);
     this->m_irBuilder->AddInstr(skipCreateInterpreterFrame, this->m_irBuilder->m_functionStartOffset);
 
     // Create interpreter stack frame
     IR::Instr* createInterpreterFrame = IR::Instr::New(Js::OpCode::GeneratorCreateInterpreterStackFrame, genFrameOpnd /* dst */, genRegOpnd /* src */, this->m_func);
     this->m_irBuilder->AddInstr(createInterpreterFrame, this->m_irBuilder->m_functionStartOffset);
 
+    // Having created the frame, skip over the jump table and start executing from the beginning of the function
     IR::BranchInstr* skipJumpTable = IR::BranchInstr::New(Js::OpCode::Br, functionBegin, this->m_func);
     this->m_irBuilder->AddInstr(skipJumpTable, this->m_irBuilder->m_functionStartOffset);
 
-    // Label to insert any initialization code
-    // $initializationCode:
-    this->m_irBuilder->AddInstr(initCode, this->m_irBuilder->m_functionStartOffset);
+    // Label for start of jumpTable - where we look for the correct Yield resume point
+    // $jumpTable:
+    this->m_irBuilder->AddInstr(jumpTable, this->m_irBuilder->m_functionStartOffset);
 
     // s3 = Ld_A s2[offset of InterpreterStackFrame::m_reader.m_currentLocation]
     IR::RegOpnd* curLocOpnd = IR::RegOpnd::New(TyMachPtr, this->m_func);
@@ -8015,46 +8016,24 @@ IRBuilder::GeneratorJumpTable::BuildJumpTable()
 
     this->m_irBuilder->AddInstr(functionBegin, this->m_irBuilder->m_functionStartOffset);
 
-    // Save these values for later use
-    this->m_initLabel = initCode;
+    // Save this value for later use
     this->m_generatorFrameOpnd = genFrameOpnd;
+    this->m_func->SetGeneratorFrameSym(genFrameOpnd->GetStackSym());
 
     return this->m_irBuilder->m_lastInstr;
 }
 
-IR::LabelInstr*
-IRBuilder::GeneratorJumpTable::GetInitLabel() const
-{
-    Assert(this->m_initLabel != nullptr);
-    return this->m_initLabel;
-}
-
 IR::RegOpnd*
-IRBuilder::GeneratorJumpTable::CreateForInEnumeratorArrayOpnd()
-{
-    Assert(this->m_initLabel != nullptr);
+IRBuilder::GeneratorJumpTable::BuildForInEnumeratorArrayOpnd(uint32 offset)
+{   
     Assert(this->m_generatorFrameOpnd != nullptr);
 
-    IR::RegOpnd* forInEnumeratorOpnd = IR::RegOpnd::New(TyMachPtr, this->m_func);
-    IR::Instr* instr = IR::Instr::New(
-        Js::OpCode::Ld_A,
-        forInEnumeratorOpnd,
+    IR::RegOpnd* forInEnumeratorArrayOpnd = IR::RegOpnd::New(TyMachPtr, this->m_func);
+    IR::Instr* instr = IR::Instr::New(Js::OpCode::Ld_A, forInEnumeratorArrayOpnd,
         POINTER_OFFSET(this->m_generatorFrameOpnd, Js::InterpreterStackFrame::GetOffsetOfForInEnumerators(), ForInEnumerators),
         this->m_func
     );
-    this->m_initLabel->InsertAfter(instr);
-
-    return forInEnumeratorOpnd;
-}
-
-IR::RegOpnd*
-IRBuilder::GeneratorJumpTable::EnsureForInEnumeratorArrayOpnd()
-{
-    if (this->m_forInEnumeratorArrayOpnd == nullptr)
-    {
-        this->m_forInEnumeratorArrayOpnd = this->CreateForInEnumeratorArrayOpnd();
-        this->m_func->SetForInEnumeratorSymForGeneratorSym(m_forInEnumeratorArrayOpnd->GetStackSym());
-    }
+    this->m_irBuilder->AddInstr(instr, offset);
 
-    return this->m_forInEnumeratorArrayOpnd;
+    return forInEnumeratorArrayOpnd;
 }

lib/Backend/IRBuilder.h

@@ -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.
 //-------------------------------------------------------------------------------------------------------
 #pragma once
@@ -227,7 +228,7 @@ class IRBuilder
     IR::RegOpnd *       BuildSrcOpnd(Js::RegSlot srcRegSlot, IRType type = TyVar);
     IR::AddrOpnd *      BuildAuxArrayOpnd(AuxArrayValue auxArrayType, uint32 auxArrayOffset);
     IR::Opnd *          BuildAuxObjectLiteralTypeRefOpnd(int objectId);
-    IR::Opnd *          BuildForInEnumeratorOpnd(uint forInLoopLevel);
+    IR::Opnd *          BuildForInEnumeratorOpnd(uint forInLoopLevel, uint32 offset);
     IR::RegOpnd *       EnsureLoopBodyForInEnumeratorArrayOpnd();
 private:
     uint                AddStatementBoundary(uint statementIndex, uint offset);
@@ -366,29 +367,12 @@ class IRBuilder
         Func* const m_func;
         IRBuilder* const m_irBuilder;
 
-        // for-in enumerators are allocated on the heap for jit'd loop body
-        // and on the stack for all other cases (the interpreter frame will
-        // reuses them when bailing out). But because we have the concept of
-        // "bailing in" for generator, reusing enumerators allocated on the stack
-        // would not work. So we have to allocate them on the generator's interpreter
-        // frame instead. This operand is loaded as part of the jump table, before we
-        // jump to any of the resume point.
-        IR::RegOpnd* m_forInEnumeratorArrayOpnd = nullptr;
-
         IR::RegOpnd* m_generatorFrameOpnd = nullptr;
 
-        // This label is used to insert any initialization code that might be needed
-        // when bailing in and before jumping to any of the resume points.
-        // As of now, we only need to load the operand for for-in enumerator.
-        IR::LabelInstr* m_initLabel = nullptr;
-
-        IR::RegOpnd* CreateForInEnumeratorArrayOpnd();
-
     public:
         GeneratorJumpTable(Func* func, IRBuilder* irBuilder);
         IR::Instr* BuildJumpTable();
-        IR::LabelInstr* GetInitLabel() const;
-        IR::RegOpnd* EnsureForInEnumeratorArrayOpnd();
+        IR::RegOpnd* BuildForInEnumeratorArrayOpnd(uint32 offset);
     };
 
     GeneratorJumpTable m_generatorJumpTable;

lib/Backend/LinearScan.cpp

@@ -1,5 +1,6 @@
 //-------------------------------------------------------------------------------------------------------
 // Copyright (C) Microsoft Corporation and contributors. 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.
 //-------------------------------------------------------------------------------------------------------
 
@@ -5301,7 +5302,7 @@ bool LinearScan::GeneratorBailIn::NeedsReloadingBackendSymWhenBailingIn(StackSym
 {
     // for-in enumerator in generator is loaded as part of the resume jump table.
     // By the same reasoning as `initializedRegs`'s, we don't have to restore this whether or not it's been spilled.
-    if (this->func->GetForInEnumeratorSymForGeneratorSym() && this->func->GetForInEnumeratorSymForGeneratorSym()->m_id == sym->m_id)
+    if (this->func->GetGeneratorFrameSym() && this->func->GetGeneratorFrameSym()->m_id == sym->m_id)
     {
         return false;
     }