Bug 1138391 - LazyLinkStub stops making a call and reuses the parent frame. r=h4writer

js/src/jit/CodeGenerator.cpp

@@ -6143,32 +6143,37 @@ JitRuntime::generateLazyLinkStub(JSContext *cx)
 {
     MacroAssembler masm(cx);
 #ifdef JS_USE_LINK_REGISTER
-    masm.push(lr);
+    masm.pushReturnAddress();
 #endif
 
-    Label call;
     GeneralRegisterSet regs = GeneralRegisterSet::Volatile();
     Register temp0 = regs.takeAny();
 
-    masm.callWithExitFrame(&call);
-#ifdef JS_USE_LINK_REGISTER
-    // sigh, this should probably attempt to bypass the push lr that starts off the block
-    // but oh well.
-    masm.pop(lr);
-#endif
-    masm.jump(ReturnReg);
+    // The caller did not push an exit frame on the stack, it pushed a
+    // JitFrameLayout.  We modify the descriptor to be a valid exit frame and
+    // restore it once the lazy link is complete.
+    Address descriptor(StackPointer, CommonFrameLayout::offsetOfDescriptor());
+    size_t convertToExitFrame = JitFrameLayout::Size() - ExitFrameLayout::Size();
+    masm.addPtr(Imm32(convertToExitFrame << FRAMESIZE_SHIFT), descriptor);
+
+    masm.enterFakeExitFrame(LazyLinkExitFrameLayout::Token());
+    masm.PushStubCode();
 
-    masm.bind(&call);
-#ifdef JS_USE_LINK_REGISTER
-        masm.push(lr);
-#endif
-    masm.enterExitFrame();
     masm.setupUnalignedABICall(1, temp0);
     masm.loadJSContext(temp0);
     masm.passABIArg(temp0);
     masm.callWithABI(JS_FUNC_TO_DATA_PTR(void *, LazyLinkTopActivation));
-    masm.leaveExitFrame();
-    masm.retn(Imm32(sizeof(ExitFrameLayout)));
+
+    masm.leaveExitFrame(/* stub code */ sizeof(JitCode*));
+
+    masm.addPtr(Imm32(- (convertToExitFrame << FRAMESIZE_SHIFT)), descriptor);
+
+#ifdef JS_USE_LINK_REGISTER
+    // Restore the return address such that the emitPrologue function of the
+    // CodeGenerator can push it back on the stack with pushReturnAddress.
+    masm.pop(lr);
+#endif
+    masm.jump(ReturnReg);
 
     Linker linker(masm);
     AutoFlushICache afc("LazyLinkStub");

js/src/jit/Ion.cpp

@@ -444,15 +444,12 @@ jit::LazyLinkTopActivation(JSContext *cx)
 
     // First frame should be an exit frame.
     JitFrameIterator it(iter);
-    MOZ_ASSERT(it.type() == JitFrame_Exit);
-
-    // Second frame is the Ion frame.
-    ++it;
-    MOZ_ASSERT(it.type() == JitFrame_IonJS);
+    LazyLinkExitFrameLayout *ll = it.exitFrame()->as<LazyLinkExitFrameLayout>();
+    JSScript *calleeScript = ScriptFromCalleeToken(ll->jsFrame()->calleeToken());
 
     // Get the pending builder from the Ion frame.
-    IonBuilder *builder = it.script()->ionScript()->pendingBuilder();
-    it.script()->setPendingIonBuilder(cx, nullptr);
+    IonBuilder *builder = calleeScript->ionScript()->pendingBuilder();
+    calleeScript->setPendingIonBuilder(cx, nullptr);
 
     AutoEnterAnalysis enterTypes(cx);
     RootedScript script(cx, builder->script());

js/src/jit/JitFrames.cpp

@@ -988,16 +988,14 @@ ReadAllocation(const JitFrameIterator &frame, const LAllocation *a)
 #endif
 
 static void
-MarkThisAndArguments(JSTracer *trc, const JitFrameIterator &frame)
+MarkThisAndArguments(JSTracer *trc, JitFrameLayout *layout)
 {
     // Mark |this| and any extra actual arguments for an Ion frame. Marking of
     // formal arguments is taken care of by the frame's safepoint/snapshot,
     // except when the script might have lazy arguments, in which case we mark
     // them as well.
 
-    JitFrameLayout *layout = frame.jsFrame();
-
-    size_t nargs = frame.numActualArgs();
+    size_t nargs = layout->numActualArgs();
     size_t nformals = 0;
     if (CalleeTokenIsFunction(layout->calleeToken())) {
         JSFunction *fun = CalleeTokenToFunction(layout->calleeToken());
@@ -1014,6 +1012,13 @@ MarkThisAndArguments(JSTracer *trc, const JitFrameIterator &frame)
         gc::MarkValueRoot(trc, &argv[i], "ion-argv");
 }
 
+static void
+MarkThisAndArguments(JSTracer *trc, const JitFrameIterator &frame)
+{
+    JitFrameLayout *layout = frame.jsFrame();
+    MarkThisAndArguments(trc, layout);
+}
+
 #ifdef JS_NUNBOX32
 static inline void
 WriteAllocation(const JitFrameIterator &frame, const LAllocation *a, uintptr_t value)
@@ -1352,6 +1357,16 @@ MarkJitExitFrame(JSTracer *trc, const JitFrameIterator &frame)
         return;
     }
 
+    if (frame.isExitFrameLayout<LazyLinkExitFrameLayout>()) {
+        LazyLinkExitFrameLayout *ll = frame.exitFrame()->as<LazyLinkExitFrameLayout>();
+        JitFrameLayout *layout = ll->jsFrame();
+
+        gc::MarkJitCodeRoot(trc, ll->stubCode(), "lazy-link-code");
+        layout->replaceCalleeToken(MarkCalleeToken(trc, layout->calleeToken()));
+        MarkThisAndArguments(trc, layout);
+        return;
+    }
+
     if (frame.isBareExit()) {
         // Nothing to mark. Fake exit frame pushed for VM functions with
         // nothing to mark on the stack.

js/src/jit/JitFrames.h

@@ -496,6 +496,7 @@ enum ExitFrameTokenValues
     IonOOLPropertyOpExitFrameLayoutToken  = 0x5,
     IonOOLSetterOpExitFrameLayoutToken    = 0x6,
     IonOOLProxyExitFrameLayoutToken       = 0x7,
+    LazyLinkExitFrameLayoutToken          = 0xFE,
     ExitFrameLayoutBareToken              = 0xFF
 };
 
@@ -874,6 +875,43 @@ struct IonDOMMethodExitFrameLayoutTraits {
         offsetof(IonDOMMethodExitFrameLayout, argv_);
 };
 
+// Cannot inherit implementation since we need to extend the top of
+// ExitFrameLayout.
+class LazyLinkExitFrameLayout
+{
+  protected: // silence clang warning about unused private fields
+    JitCode *stubCode_;
+    ExitFooterFrame footer_;
+    JitFrameLayout exit_;
+
+  public:
+    static JitCode *Token() { return (JitCode *) LazyLinkExitFrameLayoutToken; }
+
+    static inline size_t Size() {
+        return sizeof(LazyLinkExitFrameLayout);
+    }
+
+    inline JitCode **stubCode() {
+        return &stubCode_;
+    }
+    inline JitFrameLayout *jsFrame() {
+        return &exit_;
+    }
+    static size_t offsetOfExitFrame() {
+        return offsetof(LazyLinkExitFrameLayout, exit_);
+    }
+};
+
+template <>
+inline LazyLinkExitFrameLayout *
+ExitFrameLayout::as<LazyLinkExitFrameLayout>()
+{
+    MOZ_ASSERT(is<LazyLinkExitFrameLayout>());
+    uint8_t *sp = reinterpret_cast<uint8_t *>(this);
+    sp -= LazyLinkExitFrameLayout::offsetOfExitFrame();
+    return reinterpret_cast<LazyLinkExitFrameLayout *>(sp);
+}
+
 class ICStub;
 
 class BaselineStubFrameLayout : public CommonFrameLayout

js/src/jit/MacroAssembler.h

@@ -863,10 +863,14 @@ class MacroAssembler : public MacroAssemblerSpecific
     void linkExitFrame();
 
   public:
+    void PushStubCode() {
+        exitCodePatch_ = PushWithPatch(ImmWord(-1));
+    }
+
     void enterExitFrame(const VMFunction *f = nullptr) {
         linkExitFrame();
         // Push the ioncode. (Bailout or VM wrapper)
-        exitCodePatch_ = PushWithPatch(ImmWord(-1));
+        PushStubCode();
         // Push VMFunction pointer, to mark arguments.
         Push(ImmPtr(f));
     }
@@ -879,8 +883,8 @@ class MacroAssembler : public MacroAssemblerSpecific
         Push(ImmPtr(nullptr));
     }
 
-    void leaveExitFrame() {
-        freeStack(ExitFooterFrame::Size());
+    void leaveExitFrame(size_t extraFrame = 0) {
+        freeStack(ExitFooterFrame::Size() + extraFrame);
     }
 
     bool hasEnteredExitFrame() const {