Bug 2065986 - pt3. Use bailout stub info to reconstruct stack r=jandem,iain

The bailout process is modified to use bailout stub information to rebuild
the stack one frame at a time instead of the existing bulk copy.

- The bailout tail uses the bailout stub information to setup the bailout
  registers, copy contents of frames (besides the return address) and jumps
  to the appropriate bailout stub.
- The bailout stub is setup so that simply calling back into the bailout
  tail materializes the expected return address - and the bailout tail can
  continue with the next frame.

Differential Revision: https://phabricator.services.mozilla.com/D320793

js/src/jit/Bailouts.h

@@ -70,7 +70,7 @@ namespace jit {
 // 1.  Push a snapshot index and the frame size to the native stack.
 // 2.  Spill all registers.
 // 3.  Call js::jit::Bailout to reconstruct the baseline frame(s).
-// 4.  memmove() those to the right place on the native stack.
+// 4.  Copy the frame(s) to the right place on the native stack.
 // 5.  Jump into the baseline interpreter.
 //
 // When C++ code invalidates Ion code, we do on-stack invalidation, or OSI, to
@@ -97,6 +97,72 @@ namespace jit {
 //     _deopt:
 //       push imm(FrameSize)
 //       call _global_bailout_handler
+//
+// ## Bailout stack reconstruction
+//
+// BaselineStackBuilder uses the Snapshot to construct a copy (on the heap) of
+// the stack that will replace the IonFrame. Each recovered Ion inline frame
+// becomes a BaselineJS frame that resumes in the Baseline Interpreter. For each
+// reconstructed inline call, the caller and callee BaselineJS frames are
+// connected by a BaselineStub frame that resumes in a FallbackIC stub:
+//
+//    +============================+
+//    | Original IonFrame's Caller |
+//    +----------------------------+
+//    |      BaselineJS Frame      |
+//    +----------------------------+
+//    |     BaselineStub Frame     |
+//    +----------------------------+
+//    |      BaselineJS Frame      |
+//    +----------------------------+
+//    |            ...             |
+//    +----------------------------+
+//    |  Final BaselineJS Frame    |
+//    +----------------------------+
+//
+// The return addresses installed in these frames during reconstruction are to
+// specific known return points in Bailout Stubs emitted by:
+//
+// - BaselineInterpreterGenerator::emitICBailoutStub and
+// - FallbackICCodeCompiler::emitBailoutStub
+//
+// Hardware shadow stacks (e.g., CET on x86_64) use a hidden secondary stack to
+// verify that return addresses were pushed by earlier calls. To be compatible
+// with this, we can't simply write an address onto the stack. Each return
+// address in the reconstructed stack must be pushed by an actual call. We
+// therefore rebuild the stack one frame at a time.
+//
+// The BaselineStackBuilder rebuilds the stack one frame at a time, separated by
+// these return address slots. Along with the contents of each frame, it appends
+// a corresponding BailoutStubInfo containing two fields:
+//
+// - frameBoundary: the native stack boundary down to which (assuming the stack
+// grows downwards) the current frame's contents must be copied. This excludes
+// the return address slot representing the call into the next (outer) frame.
+//
+// - bailoutStub: the address of the bailout stub that actually contains the
+// return address for the reconstructed frame and pushes it on the stack.
+//
+// The last BailoutStubInfo corresponding to the final BaselineJS frame where
+// execution resumes in the BaselineInterpreter has a NULL bailoutStub since it
+// doesn't call into another frame and so doesn't have a return address slot.
+//
+// The bailoutTail generated by MacroAssembler::generateBailoutTail loops
+// through these BailoutStubInfo entries and copies the reconstructed frames on
+// the native stack from the outer-most to inner-most frame. The process begins
+// at the return address slot of the original IonFrame's caller.
+//
+// For each BailoutStubInfo entry, the bailoutTail:
+//
+// 1. Copies the contents of the frame into the native stack up to
+// frameBoundary.
+//
+// 2. If bailoutStub is NULL the bailout stack has been fully copied and the
+// bailoutTail proceeds with a call to js::jit::FinishBailoutToBaseline.
+//
+// 3. Otherwise, jump to the bailout stub at bailoutStub. The bailout stub
+// performs an indirect call back to bailoutTail to push the expected return
+// address on the stack and continue copying the next frame.
 
 // BailoutStack is an architecture specific pointer to the stack, given by the
 // bailout handler.

js/src/jit/BaselineCodeGen.cpp

@@ -7157,8 +7157,9 @@ void BaselineInterpreterGenerator::emitICBailoutStub() {
   icReturn.bind(entry.offset);
   entry.bailoutStubOffset = masm.currentOffset();
   // The bailoutTail jumps here when performing bailout stack
-  // reconstruction.
-  // TODO: call the bailout stub handler.
+  // reconstruction. The Baseline frame has been rebuilt.
+  // Only the return address remains to be pushed.
+  entry.offset = masm.call(BailoutStubHandlerReg).offset();
   masm.jump(&icReturn);
 }
 

js/src/jit/BaselineIC.cpp

@@ -650,9 +650,9 @@ void FallbackICCodeCompiler::enterStubFrame(MacroAssembler& masm,
 void FallbackICCodeCompiler::emitBailoutStub(BailoutReturnKind kind) {
   code.initBailoutStubOffset(kind, masm.currentOffset());
   // The bailoutTail jumps here when performing bailout stack
-  // reconstruction.
-  // TODO: call the bailout stub handler.
-  code.initBailoutReturnOffset(kind, masm.currentOffset());
+  // reconstruction. The BaselineStub frame has been rebuilt.
+  // Only the return address remains to be pushed.
+  code.initBailoutReturnOffset(kind, masm.call(BailoutStubHandlerReg).offset());
 }
 
 void FallbackICCodeCompiler::assumeStubFrame() {

js/src/jit/MacroAssembler.cpp

@@ -4115,42 +4115,85 @@ void MacroAssembler::generateBailoutTail(Register scratch,
     AllocatableGeneralRegisterSet regs(GeneralRegisterSet::All());
     MOZ_ASSERT_IF(!IsHiddenSP(getStackPointer()),
                   !regs.has(AsRegister(getStackPointer())));
+    regs.take(scratch);
     regs.take(bailoutInfo);
+    regs.take(BailoutStubHandlerReg);
 
-    Register temp = regs.takeAny();
+    Register copyCur = regs.takeAny();
+    Register copyEnd = regs.takeAny();
+    Register stubInfo = regs.takeAny();
 
 #ifdef DEBUG
     // Assert the stack pointer points to the JitFrameLayout header. Copying
     // starts here.
     Label ok;
     loadPtr(Address(bailoutInfo, offsetof(BaselineBailoutInfo, incomingStack)),
-            temp);
-    branchStackPtr(Assembler::Equal, temp, &ok);
+            scratch);
+    branchStackPtr(Assembler::Equal, scratch, &ok);
     assumeUnreachable("Unexpected stack pointer value");
     bind(&ok);
 #endif
 
-    Register copyCur = regs.takeAny();
-    Register copyEnd = regs.takeAny();
-
     // Copy data onto stack.
     loadPtr(Address(bailoutInfo, offsetof(BaselineBailoutInfo, copyStackTop)),
             copyCur);
-    loadPtr(
-        Address(bailoutInfo, offsetof(BaselineBailoutInfo, copyStackBottom)),
-        copyEnd);
-    {
-      Label copyLoop;
-      Label endOfCopy;
-      bind(&copyLoop);
-      branchPtr(Assembler::BelowOrEqual, copyCur, copyEnd, &endOfCopy);
-      subPtr(Imm32(sizeof(uintptr_t)), copyCur);
-      subFromStackPtr(Imm32(sizeof(uintptr_t)));
-      loadPtr(Address(copyCur, 0), temp);
-      storePtr(temp, Address(getStackPointer(), 0));
-      jump(&copyLoop);
-      bind(&endOfCopy);
-    }
+    // Instead of directly copying the entire bailout stack in one go, we
+    // materialize the frames one at a time using calls to bailout stubs for the
+    // return addresses.
+    computeEffectiveAddress(Address(bailoutInfo, sizeof(BaselineBailoutInfo)),
+                            stubInfo);
+
+    CodeLabel bailoutStubHandler;
+    // Pin the BailoutStubHandlerReg that each bailout stub will call.
+    mov(&bailoutStubHandler, BailoutStubHandlerReg);
+
+#ifdef JS_USE_LINK_REGISTER
+    // On the first iteration into the bailout stub handler,
+    // the link register does not need to be pushed.
+    Label copyFrame;
+    jump(&copyFrame);
+#endif
+
+    bind(&bailoutStubHandler);
+    addCodeLabel(bailoutStubHandler);
+
+#ifdef JS_USE_LINK_REGISTER
+    // The call from the bailout stub to the bailout stub handler
+    // does not actually push the return address onto the stack.
+    pushReturnAddress();
+    bind(&copyFrame);
+#endif
+
+    // Bailout segment handler starts here. The baseline bailout stubs
+    // repeatedly call into this to push the return address and begin
+    // copying the next frame on the stack.
+    loadPtr(Address(stubInfo, offsetof(BailoutStubInfo, frameBoundary)),
+            copyEnd);
+
+    Label copyLoop;
+    Label endOfCopy;
+    bind(&copyLoop);
+    branchStackPtr(Assembler::BelowOrEqual, copyEnd, &endOfCopy);
+    subPtr(Imm32(sizeof(uintptr_t)), copyCur);
+    subFromStackPtr(Imm32(sizeof(uintptr_t)));
+    loadPtr(Address(copyCur, 0), scratch);
+    storePtr(scratch, Address(getStackPointer(), 0));
+    jump(&copyLoop);
+
+    bind(&endOfCopy);
+    // all of the frame's contents have been copied except the return address.
+    loadPtr(Address(stubInfo, offsetof(BailoutStubInfo, bailoutStub)), scratch);
+    // move to the next bailout stub info.
+    addPtr(Imm32(sizeof(BailoutStubInfo)), stubInfo);
+    Label copyDone;
+    // A null stubAddr marks the final stack boundary.
+    branchTestPtr(Assembler::Zero, scratch, scratch, &copyDone);
+
+    // The call in each bailout stub recreates the return address skipped here.
+    subPtr(Imm32(sizeof(uintptr_t)), copyCur);
+    jump(scratch);
+
+    bind(&copyDone);
 
     loadPtr(Address(bailoutInfo, offsetof(BaselineBailoutInfo, resumeFramePtr)),
             FramePointer);
@@ -4168,7 +4211,7 @@ void MacroAssembler::generateBailoutTail(Register scratch,
 
     // Call a stub to free allocated memory and create arguments objects.
     using Fn = bool (*)(BaselineBailoutInfo* bailoutInfoArg);
-    setupUnalignedABICall(temp);
+    setupUnalignedABICall(scratch);
     passABIArg(bailoutInfo);
     callWithABI<Fn, FinishBailoutToBaseline>(
         ABIType::General, CheckUnsafeCallWithABI::DontCheckHasExitFrame);

js/src/jit/arm/Assembler-arm.h

@@ -275,6 +275,10 @@ static constexpr Register RegExpSearcherRegExpReg = CallTempReg0;
 static constexpr Register RegExpSearcherStringReg = CallTempReg1;
 static constexpr Register RegExpSearcherLastIndexReg = CallTempReg2;
 
+// Register used by the bailout tail and bailout stubs during stack
+// reconstruction.
+static constexpr Register BailoutStubHandlerReg = CallTempReg0;
+
 static constexpr FloatRegister d0 = {FloatRegisters::d0, VFPRegister::Double};
 static constexpr FloatRegister d1 = {FloatRegisters::d1, VFPRegister::Double};
 static constexpr FloatRegister d2 = {FloatRegisters::d2, VFPRegister::Double};

js/src/jit/arm64/Assembler-arm64.h

@@ -402,6 +402,10 @@ static constexpr Register RegExpSearcherRegExpReg = CallTempReg0;
 static constexpr Register RegExpSearcherStringReg = CallTempReg1;
 static constexpr Register RegExpSearcherLastIndexReg = CallTempReg2;
 
+// Register used by the bailout tail and bailout stubs during stack
+// reconstruction.
+static constexpr Register BailoutStubHandlerReg = CallTempReg0;
+
 static constexpr Register JSReturnReg_Type = r3;
 static constexpr Register JSReturnReg_Data = r2;
 

js/src/jit/loong64/Assembler-loong64.h

@@ -176,6 +176,10 @@ static constexpr Register RegExpSearcherRegExpReg = CallTempReg0;
 static constexpr Register RegExpSearcherStringReg = CallTempReg1;
 static constexpr Register RegExpSearcherLastIndexReg = CallTempReg2;
 
+// Register used by the bailout tail and bailout stubs during stack
+// reconstruction.
+static constexpr Register BailoutStubHandlerReg = CallTempReg0;
+
 static constexpr Register JSReturnReg_Type = a3;
 static constexpr Register JSReturnReg_Data = a2;
 static constexpr Register JSReturnReg = a2;

js/src/jit/mips-shared/Assembler-mips-shared.h

@@ -128,6 +128,10 @@ static constexpr Register RegExpSearcherRegExpReg = CallTempReg0;
 static constexpr Register RegExpSearcherStringReg = CallTempReg1;
 static constexpr Register RegExpSearcherLastIndexReg = CallTempReg2;
 
+// Register used by the bailout tail and bailout stubs during stack
+// reconstruction.
+static constexpr Register BailoutStubHandlerReg = CallTempReg0;
+
 static constexpr uint32_t CodeAlignment = 8;
 
 /* clang-format off */

js/src/jit/none/Assembler-none.h

@@ -67,6 +67,8 @@ static constexpr Register RegExpSearcherRegExpReg{Registers::invalid_reg};
 static constexpr Register RegExpSearcherStringReg{Registers::invalid_reg};
 static constexpr Register RegExpSearcherLastIndexReg{Registers::invalid_reg};
 
+static constexpr Register BailoutStubHandlerReg{Registers::invalid_reg};
+
 // Uses |invalid_reg2| to avoid static_assert failures.
 static constexpr Register JSReturnReg_Type{Registers::invalid_reg2};
 static constexpr Register JSReturnReg_Data{Registers::invalid_reg2};

js/src/jit/riscv64/Register-riscv64.h

@@ -143,6 +143,10 @@ static constexpr Register RegExpSearcherRegExpReg = CallTempReg0;
 static constexpr Register RegExpSearcherStringReg = CallTempReg1;
 static constexpr Register RegExpSearcherLastIndexReg = CallTempReg2;
 
+// Register used by the bailout tail and bailout stubs during stack
+// reconstruction.
+static constexpr Register BailoutStubHandlerReg = CallTempReg0;
+
 static constexpr Register JSReturnReg_Type{Registers::a3};
 static constexpr Register JSReturnReg_Data{Registers::s2};
 static constexpr Register JSReturnReg{Registers::a2};

js/src/jit/wasm32/Assembler-wasm32.h

@@ -76,6 +76,8 @@ static constexpr Register RegExpSearcherRegExpReg{Registers::invalid_reg};
 static constexpr Register RegExpSearcherStringReg{Registers::invalid_reg};
 static constexpr Register RegExpSearcherLastIndexReg{Registers::invalid_reg};
 
+static constexpr Register BailoutStubHandlerReg{Registers::invalid_reg};
+
 // Uses |invalid_reg2| to avoid static_assert failures.
 static constexpr Register JSReturnReg_Type{Registers::invalid_reg2};
 static constexpr Register JSReturnReg_Data{Registers::invalid_reg2};

js/src/jit/x64/Assembler-x64.h

@@ -178,6 +178,10 @@ static constexpr Register RegExpSearcherRegExpReg = CallTempReg1;
 static constexpr Register RegExpSearcherStringReg = CallTempReg2;
 static constexpr Register RegExpSearcherLastIndexReg = CallTempReg3;
 
+// Register used by the bailout tail and bailout stubs during stack
+// reconstruction.
+static constexpr Register BailoutStubHandlerReg = CallTempReg0;
+
 class ABIArgGenerator : public ABIArgGeneratorShared {
 #if defined(XP_WIN)
   unsigned regIndex_;

js/src/jit/x86/Assembler-x86.h

@@ -165,6 +165,10 @@ static constexpr Register RegExpSearcherRegExpReg = CallTempReg0;
 static constexpr Register RegExpSearcherStringReg = CallTempReg2;
 static constexpr Register RegExpSearcherLastIndexReg = CallTempReg3;
 
+// Register used by the bailout tail and bailout stubs during stack
+// reconstruction.
+static constexpr Register BailoutStubHandlerReg = CallTempReg0;
+
 // GCC stack is aligned on 16 bytes. Ion does not maintain this for internal
 // calls. wasm code does.
 #if defined(__GNUC__) && !defined(__MINGW32__)