[MERGE #6293 @zenparsing] Rework how captured values are written to G…

…eneratorBailInInstrs

Merge pull request #6293 from zenparsing:bailin-captured-values

This PR fixes an existing bug in the generator JIT implementation found by @rhuanjl.

In order to generate bail-in code, the backward pass copies "captured value" information from a yield's `BailOutInfo` into the associated `GeneratorBailInInstr`. The current implementation was copying `bailOutInfo->capturedValues->constantValues` on every pass over the yield instr. However, in a loop the yield instr is passed over twice, and each time entries are moved from `bailOutInfo->capturedValues` into `bailOutInfo->usedCapturedValues`. On the second pass over the instr, the current implementation is overwriting the constant values list with an empty list.

With this change, the constant values are transferred to the bailin instr directly in the code that processes bailout constant values and copy prop syms.

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>