Rework how captured values are written to GeneratorBailInInstrs

lib/Backend/BackwardPass.cpp

@@ -1770,9 +1770,17 @@ BackwardPass::ProcessBailOutConstants(BailOutInfo * bailOutInfo, BVSparse<JitAre
     }
     NEXT_SLISTBASE_ENTRY;
 
+    IR::GeneratorBailInInstr* bailInInstr = bailOutInfo->bailInInstr;
+
     // Find other constants that we need to restore
     FOREACH_SLISTBASE_ENTRY_EDITING(ConstantStackSymValue, value, &bailOutInfo->capturedValues->constantValues, iter)
     {
+        if (bailInInstr)
+        {
+            // Store all captured constant values for the corresponding bailin instr
+            bailInInstr->capturedValues.constantValues.PrependNode(this->func->m_alloc, value);
+        }
+
         if (byteCodeUpwardExposedUsed->TestAndClear(value.Key()->m_id) || bailoutReferencedArgSymsBv->TestAndClear(value.Key()->m_id))
         {
             // Constant need to be restore, move it to the restore list
@@ -1805,6 +1813,7 @@ BackwardPass::ProcessBailOutCopyProps(BailOutInfo * bailOutInfo, BVSparse<JitAre
     JitArenaAllocator * allocator = this->func->m_alloc;
     BasicBlock * block = this->currentBlock;
     BVSparse<JitArenaAllocator> * upwardExposedUses = block->upwardExposedUses;
+    IR::GeneratorBailInInstr* bailInInstr = bailOutInfo->bailInInstr;
 
     // Find other copy prop that we need to restore
     FOREACH_SLISTBASE_ENTRY_EDITING(CopyPropSyms, copyPropSyms, &bailOutInfo->capturedValues->copyPropSyms, iter)
@@ -1814,6 +1823,13 @@ BackwardPass::ProcessBailOutCopyProps(BailOutInfo * bailOutInfo, BVSparse<JitAre
         Assert(!copyPropSyms.Value()->IsTypeSpec());
         if (byteCodeUpwardExposedUsed->TestAndClear(copyPropSyms.Key()->m_id) || bailoutReferencedArgSymsBv->TestAndClear(copyPropSyms.Key()->m_id))
         {
+            if (bailInInstr)
+            {
+                // Copy all copyprop syms into the corresponding bail-in instr so that
+                // we can map the correct symbols to restore during bail-in
+                bailInInstr->capturedValues.copyPropSyms.PrependNode(allocator, copyPropSyms);
+            }
+
             // This copy-prop sym needs to be restored; add it to the restore list.
 
             /*
@@ -2561,30 +2577,6 @@ BackwardPass::NeedBailOutOnImplicitCallsForTypedArrayStore(IR::Instr* instr)
     return false;
 }
 
-IR::Instr*
-BackwardPass::ProcessPendingPreOpBailOutInfoForYield(IR::Instr* const currentInstr)
-{
-    Assert(currentInstr->m_opcode == Js::OpCode::Yield);
-    IR::GeneratorBailInInstr* bailInInstr = currentInstr->m_next->m_next->AsGeneratorBailInInstr();
-
-    BailOutInfo* bailOutInfo = currentInstr->GetBailOutInfo();
-
-    // Make a copy of all detected constant values before we actually process
-    // the bailout info since we will then remove any values that don't need
-    // to be restored for the normal bailout cases. As for yields, we still
-    // need them for our bailin code.
-    bailInInstr->SetConstantValues(bailOutInfo->capturedValues->constantValues);
-
-    IR::Instr* ret = this->ProcessPendingPreOpBailOutInfo(currentInstr);
-
-    // We will need list of symbols that have been copy-prop'd to map the correct
-    // symbols to restore during bail-in. Since this list is cleared during
-    // FillBailOutRecord, make a copy of it now.
-    bailInInstr->SetCopyPropSyms(bailOutInfo->usedCapturedValues->copyPropSyms);
-
-    return ret;
-}
-
 IR::Instr*
 BackwardPass::ProcessPendingPreOpBailOutInfo(IR::Instr *const currentInstr)
 {
@@ -3887,17 +3879,10 @@ BackwardPass::ProcessBlock(BasicBlock * block)
         // bail-in
         if (instr->IsGeneratorBailInInstr() && this->currentBlock->upwardExposedUses)
         {
-            instr->AsGeneratorBailInInstr()->SetUpwardExposedUses(*this->currentBlock->upwardExposedUses);
+            instr->AsGeneratorBailInInstr()->upwardExposedUses.Copy(this->currentBlock->upwardExposedUses);
         }
 
-        if (instr->m_opcode == Js::OpCode::Yield)
-        {
-            instrPrev = ProcessPendingPreOpBailOutInfoForYield(instr);
-        }
-        else
-        {
-            instrPrev = ProcessPendingPreOpBailOutInfo(instr);
-        }
+        instrPrev = ProcessPendingPreOpBailOutInfo(instr);
 
 #if DBG_DUMP
         TraceInstrUses(block, instr, false);

lib/Backend/BackwardPass.h

@@ -45,7 +45,6 @@ class BackwardPass
     bool ProcessByteCodeUsesInstr(IR::Instr * instr);
     bool ProcessBailOutInfo(IR::Instr * instr);
     void ProcessBailOutInfo(IR::Instr * instr, BailOutInfo * bailOutInfo);
-    IR::Instr* ProcessPendingPreOpBailOutInfoForYield(IR::Instr* const currentInstr);
     IR::Instr* ProcessPendingPreOpBailOutInfo(IR::Instr *const currentInstr);
     void ClearDstUseForPostOpLazyBailOut(IR::Instr *instr);
     void ProcessBailOutArgObj(BailOutInfo * bailOutInfo, BVSparse<JitArenaAllocator> * byteCodeUpwardExposedUsed);

lib/Backend/BailOut.cpp

@@ -75,6 +75,7 @@ void BailOutInfo::PartialDeepCopyTo(BailOutInfo * const other) const
 #endif
 
     other->bailOutInstr = this->bailOutInstr;
+    other->bailInInstr = this->bailInInstr;
 
 #if ENABLE_DEBUG_CONFIG_OPTIONS
     other->bailOutOpcode = this->bailOutOpcode;

lib/Backend/BailOut.h

@@ -26,7 +26,7 @@ class BailOutInfo
 
     BailOutInfo(uint32 bailOutOffset, Func* bailOutFunc) :
         bailOutOffset(bailOutOffset), bailOutFunc(bailOutFunc),
-        byteCodeUpwardExposedUsed(nullptr), polymorphicCacheIndex((uint)-1), startCallCount(0), startCallInfo(nullptr), bailOutInstr(nullptr),
+        byteCodeUpwardExposedUsed(nullptr), polymorphicCacheIndex((uint)-1), startCallCount(0), startCallInfo(nullptr), bailOutInstr(nullptr), bailInInstr(nullptr),
         totalOutParamCount(0), argOutSyms(nullptr), bailOutRecord(nullptr), wasCloned(false), isInvertedBranch(false), sharedBailOutKind(true), isLoopTopBailOutInfo(false), canDeadStore(true),
         outParamInlinedArgSlot(nullptr), liveVarSyms(nullptr), liveLosslessInt32Syms(nullptr), liveFloat64Syms(nullptr),
         branchConditionOpnd(nullptr),
@@ -160,6 +160,8 @@ class BailOutInfo
     // 2) After we generated bailout, this becomes label instr. In case of shared bailout other instrs JMP to this label.
     IR::Instr * bailOutInstr;
 
+    IR::GeneratorBailInInstr * bailInInstr;
+
 #if ENABLE_DEBUG_CONFIG_OPTIONS
     Js::OpCode bailOutOpcode;
 #endif

lib/Backend/IR.h

@@ -1117,53 +1117,20 @@ class PragmaInstr : public Instr
 class GeneratorBailInInstr : public LabelInstr
 {
 private:
-    GeneratorBailInInstr(JitArenaAllocator* allocator, IR::Instr* yieldInstr):
-        LabelInstr(allocator), allocator(allocator), yieldInstr(yieldInstr), upwardExposedUses(allocator)
+    GeneratorBailInInstr(JitArenaAllocator* allocator, IR::Instr* yieldInstr) :
+        LabelInstr(allocator),
+        yieldInstr(yieldInstr),
+        upwardExposedUses(allocator)
     {
         Assert(yieldInstr != nullptr && yieldInstr->m_opcode == Js::OpCode::Yield);
-        this->usedCapturedValues = JitAnew(allocator, CapturedValues);
     }
 
-    JitArenaAllocator* const allocator;
-    IR::Instr* const yieldInstr;
-    CapturedValues* usedCapturedValues;
+public:
+    IR::Instr* yieldInstr;
+    CapturedValues capturedValues;
     BVSparse<JitArenaAllocator> upwardExposedUses;
 
-public:
     static GeneratorBailInInstr* New(IR::Instr* yieldInstr, Func* func);
-
-    IR::Instr* GetYieldInstr() const
-    {
-        return this->yieldInstr;
-    }
-
-    const CapturedValues& GetCapturedValues() const
-    {
-        return *this->usedCapturedValues;
-    }
-
-    const BVSparse<JitArenaAllocator>& GetUpwardExposedUses() const
-    {
-        return this->upwardExposedUses;
-    }
-
-    void SetCopyPropSyms(const SListBase<CopyPropSyms>& copyPropSyms)
-    {
-        this->usedCapturedValues->copyPropSyms.Clear(this->allocator);
-        copyPropSyms.CopyTo(this->allocator , this->usedCapturedValues->copyPropSyms);
-    }
-
-    void SetConstantValues(const SListBase<ConstantStackSymValue>& constantValues)
-    {
-        this->usedCapturedValues->constantValues.Clear(this->allocator);
-        constantValues.CopyTo(this->allocator, this->usedCapturedValues->constantValues);
-    }
-
-    void SetUpwardExposedUses(const BVSparse<JitArenaAllocator>& other)
-    {
-        this->upwardExposedUses.ClearAll();
-        this->upwardExposedUses.Or(&other);
-    }
 };
 
 template <typename InstrType>

lib/Backend/IRBuilder.cpp

@@ -1900,15 +1900,18 @@ IRBuilder::BuildReg2(Js::OpCode newOpcode, uint32 offset, Js::RegSlot R0, Js::Re
     case Js::OpCode::Yield:
         instr = IR::Instr::New(newOpcode, dstOpnd, src1Opnd, m_func);
         this->AddInstr(instr, offset);
-        this->m_lastInstr = instr->ConvertToBailOutInstr(instr, IR::BailOutForGeneratorYield);
+        IR::Instr* yieldInstr = instr->ConvertToBailOutInstr(instr, IR::BailOutForGeneratorYield);
+        this->m_lastInstr = yieldInstr;
 
         // This label indicates the bail-in section that we will jump to from the generator jump table
-        IR::LabelInstr* bailInLabel = IR::GeneratorBailInInstr::New(this->m_lastInstr /* yieldInstr */, m_func);
+        auto* bailInLabel = IR::GeneratorBailInInstr::New(yieldInstr, m_func);
         bailInLabel->m_hasNonBranchRef = true;              // set to true so that we don't move this label around
         LABELNAMESET(bailInLabel, "GeneratorBailInLabel");
         this->AddInstr(bailInLabel, offset);
         this->m_func->AddYieldOffsetResumeLabel(nextOffset, bailInLabel);
 
+        yieldInstr->GetBailOutInfo()->bailInInstr = bailInLabel;
+
 #ifdef ENABLE_DEBUG_CONFIG_OPTIONS
         if (PHASE_TRACE(Js::Phase::BailInPhase, this->m_func))
         {

lib/Backend/LinearScan.cpp

@@ -5004,7 +5004,7 @@ void LinearScan::GeneratorBailIn::SpillRegsForBailIn()
 //
 IR::Instr* LinearScan::GeneratorBailIn::GenerateBailIn(IR::GeneratorBailInInstr* bailInInstr)
 {
-    BailOutInfo* bailOutInfo = bailInInstr->GetYieldInstr()->GetBailOutInfo();
+    BailOutInfo* bailOutInfo = bailInInstr->yieldInstr->GetBailOutInfo();
 
     Assert(!bailOutInfo->capturedValues || bailOutInfo->capturedValues->constantValues.Empty());
     Assert(!bailOutInfo->capturedValues || bailOutInfo->capturedValues->copyPropSyms.Empty());
@@ -5045,14 +5045,14 @@ IR::Instr* LinearScan::GeneratorBailIn::GenerateBailIn(IR::GeneratorBailInInstr*
 
     this->BuildBailInSymbolList(
         *bailOutInfo->byteCodeUpwardExposedUsed,
-        bailInInstr->GetUpwardExposedUses(),
-        bailInInstr->GetCapturedValues()
+        bailInInstr->upwardExposedUses,
+        bailInInstr->capturedValues
     );
 
     this->InsertRestoreSymbols(
         *bailOutInfo->byteCodeUpwardExposedUsed,
-        bailInInstr->GetUpwardExposedUses(),
-        bailInInstr->GetCapturedValues(),
+        bailInInstr->upwardExposedUses,
+        bailInInstr->capturedValues,
         insertionPoint
     );
     Assert(!this->func->IsStackArgsEnabled());

lib/Backend/Lower.cpp

@@ -29085,7 +29085,7 @@ void Lowerer::LowerGeneratorHelper::InsertNullOutGeneratorFrameInEpilogue(IR::La
 
     IR::IndirOpnd* indirOpnd = IR::IndirOpnd::New(dstOpnd, Js::JavascriptGenerator::GetFrameOffset(), TyMachPtr, this->func);
     IR::AddrOpnd* addrOpnd = IR::AddrOpnd::NewNull(this->func);
-    InsertMove(indirOpnd, addrOpnd, insertionPoint);
+    InsertMove(indirOpnd, addrOpnd, insertionPoint, false /* generateWriteBarrier */);
 }
 
 void

test/es6/generator-jit-bugs.js

@@ -0,0 +1,20 @@
+//-------------------------------------------------------------------------------------------------------
+// Copyright (C) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
+//-------------------------------------------------------------------------------------------------------
+
+// Test 1 - const that is unused/replaced with undefined
+function* foo() {
+    const temp2 = null;
+    while (true) {
+        yield temp2;
+    }
+}
+
+const gen = foo();
+
+gen.next();
+gen.next();
+gen.next();
+
+print("Pass");

test/es6/rlexe.xml

@@ -126,6 +126,13 @@
       <tags>exclude_dynapogo</tags>
     </default>
   </test>
+  <test>
+    <default>
+      <files>generator-jit-bugs.js</files>
+      <compile-flags>-JitES6Generators</compile-flags>
+      <tags>exclude_nonative</tags>
+    </default>
+  </test>
   <test>
     <default>
       <files>proto_basic.js</files>