Fix KBASE for Lua functions below stitched fast functions.

4 files changed, 79 insertions, 10 deletions

diff --git a/src/vm_arm.dasc b/src/vm_arm.dasc
index 559d20bd..8a48016e 100644
--- a/src/vm_arm.dasc
+++ b/src/vm_arm.dasc
@@ -2109,6 +2109,8 @@ static void build_subroutines(BuildCtx *ctx)
   |   bhi >9                            // More results wanted?
   |
   |  ldr TRACE:RA, [CARG1, CARG3, lsl #2]
+  |  cmp TRACE:RA, #0
+  |  beq ->cont_nop
   |  ldrh RC, TRACE:RA->link
   |  cmp RC, CARG3
   |  beq ->cont_nop                     // Blacklisted.
@@ -2193,7 +2195,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  ldr L, SAVE_L
   |1:
   |  cmp CARG1, #0
-  |  blt >3                             // Check for error from exit.
+  |  blt >9                             // Check for error from exit.
   |   lsl RC, CARG1, #3
   |  ldr LFUNC:CARG2, [BASE, FRAME_FUNC]
   |   str RC, SAVE_MULTRES
@@ -2209,6 +2211,9 @@ static void build_subroutines(BuildCtx *ctx)
   |   ldr INS, [PC], #4
   |     lsl MASKR8, MASKR8, #3          // MASKR8 = 255*8.
   |    st_vmstate CARG4
+  |  cmp OP, #BC_FUNCC+2                // Fast function?
+  |  bhs >4
+  |2:
   |  cmp OP, #BC_FUNCF                  // Function header?
   |  ldr OP, [DISPATCH, OP, lsl #2]
   |   decode_RA8 RA, INS
@@ -2218,7 +2223,20 @@ static void build_subroutines(BuildCtx *ctx)
   |   ldrhs CARG3, [BASE, FRAME_FUNC]
   |  bx OP
   |
-  |3:  // Rethrow error from the right C frame.
+  |4:  // Check frame below fast function.
+  |  ldr CARG1, [BASE, FRAME_PC]
+  |  ands CARG2, CARG1, #FRAME_TYPE
+  |  bne <2                     // Trace stitching continuation?
+  |  // Otherwise set KBASE for Lua function below fast function.
+  |  ldr CARG3, [CARG1, #-4]
+  |  decode_RA8 CARG1, CARG3
+  |  sub CARG2, BASE, CARG1
+  |  ldr LFUNC:CARG3, [CARG2, #-16]
+  |  ldr CARG3, LFUNC:CARG3->field_pc
+  |  ldr KBASE, [CARG3, #PC2PROTO(k)]
+  |  b <2
+  |
+  |9:  // Rethrow error from the right C frame.
   |  rsb CARG2, CARG1, #0
   |  mov CARG1, L
   |  bl extern lj_err_throw             // (lua_State *L, int errcode)
diff --git a/src/vm_mips.dasc b/src/vm_mips.dasc
index 094ffe38..4aaff4b0 100644
--- a/src/vm_mips.dasc
+++ b/src/vm_mips.dasc
@@ -2043,6 +2043,8 @@ static void build_subroutines(BuildCtx *ctx)
   |
   |  addu TMP2, TMP1, TMP2
   |  lw TRACE:TMP2, 0(TMP2)
+  |  beqz TRACE:TMP2, ->cont_nop
+  |.  nop
   |  lhu RD, TRACE:TMP2->link
   |  beq RD, TMP3, ->cont_nop           // Blacklisted.
   |.  load_got lj_dispatch_stitch
@@ -2144,7 +2146,7 @@ static void build_subroutines(BuildCtx *ctx)
   |   addiu DISPATCH, JGL, -GG_DISP2G-32768
   |  sw BASE, L->base
   |1:
-  |  bltz CRET1, >3                     // Check for error from exit.
+  |  bltz CRET1, >9                     // Check for error from exit.
   |.  lw LFUNC:RB, FRAME_FUNC(BASE)
   |    lui TMP3, 0x59c0                 // TOBIT = 2^52 + 2^51 (float).
   |  sll MULTRES, CRET1, 3
@@ -2172,11 +2174,27 @@ static void build_subroutines(BuildCtx *ctx)
   |  jr AT
   |.  decode_RD8b RD
   |2:
+  |  sltiu TMP2, TMP1, (BC_FUNCC+2)*4   // Fast function?
+  |  bnez TMP2, >3
+  |.  lw TMP1, FRAME_PC(BASE)
+  |  // Check frame below fast function.
+  |  andi TMP0, TMP1, FRAME_TYPE
+  |  bnez TMP0, >3                      // Trace stitching continuation?
+  |.  nop
+  |  // Otherwise set KBASE for Lua function below fast function.
+  |  lw TMP2, -4(TMP1)
+  |  decode_RA8a TMP0, TMP2
+  |  decode_RA8b TMP0
+  |  subu TMP1, BASE, TMP0
+  |  lw LFUNC:TMP2, -8+FRAME_FUNC(TMP1)
+  |  lw TMP1, LFUNC:TMP2->pc
+  |  lw KBASE, PC2PROTO(k)(TMP1)
+  |3:
   |  addiu RC, MULTRES, -8
   |  jr AT
   |.  addu RA, RA, BASE
   |
-  |3:  // Rethrow error from the right C frame.
+  |9:  // Rethrow error from the right C frame.
   |  load_got lj_err_throw
   |  negu CARG2, CRET1
   |  call_intern lj_err_throw           // (lua_State *L, int errcode)
diff --git a/src/vm_ppc.dasc b/src/vm_ppc.dasc
index c21f5c43..1d5446c2 100644
--- a/src/vm_ppc.dasc
+++ b/src/vm_ppc.dasc
@@ -2531,6 +2531,8 @@ static void build_subroutines(BuildCtx *ctx)
   |
   |  slwi TMP2, TMP3, 2
   |  lwzx TRACE:TMP2, TMP1, TMP2
+  |  cmpwi TRACE:TMP2, 0
+  |  beq ->cont_nop
   |  lhz RD, TRACE:TMP2->link
   |  cmpw RD, TMP3
   |   cmpwi cr1, RD, 0
@@ -2637,7 +2639,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  stp BASE, L->base
   |1:
   |  cmpwi CARG1, 0
-  |  blt >3                             // Check for error from exit.
+  |  blt >9                             // Check for error from exit.
   |  lwz LFUNC:RB, FRAME_FUNC(BASE)
   |   slwi MULTRES, CARG1, 3
   |    li TMP2, 0
@@ -2673,11 +2675,25 @@ static void build_subroutines(BuildCtx *ctx)
   |   decode_RC8 RC, INS
   |  bctr
   |2:
+  |  cmplwi TMP1, (BC_FUNCC+2)*4        // Fast function?
+  |  blt >3
+  |  // Check frame below fast function.
+  |  lwz TMP1, FRAME_PC(BASE)
+  |  andix. TMP0, TMP1, FRAME_TYPE
+  |  bney >3                            // Trace stitching continuation?
+  |  // Otherwise set KBASE for Lua function below fast function.
+  |  lwz TMP2, -4(TMP1)
+  |  decode_RA8 TMP0, TMP2
+  |  sub TMP1, BASE, TMP0
+  |  lwz LFUNC:TMP2, -12(TMP1)
+  |  lwz TMP1, LFUNC:TMP2->pc
+  |  lwz KBASE, PC2PROTO(k)(TMP1)
+  |3:
   |   subi RC, MULTRES, 8
   |   add RA, RA, BASE
   |  bctr
   |
-  |3:  // Rethrow error from the right C frame.
+  |9:  // Rethrow error from the right C frame.
   |  neg CARG2, CARG1
   |  mr CARG1, L
   |  bl extern lj_err_throw             // (lua_State *L, int errcode)
diff --git a/src/vm_x86.dasc b/src/vm_x86.dasc
index c2f03d80..451c246b 100644
--- a/src/vm_x86.dasc
+++ b/src/vm_x86.dasc
@@ -2697,6 +2697,8 @@ static void build_subroutines(BuildCtx *ctx)
   |  mov RB, TMP1                       // Get previous trace number.
   |  mov RA, [DISPATCH+DISPATCH_J(trace)]
   |  mov TRACE:RD, [RA+RB*4]
+  |  test TRACE:RD, TRACE:RD
+  |  jz ->cont_nop
   |  movzx RD, word TRACE:RD->link
   |  cmp RD, RB
   |  je ->cont_nop                      // Blacklisted.
@@ -2844,7 +2846,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  mov r13, TMPa
   |  mov r12, TMPQ
   |.endif
-  |  test RD, RD; js >3                 // Check for error from exit.
+  |  test RD, RD; js >9                 // Check for error from exit.
   |  mov L:RB, SAVE_L
   |  mov MULTRES, RD
   |  mov LFUNC:KBASE, [BASE-8]
@@ -2860,16 +2862,31 @@ static void build_subroutines(BuildCtx *ctx)
   |  add PC, 4
   |  shr RC, 16
   |  cmp OP, BC_FUNCF                   // Function header?
-  |  jb >2
-  |  mov RC, MULTRES                    // RC/RD holds nres+1.
+  |  jb >3
+  |  cmp OP, BC_FUNCC+2                 // Fast function?
+  |  jae >4
   |2:
+  |  mov RC, MULTRES                    // RC/RD holds nres+1.
+  |3:
   |.if X64
   |  jmp aword [DISPATCH+OP*8]
   |.else
   |  jmp aword [DISPATCH+OP*4]
   |.endif
   |
-  |3:  // Rethrow error from the right C frame.
+  |4:  // Check frame below fast function.
+  |  mov RC, [BASE-4]
+  |  test RC, FRAME_TYPE
+  |  jnz <2                             // Trace stitching continuation?
+  |  // Otherwise set KBASE for Lua function below fast function.
+  |  movzx RC, byte [RC-3]
+  |  not RCa
+  |  mov LFUNC:KBASE, [BASE+RC*8-8]
+  |  mov KBASE, LFUNC:KBASE->pc
+  |  mov KBASE, [KBASE+PC2PROTO(k)]
+  |  jmp <2
+  |
+  |9:  // Rethrow error from the right C frame.
   |  neg RD
   |  mov FCARG1, L:RB
   |  mov FCARG2, RD