wazevo: shares the stack grow sequence (#1644)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/backend/isa/arm64/abi_go_call.go

@@ -13,7 +13,7 @@ var calleeSavedRegistersPlusLinkRegSorted = []regalloc.VReg{
 }
 
 // CompileGoFunctionTrampoline implements backend.Machine.
-func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool) {
+func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool) []byte {
 	cur := m.allocateInstr()
 	cur.asNop0()
 	m.rootInstr = cur
@@ -145,6 +145,9 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 	ret.asRet(nil)
 	ret.prev = cur
 	cur.next = ret
+
+	m.encode(m.rootInstr)
+	return m.compiler.Buf()
 }
 
 func (m *machine) saveRegistersInExecutionContext(cur *instruction, regs []regalloc.VReg) *instruction {

internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go

@@ -42,32 +42,32 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 			},
 			needModuleContextPtr: true,
 			exp: `
-	str x19, [x0, #0x50]
-	str x20, [x0, #0x60]
-	str x21, [x0, #0x70]
-	str x22, [x0, #0x80]
-	str x23, [x0, #0x90]
-	str x24, [x0, #0xa0]
-	str x25, [x0, #0xb0]
-	str x26, [x0, #0xc0]
-	str x28, [x0, #0xd0]
-	str x30, [x0, #0xe0]
-	str q18, [x0, #0xf0]
-	str q19, [x0, #0x100]
-	str q20, [x0, #0x110]
-	str q21, [x0, #0x120]
-	str q22, [x0, #0x130]
-	str q23, [x0, #0x140]
-	str q24, [x0, #0x150]
-	str q25, [x0, #0x160]
-	str q26, [x0, #0x170]
-	str q27, [x0, #0x180]
-	str q28, [x0, #0x190]
-	str q29, [x0, #0x1a0]
-	str q30, [x0, #0x1b0]
-	str q31, [x0, #0x1c0]
-	str x1, [x0, #0x450]
-	add x15, x0, #0x458
+	str x19, [x0, #0x60]
+	str x20, [x0, #0x70]
+	str x21, [x0, #0x80]
+	str x22, [x0, #0x90]
+	str x23, [x0, #0xa0]
+	str x24, [x0, #0xb0]
+	str x25, [x0, #0xc0]
+	str x26, [x0, #0xd0]
+	str x28, [x0, #0xe0]
+	str x30, [x0, #0xf0]
+	str q18, [x0, #0x100]
+	str q19, [x0, #0x110]
+	str q20, [x0, #0x120]
+	str q21, [x0, #0x130]
+	str q22, [x0, #0x140]
+	str q23, [x0, #0x150]
+	str q24, [x0, #0x160]
+	str q25, [x0, #0x170]
+	str q26, [x0, #0x180]
+	str q27, [x0, #0x190]
+	str q28, [x0, #0x1a0]
+	str q29, [x0, #0x1b0]
+	str q30, [x0, #0x1c0]
+	str q31, [x0, #0x1d0]
+	str x1, [x0, #0x460]
+	add x15, x0, #0x468
 	str d0, [x15], #0x8
 	movz w17, #0x6406, LSL 0
 	str w17, [x0]
@@ -76,31 +76,31 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 	adr x27, #0x1c
 	str x27, [x0, #0x30]
 	exit_sequence x0
-	ldr x19, [x0, #0x50]
-	ldr x20, [x0, #0x60]
-	ldr x21, [x0, #0x70]
-	ldr x22, [x0, #0x80]
-	ldr x23, [x0, #0x90]
-	ldr x24, [x0, #0xa0]
-	ldr x25, [x0, #0xb0]
-	ldr x26, [x0, #0xc0]
-	ldr x28, [x0, #0xd0]
-	ldr x30, [x0, #0xe0]
-	ldr q18, [x0, #0xf0]
-	ldr q19, [x0, #0x100]
-	ldr q20, [x0, #0x110]
-	ldr q21, [x0, #0x120]
-	ldr q22, [x0, #0x130]
-	ldr q23, [x0, #0x140]
-	ldr q24, [x0, #0x150]
-	ldr q25, [x0, #0x160]
-	ldr q26, [x0, #0x170]
-	ldr q27, [x0, #0x180]
-	ldr q28, [x0, #0x190]
-	ldr q29, [x0, #0x1a0]
-	ldr q30, [x0, #0x1b0]
-	ldr q31, [x0, #0x1c0]
-	add x15, x0, #0x458
+	ldr x19, [x0, #0x60]
+	ldr x20, [x0, #0x70]
+	ldr x21, [x0, #0x80]
+	ldr x22, [x0, #0x90]
+	ldr x23, [x0, #0xa0]
+	ldr x24, [x0, #0xb0]
+	ldr x25, [x0, #0xc0]
+	ldr x26, [x0, #0xd0]
+	ldr x28, [x0, #0xe0]
+	ldr x30, [x0, #0xf0]
+	ldr q18, [x0, #0x100]
+	ldr q19, [x0, #0x110]
+	ldr q20, [x0, #0x120]
+	ldr q21, [x0, #0x130]
+	ldr q22, [x0, #0x140]
+	ldr q23, [x0, #0x150]
+	ldr q24, [x0, #0x160]
+	ldr q25, [x0, #0x170]
+	ldr q26, [x0, #0x180]
+	ldr q27, [x0, #0x190]
+	ldr q28, [x0, #0x1a0]
+	ldr q29, [x0, #0x1b0]
+	ldr q30, [x0, #0x1c0]
+	ldr q31, [x0, #0x1d0]
+	add x15, x0, #0x468
 	ldr w27, [x15], #0x8
 	mov w0, w27
 	ldr x27, [x15], #0x8
@@ -121,32 +121,32 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 			},
 			needModuleContextPtr: true,
 			exp: `
-	str x19, [x0, #0x50]
-	str x20, [x0, #0x60]
-	str x21, [x0, #0x70]
-	str x22, [x0, #0x80]
-	str x23, [x0, #0x90]
-	str x24, [x0, #0xa0]
-	str x25, [x0, #0xb0]
-	str x26, [x0, #0xc0]
-	str x28, [x0, #0xd0]
-	str x30, [x0, #0xe0]
-	str q18, [x0, #0xf0]
-	str q19, [x0, #0x100]
-	str q20, [x0, #0x110]
-	str q21, [x0, #0x120]
-	str q22, [x0, #0x130]
-	str q23, [x0, #0x140]
-	str q24, [x0, #0x150]
-	str q25, [x0, #0x160]
-	str q26, [x0, #0x170]
-	str q27, [x0, #0x180]
-	str q28, [x0, #0x190]
-	str q29, [x0, #0x1a0]
-	str q30, [x0, #0x1b0]
-	str q31, [x0, #0x1c0]
-	str x1, [x0, #0x450]
-	add x15, x0, #0x458
+	str x19, [x0, #0x60]
+	str x20, [x0, #0x70]
+	str x21, [x0, #0x80]
+	str x22, [x0, #0x90]
+	str x23, [x0, #0xa0]
+	str x24, [x0, #0xb0]
+	str x25, [x0, #0xc0]
+	str x26, [x0, #0xd0]
+	str x28, [x0, #0xe0]
+	str x30, [x0, #0xf0]
+	str q18, [x0, #0x100]
+	str q19, [x0, #0x110]
+	str q20, [x0, #0x120]
+	str q21, [x0, #0x130]
+	str q22, [x0, #0x140]
+	str q23, [x0, #0x150]
+	str q24, [x0, #0x160]
+	str q25, [x0, #0x170]
+	str q26, [x0, #0x180]
+	str q27, [x0, #0x190]
+	str q28, [x0, #0x1a0]
+	str q29, [x0, #0x1b0]
+	str q30, [x0, #0x1c0]
+	str q31, [x0, #0x1d0]
+	str x1, [x0, #0x460]
+	add x15, x0, #0x468
 	str d0, [x15], #0x8
 	str d1, [x15], #0x8
 	str x2, [x15], #0x8
@@ -158,31 +158,31 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 	adr x27, #0x1c
 	str x27, [x0, #0x30]
 	exit_sequence x0
-	ldr x19, [x0, #0x50]
-	ldr x20, [x0, #0x60]
-	ldr x21, [x0, #0x70]
-	ldr x22, [x0, #0x80]
-	ldr x23, [x0, #0x90]
-	ldr x24, [x0, #0xa0]
-	ldr x25, [x0, #0xb0]
-	ldr x26, [x0, #0xc0]
-	ldr x28, [x0, #0xd0]
-	ldr x30, [x0, #0xe0]
-	ldr q18, [x0, #0xf0]
-	ldr q19, [x0, #0x100]
-	ldr q20, [x0, #0x110]
-	ldr q21, [x0, #0x120]
-	ldr q22, [x0, #0x130]
-	ldr q23, [x0, #0x140]
-	ldr q24, [x0, #0x150]
-	ldr q25, [x0, #0x160]
-	ldr q26, [x0, #0x170]
-	ldr q27, [x0, #0x180]
-	ldr q28, [x0, #0x190]
-	ldr q29, [x0, #0x1a0]
-	ldr q30, [x0, #0x1b0]
-	ldr q31, [x0, #0x1c0]
-	add x15, x0, #0x458
+	ldr x19, [x0, #0x60]
+	ldr x20, [x0, #0x70]
+	ldr x21, [x0, #0x80]
+	ldr x22, [x0, #0x90]
+	ldr x23, [x0, #0xa0]
+	ldr x24, [x0, #0xb0]
+	ldr x25, [x0, #0xc0]
+	ldr x26, [x0, #0xd0]
+	ldr x28, [x0, #0xe0]
+	ldr x30, [x0, #0xf0]
+	ldr q18, [x0, #0x100]
+	ldr q19, [x0, #0x110]
+	ldr q20, [x0, #0x120]
+	ldr q21, [x0, #0x130]
+	ldr q22, [x0, #0x140]
+	ldr q23, [x0, #0x150]
+	ldr q24, [x0, #0x160]
+	ldr q25, [x0, #0x170]
+	ldr q26, [x0, #0x180]
+	ldr q27, [x0, #0x190]
+	ldr q28, [x0, #0x1a0]
+	ldr q29, [x0, #0x1b0]
+	ldr q30, [x0, #0x1c0]
+	ldr q31, [x0, #0x1d0]
+	add x15, x0, #0x468
 	ret
 `,
 		},
@@ -194,31 +194,31 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 				Results: []ssa.Type{ssa.TypeI32},
 			},
 			exp: `
-	str x19, [x0, #0x50]
-	str x20, [x0, #0x60]
-	str x21, [x0, #0x70]
-	str x22, [x0, #0x80]
-	str x23, [x0, #0x90]
-	str x24, [x0, #0xa0]
-	str x25, [x0, #0xb0]
-	str x26, [x0, #0xc0]
-	str x28, [x0, #0xd0]
-	str x30, [x0, #0xe0]
-	str q18, [x0, #0xf0]
-	str q19, [x0, #0x100]
-	str q20, [x0, #0x110]
-	str q21, [x0, #0x120]
-	str q22, [x0, #0x130]
-	str q23, [x0, #0x140]
-	str q24, [x0, #0x150]
-	str q25, [x0, #0x160]
-	str q26, [x0, #0x170]
-	str q27, [x0, #0x180]
-	str q28, [x0, #0x190]
-	str q29, [x0, #0x1a0]
-	str q30, [x0, #0x1b0]
-	str q31, [x0, #0x1c0]
-	add x15, x0, #0x458
+	str x19, [x0, #0x60]
+	str x20, [x0, #0x70]
+	str x21, [x0, #0x80]
+	str x22, [x0, #0x90]
+	str x23, [x0, #0xa0]
+	str x24, [x0, #0xb0]
+	str x25, [x0, #0xc0]
+	str x26, [x0, #0xd0]
+	str x28, [x0, #0xe0]
+	str x30, [x0, #0xf0]
+	str q18, [x0, #0x100]
+	str q19, [x0, #0x110]
+	str q20, [x0, #0x120]
+	str q21, [x0, #0x130]
+	str q22, [x0, #0x140]
+	str q23, [x0, #0x150]
+	str q24, [x0, #0x160]
+	str q25, [x0, #0x170]
+	str q26, [x0, #0x180]
+	str q27, [x0, #0x190]
+	str q28, [x0, #0x1a0]
+	str q29, [x0, #0x1b0]
+	str q30, [x0, #0x1c0]
+	str q31, [x0, #0x1d0]
+	add x15, x0, #0x468
 	str x1, [x15], #0x8
 	orr w17, wzr, #0x2
 	str w17, [x0]
@@ -227,31 +227,31 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 	adr x27, #0x1c
 	str x27, [x0, #0x30]
 	exit_sequence x0
-	ldr x19, [x0, #0x50]
-	ldr x20, [x0, #0x60]
-	ldr x21, [x0, #0x70]
-	ldr x22, [x0, #0x80]
-	ldr x23, [x0, #0x90]
-	ldr x24, [x0, #0xa0]
-	ldr x25, [x0, #0xb0]
-	ldr x26, [x0, #0xc0]
-	ldr x28, [x0, #0xd0]
-	ldr x30, [x0, #0xe0]
-	ldr q18, [x0, #0xf0]
-	ldr q19, [x0, #0x100]
-	ldr q20, [x0, #0x110]
-	ldr q21, [x0, #0x120]
-	ldr q22, [x0, #0x130]
-	ldr q23, [x0, #0x140]
-	ldr q24, [x0, #0x150]
-	ldr q25, [x0, #0x160]
-	ldr q26, [x0, #0x170]
-	ldr q27, [x0, #0x180]
-	ldr q28, [x0, #0x190]
-	ldr q29, [x0, #0x1a0]
-	ldr q30, [x0, #0x1b0]
-	ldr q31, [x0, #0x1c0]
-	add x15, x0, #0x458
+	ldr x19, [x0, #0x60]
+	ldr x20, [x0, #0x70]
+	ldr x21, [x0, #0x80]
+	ldr x22, [x0, #0x90]
+	ldr x23, [x0, #0xa0]
+	ldr x24, [x0, #0xb0]
+	ldr x25, [x0, #0xc0]
+	ldr x26, [x0, #0xd0]
+	ldr x28, [x0, #0xe0]
+	ldr x30, [x0, #0xf0]
+	ldr q18, [x0, #0x100]
+	ldr q19, [x0, #0x110]
+	ldr q20, [x0, #0x120]
+	ldr q21, [x0, #0x130]
+	ldr q22, [x0, #0x140]
+	ldr q23, [x0, #0x150]
+	ldr q24, [x0, #0x160]
+	ldr q25, [x0, #0x170]
+	ldr q26, [x0, #0x180]
+	ldr q27, [x0, #0x190]
+	ldr q28, [x0, #0x1a0]
+	ldr q29, [x0, #0x1b0]
+	ldr q30, [x0, #0x1c0]
+	ldr q31, [x0, #0x1d0]
+	add x15, x0, #0x468
 	ldr w27, [x15], #0x8
 	mov w0, w27
 	ret

internal/engine/wazevo/backend/isa/arm64/instr.go

@@ -676,8 +676,6 @@ func (i *instruction) String() (str string) {
 	switch i.kind {
 	case nop0:
 		str = "nop0"
-	case nop4:
-		panic("TODO")
 	case aluRRR:
 		size := is64SizeBitToSize(i.u3)
 		str = fmt.Sprintf("%s %s, %s, %s", aluOp(i.u1).String(),
@@ -973,8 +971,6 @@ func (i *instruction) asAdr(rd regalloc.VReg, offset int64) {
 const (
 	// nop0 represents a no-op of zero size.
 	nop0 instructionKind = iota + 1
-	// nop4 represents a no-op that is one instruction large.
-	nop4
 	// aluRRR represents an ALU operation with two register sources and a register destination.
 	aluRRR
 	// aluRRRR represents an ALU operation with three register sources and a register destination.

internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go

@@ -12,10 +12,6 @@ func (m *machine) SetupPrologue() {
 	cur := m.rootInstr
 	prevInitInst := cur.next
 
-	if !m.stackBoundsCheckDisabled {
-		cur = m.insertStackBoundsCheck(m.requiredStackSize(), cur)
-	}
-
 	//
 	//                   (high address)                    (high address)
 	//                 +-----------------+               +------------------+
@@ -40,6 +36,10 @@ func (m *machine) SetupPrologue() {
 	str.prev = cur
 	cur = str
 
+	if !m.stackBoundsCheckDisabled {
+		cur = m.insertStackBoundsCheck(m.requiredStackSize(), cur)
+	}
+
 	// Decrement SP if spillSlotSize > 0.
 	if size := m.spillSlotSize; size > 0 {
 		// Check if size is 16-byte aligned.
@@ -251,7 +251,7 @@ func (m *machine) setupEpilogueAfter(cur *instruction) {
 // which always points to the execution context whenever the native code is entered from Go.
 var saveRequiredRegs = []regalloc.VReg{
 	x1VReg, x2VReg, x3VReg, x4VReg, x5VReg, x6VReg, x7VReg,
-	x18VReg, x19VReg, x20VReg, x21VReg, x22VReg, x23VReg, x24VReg, x25VReg, x26VReg, x28VReg, lrVReg,
+	x19VReg, x20VReg, x21VReg, x22VReg, x23VReg, x24VReg, x25VReg, x26VReg, x28VReg, lrVReg,
 	v0VReg, v1VReg, v2VReg, v3VReg, v4VReg, v5VReg, v6VReg, v7VReg,
 	v18VReg, v19VReg, v20VReg, v21VReg, v22VReg, v23VReg, v24VReg, v25VReg, v26VReg, v27VReg, v28VReg, v29VReg, v30VReg, v31VReg,
 }
@@ -319,15 +319,6 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 	cur.next = cbr
 	cur = cbr
 
-	// Save the callee saved and argument registers.
-	cur = m.saveRegistersInExecutionContext(cur, saveRequiredRegs)
-
-	// Save the current stack pointer.
-	cur = m.saveCurrentStackPointer(cur, x0VReg)
-
-	// Set the exit status on the execution context.
-	cur = m.setExitCode(cur, x0VReg, wazevoapi.ExitCodeGrowStack)
-
 	// Set the required stack size and set it to the exec context.
 	{
 		// First load the requiredStackSize into the temporary register,
@@ -352,11 +343,23 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 		cur = setRequiredStackSize
 	}
 
-	// Exit the execution.
-	cur = m.storeReturnAddressAndExit(cur)
+	ldrAddress := m.allocateInstr()
+	ldrAddress.asULoad(operandNR(tmpRegVReg), addressMode{
+		kind: addressModeKindRegUnsignedImm12,
+		rn:   x0VReg, // execution context is always the first argument
+		imm:  wazevoapi.ExecutionContextOffsets.StackGrowCallSequenceAddress.I64(),
+	}, 64)
+	ldrAddress.prev = cur
+	cur.next = ldrAddress
+	cur = ldrAddress
 
-	// After the exit, restore the saved registers.
-	cur = m.restoreRegistersInExecutionContext(cur, saveRequiredRegs)
+	// Then jumps to the stack grow call sequence's address, meaning
+	// transferring the control to the code compiled by CompileStackGrowCallSequence.
+	bl := m.allocateInstr()
+	bl.asCallIndirect(tmpRegVReg, nil)
+	bl.prev = cur
+	cur.next = bl
+	cur = bl
 
 	// Now that we know the entire code, we can finalize how many bytes
 	// we have to skip when the stack size is sufficient.
@@ -370,3 +373,35 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 	cbr.condBrOffsetResolve(cbrOffset)
 	return cur
 }
+
+// CompileStackGrowCallSequence implements backend.Machine.
+func (m *machine) CompileStackGrowCallSequence() []byte {
+	cur := m.allocateInstr()
+	cur.asNop0()
+	m.rootInstr = cur
+
+	// Save the callee saved and argument registers.
+	cur = m.saveRegistersInExecutionContext(cur, saveRequiredRegs)
+
+	// Save the current stack pointer.
+	cur = m.saveCurrentStackPointer(cur, x0VReg)
+
+	// Set the exit status on the execution context.
+	cur = m.setExitCode(cur, x0VReg, wazevoapi.ExitCodeGrowStack)
+
+	// Exit the execution.
+	cur = m.storeReturnAddressAndExit(cur)
+
+	// After the exit, restore the saved registers.
+	cur = m.restoreRegistersInExecutionContext(cur, saveRequiredRegs)
+
+	// Then goes back the original address of this stack grow call.
+	ret := m.allocateInstr()
+	ret.asRet(nil)
+	ret.prev = cur
+	cur.next = ret
+	cur = ret
+
+	m.encode(m.rootInstr)
+	return m.compiler.Buf()
+}

internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go

@@ -174,96 +174,11 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
 	sub x27, sp, x27
 	ldr x11, [x0, #0x28]
 	subs xzr, x27, x11
-	b.ge #0x178
-	str x1, [x0, #0x50]
-	str x2, [x0, #0x60]
-	str x3, [x0, #0x70]
-	str x4, [x0, #0x80]
-	str x5, [x0, #0x90]
-	str x6, [x0, #0xa0]
-	str x7, [x0, #0xb0]
-	str x18, [x0, #0xc0]
-	str x19, [x0, #0xd0]
-	str x20, [x0, #0xe0]
-	str x21, [x0, #0xf0]
-	str x22, [x0, #0x100]
-	str x23, [x0, #0x110]
-	str x24, [x0, #0x120]
-	str x25, [x0, #0x130]
-	str x26, [x0, #0x140]
-	str x28, [x0, #0x150]
-	str x30, [x0, #0x160]
-	str q0, [x0, #0x170]
-	str q1, [x0, #0x180]
-	str q2, [x0, #0x190]
-	str q3, [x0, #0x1a0]
-	str q4, [x0, #0x1b0]
-	str q5, [x0, #0x1c0]
-	str q6, [x0, #0x1d0]
-	str q7, [x0, #0x1e0]
-	str q18, [x0, #0x1f0]
-	str q19, [x0, #0x200]
-	str q20, [x0, #0x210]
-	str q21, [x0, #0x220]
-	str q22, [x0, #0x230]
-	str q23, [x0, #0x240]
-	str q24, [x0, #0x250]
-	str q25, [x0, #0x260]
-	str q26, [x0, #0x270]
-	str q27, [x0, #0x280]
-	str q28, [x0, #0x290]
-	str q29, [x0, #0x2a0]
-	str q30, [x0, #0x2b0]
-	str q31, [x0, #0x2c0]
-	mov x27, sp
-	str x27, [x0, #0x38]
-	orr w17, wzr, #0x1
-	str w17, [x0]
+	b.ge #0x14
 	movz x27, #0xfff0, LSL 0
 	str x27, [x0, #0x40]
-	adr x27, #0x1c
-	str x27, [x0, #0x30]
-	exit_sequence x0
-	ldr x1, [x0, #0x50]
-	ldr x2, [x0, #0x60]
-	ldr x3, [x0, #0x70]
-	ldr x4, [x0, #0x80]
-	ldr x5, [x0, #0x90]
-	ldr x6, [x0, #0xa0]
-	ldr x7, [x0, #0xb0]
-	ldr x18, [x0, #0xc0]
-	ldr x19, [x0, #0xd0]
-	ldr x20, [x0, #0xe0]
-	ldr x21, [x0, #0xf0]
-	ldr x22, [x0, #0x100]
-	ldr x23, [x0, #0x110]
-	ldr x24, [x0, #0x120]
-	ldr x25, [x0, #0x130]
-	ldr x26, [x0, #0x140]
-	ldr x28, [x0, #0x150]
-	ldr x30, [x0, #0x160]
-	ldr q0, [x0, #0x170]
-	ldr q1, [x0, #0x180]
-	ldr q2, [x0, #0x190]
-	ldr q3, [x0, #0x1a0]
-	ldr q4, [x0, #0x1b0]
-	ldr q5, [x0, #0x1c0]
-	ldr q6, [x0, #0x1d0]
-	ldr q7, [x0, #0x1e0]
-	ldr q18, [x0, #0x1f0]
-	ldr q19, [x0, #0x200]
-	ldr q20, [x0, #0x210]
-	ldr q21, [x0, #0x220]
-	ldr q22, [x0, #0x230]
-	ldr q23, [x0, #0x240]
-	ldr q24, [x0, #0x250]
-	ldr q25, [x0, #0x260]
-	ldr q26, [x0, #0x270]
-	ldr q27, [x0, #0x280]
-	ldr q28, [x0, #0x290]
-	ldr q29, [x0, #0x2a0]
-	ldr q30, [x0, #0x2b0]
-	ldr q31, [x0, #0x2c0]
+	ldr x27, [x0, #0x50]
+	bl w27
 `,
 		},
 		{
@@ -272,96 +187,11 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
 	sub x27, sp, #0x10
 	ldr x11, [x0, #0x28]
 	subs xzr, x27, x11
-	b.ge #0x178
-	str x1, [x0, #0x50]
-	str x2, [x0, #0x60]
-	str x3, [x0, #0x70]
-	str x4, [x0, #0x80]
-	str x5, [x0, #0x90]
-	str x6, [x0, #0xa0]
-	str x7, [x0, #0xb0]
-	str x18, [x0, #0xc0]
-	str x19, [x0, #0xd0]
-	str x20, [x0, #0xe0]
-	str x21, [x0, #0xf0]
-	str x22, [x0, #0x100]
-	str x23, [x0, #0x110]
-	str x24, [x0, #0x120]
-	str x25, [x0, #0x130]
-	str x26, [x0, #0x140]
-	str x28, [x0, #0x150]
-	str x30, [x0, #0x160]
-	str q0, [x0, #0x170]
-	str q1, [x0, #0x180]
-	str q2, [x0, #0x190]
-	str q3, [x0, #0x1a0]
-	str q4, [x0, #0x1b0]
-	str q5, [x0, #0x1c0]
-	str q6, [x0, #0x1d0]
-	str q7, [x0, #0x1e0]
-	str q18, [x0, #0x1f0]
-	str q19, [x0, #0x200]
-	str q20, [x0, #0x210]
-	str q21, [x0, #0x220]
-	str q22, [x0, #0x230]
-	str q23, [x0, #0x240]
-	str q24, [x0, #0x250]
-	str q25, [x0, #0x260]
-	str q26, [x0, #0x270]
-	str q27, [x0, #0x280]
-	str q28, [x0, #0x290]
-	str q29, [x0, #0x2a0]
-	str q30, [x0, #0x2b0]
-	str q31, [x0, #0x2c0]
-	mov x27, sp
-	str x27, [x0, #0x38]
-	orr w17, wzr, #0x1
-	str w17, [x0]
+	b.ge #0x14
 	orr x27, xzr, #0x10
 	str x27, [x0, #0x40]
-	adr x27, #0x1c
-	str x27, [x0, #0x30]
-	exit_sequence x0
-	ldr x1, [x0, #0x50]
-	ldr x2, [x0, #0x60]
-	ldr x3, [x0, #0x70]
-	ldr x4, [x0, #0x80]
-	ldr x5, [x0, #0x90]
-	ldr x6, [x0, #0xa0]
-	ldr x7, [x0, #0xb0]
-	ldr x18, [x0, #0xc0]
-	ldr x19, [x0, #0xd0]
-	ldr x20, [x0, #0xe0]
-	ldr x21, [x0, #0xf0]
-	ldr x22, [x0, #0x100]
-	ldr x23, [x0, #0x110]
-	ldr x24, [x0, #0x120]
-	ldr x25, [x0, #0x130]
-	ldr x26, [x0, #0x140]
-	ldr x28, [x0, #0x150]
-	ldr x30, [x0, #0x160]
-	ldr q0, [x0, #0x170]
-	ldr q1, [x0, #0x180]
-	ldr q2, [x0, #0x190]
-	ldr q3, [x0, #0x1a0]
-	ldr q4, [x0, #0x1b0]
-	ldr q5, [x0, #0x1c0]
-	ldr q6, [x0, #0x1d0]
-	ldr q7, [x0, #0x1e0]
-	ldr q18, [x0, #0x1f0]
-	ldr q19, [x0, #0x200]
-	ldr q20, [x0, #0x210]
-	ldr q21, [x0, #0x220]
-	ldr q22, [x0, #0x230]
-	ldr q23, [x0, #0x240]
-	ldr q24, [x0, #0x250]
-	ldr q25, [x0, #0x260]
-	ldr q26, [x0, #0x270]
-	ldr q27, [x0, #0x280]
-	ldr q28, [x0, #0x290]
-	ldr q29, [x0, #0x2a0]
-	ldr q30, [x0, #0x2b0]
-	ldr q31, [x0, #0x2c0]
+	ldr x27, [x0, #0x50]
+	bl w27
 `,
 		},
 	} {
@@ -378,3 +208,98 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
 		})
 	}
 }
+
+func TestMachine_CompileStackGrowCallSequence(t *testing.T) {
+	_, _, m := newSetupWithMockContext()
+	src := m.CompileStackGrowCallSequence()
+	fmt.Println(hex.EncodeToString(src))
+
+	require.Equal(t, `
+	str x1, [x0, #0x60]
+	str x2, [x0, #0x70]
+	str x3, [x0, #0x80]
+	str x4, [x0, #0x90]
+	str x5, [x0, #0xa0]
+	str x6, [x0, #0xb0]
+	str x7, [x0, #0xc0]
+	str x19, [x0, #0xd0]
+	str x20, [x0, #0xe0]
+	str x21, [x0, #0xf0]
+	str x22, [x0, #0x100]
+	str x23, [x0, #0x110]
+	str x24, [x0, #0x120]
+	str x25, [x0, #0x130]
+	str x26, [x0, #0x140]
+	str x28, [x0, #0x150]
+	str x30, [x0, #0x160]
+	str q0, [x0, #0x170]
+	str q1, [x0, #0x180]
+	str q2, [x0, #0x190]
+	str q3, [x0, #0x1a0]
+	str q4, [x0, #0x1b0]
+	str q5, [x0, #0x1c0]
+	str q6, [x0, #0x1d0]
+	str q7, [x0, #0x1e0]
+	str q18, [x0, #0x1f0]
+	str q19, [x0, #0x200]
+	str q20, [x0, #0x210]
+	str q21, [x0, #0x220]
+	str q22, [x0, #0x230]
+	str q23, [x0, #0x240]
+	str q24, [x0, #0x250]
+	str q25, [x0, #0x260]
+	str q26, [x0, #0x270]
+	str q27, [x0, #0x280]
+	str q28, [x0, #0x290]
+	str q29, [x0, #0x2a0]
+	str q30, [x0, #0x2b0]
+	str q31, [x0, #0x2c0]
+	mov x27, sp
+	str x27, [x0, #0x38]
+	orr w17, wzr, #0x1
+	str w17, [x0]
+	adr x27, #0x1c
+	str x27, [x0, #0x30]
+	exit_sequence x0
+	ldr x1, [x0, #0x60]
+	ldr x2, [x0, #0x70]
+	ldr x3, [x0, #0x80]
+	ldr x4, [x0, #0x90]
+	ldr x5, [x0, #0xa0]
+	ldr x6, [x0, #0xb0]
+	ldr x7, [x0, #0xc0]
+	ldr x19, [x0, #0xd0]
+	ldr x20, [x0, #0xe0]
+	ldr x21, [x0, #0xf0]
+	ldr x22, [x0, #0x100]
+	ldr x23, [x0, #0x110]
+	ldr x24, [x0, #0x120]
+	ldr x25, [x0, #0x130]
+	ldr x26, [x0, #0x140]
+	ldr x28, [x0, #0x150]
+	ldr x30, [x0, #0x160]
+	ldr q0, [x0, #0x170]
+	ldr q1, [x0, #0x180]
+	ldr q2, [x0, #0x190]
+	ldr q3, [x0, #0x1a0]
+	ldr q4, [x0, #0x1b0]
+	ldr q5, [x0, #0x1c0]
+	ldr q6, [x0, #0x1d0]
+	ldr q7, [x0, #0x1e0]
+	ldr q18, [x0, #0x1f0]
+	ldr q19, [x0, #0x200]
+	ldr q20, [x0, #0x210]
+	ldr q21, [x0, #0x220]
+	ldr q22, [x0, #0x230]
+	ldr q23, [x0, #0x240]
+	ldr q24, [x0, #0x250]
+	ldr q25, [x0, #0x260]
+	ldr q26, [x0, #0x270]
+	ldr q27, [x0, #0x280]
+	ldr q28, [x0, #0x290]
+	ldr q29, [x0, #0x2a0]
+	ldr q30, [x0, #0x2b0]
+	ldr q31, [x0, #0x2c0]
+	ret
+`, m.Format())
+}

internal/engine/wazevo/backend/machine.go

@@ -9,6 +9,7 @@ import (
 type (
 	// Machine is a backend for a specific ISA machine.
 	Machine interface {
+		// DisableStackCheck disables the stack check for the current compilation for debugging/testing.
 		DisableStackCheck()
 
 		// RegisterInfo returns the set of registers that can be used for register allocation.
@@ -97,6 +98,10 @@ type (
 		Encode()
 
 		// CompileGoFunctionTrampoline compiles the trampoline function  to call a Go function of the given exit code and signature.
-		CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool)
+		CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool) []byte
+
+		// CompileStackGrowCallSequence returns the sequence of instructions shared by all functions to
+		// call the stack grow builtin function.
+		CompileStackGrowCallSequence() []byte
 	}
 )

internal/engine/wazevo/backend/machine_test.go

@@ -24,14 +24,17 @@ type mockMachine struct {
 	rinfo                  *regalloc.RegisterInfo
 }
 
+func (m mockMachine) CompileStackGrowCallSequence() []byte {
+	panic("TODO")
+}
+
 // CompileGoFunctionTrampoline implements Machine.CompileGoFunctionTrampoline.
-func (m mockMachine) CompileGoFunctionTrampoline(wazevoapi.ExitCode, *ssa.Signature, bool) {}
+func (m mockMachine) CompileGoFunctionTrampoline(wazevoapi.ExitCode, *ssa.Signature, bool) []byte {
+	panic("TODO")
+}
 
 // Encode implements Machine.Encode.
-func (m mockMachine) Encode() {
-	// TODO implement me
-	panic("implement me")
-}
+func (m mockMachine) Encode() {}
 
 // ResolveRelocations implements Machine.ResolveRelocations.
 func (m mockMachine) ResolveRelocations(map[ssa.FuncRef]int, []byte, []RelocationInfo) {}

internal/engine/wazevo/call_engine.go

@@ -3,6 +3,7 @@ package wazevo
 import (
 	"context"
 	"encoding/binary"
+	"fmt"
 	"reflect"
 	"unsafe"
 
@@ -59,8 +60,11 @@ type (
 		stackGrowRequiredSize uintptr
 		// memoryGrowTrampolineAddress holds the address of memory grow trampoline function.
 		memoryGrowTrampolineAddress *byte
+		// stackGrowCallSequenceAddress holds the address of stack grow call sequence function.
+		stackGrowCallSequenceAddress *byte
 		// savedRegisters is the opaque spaces for save/restore registers.
 		// We want to align 16 bytes for each register, so we use [64][2]uint64.
+		_              uint64
 		savedRegisters [64][2]uint64
 		// goFunctionCallCalleeModuleContextOpaque is the pointer to the target Go function's moduleContextOpaque.
 		goFunctionCallCalleeModuleContextOpaque uintptr
@@ -114,6 +118,8 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
 		paramResultPtr = &paramResultStack[0]
 	}
 
+	fmt.Printf("stackGrowCallSequenceAddress ===== %#x\n", c.execCtx.stackGrowCallSequenceAddress)
+
 	entrypoint(c.executable, c.execCtxPtr, c.parent.opaquePtr, paramResultPtr, c.stackTop)
 	for {
 		switch ec := c.execCtx.exitCode; ec & wazevoapi.ExitCodeMask {

internal/engine/wazevo/engine.go

@@ -32,6 +32,8 @@ type (
 	builtinFunctions struct {
 		// memoryGrowExecutable is a compiled executable for memory.grow builtin function.
 		memoryGrowExecutable []byte
+		// stackGrowExecutable is a compiled executable for growing stack builtin function.
+		stackGrowExecutable []byte
 	}
 
 	// compiledModule is a compiled variant of a wasm.Module and ready to be used for instantiation.
@@ -341,34 +343,42 @@ func (e *engine) NewModuleEngine(m *wasm.Module, mi *wasm.ModuleInstance) (wasm.
 }
 
 func (e *engine) compileBuiltinFunctions() {
+	machine := newMachine()
+	be := backend.NewCompiler(machine, ssa.NewBuilder())
 	e.builtinFunctions = &builtinFunctions{}
+
 	{
-		machine := newMachine()
-		be := backend.NewCompiler(machine, ssa.NewBuilder())
-		machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeGrowMemory, &ssa.Signature{
+		src := machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeGrowMemory, &ssa.Signature{
 			Params:  []ssa.Type{ssa.TypeI32 /* exec context */, ssa.TypeI32},
 			Results: []ssa.Type{ssa.TypeI32},
 		}, false)
-		be.Encode()
-		src := be.Buf()
-
-		executable, err := platform.MmapCodeSegment(len(src))
-		if err != nil {
-			panic(err)
-		}
-
-		copy(executable, src)
-
-		if runtime.GOARCH == "arm64" {
-			// On arm64, we cannot give all of rwx at the same time, so we change it to exec.
-			if err = platform.MprotectRX(executable); err != nil {
-				panic(err)
-			}
-		}
-		e.builtinFunctions.memoryGrowExecutable = executable
+		e.builtinFunctions.memoryGrowExecutable = mmapExecutable(src)
 	}
 
 	// TODO: table grow, etc.
 
+	be.Init(false)
+	{
+		src := machine.CompileStackGrowCallSequence()
+		e.builtinFunctions.stackGrowExecutable = mmapExecutable(src)
+	}
+
 	// TODO: finalizer.
 }
+
+func mmapExecutable(src []byte) []byte {
+	executable, err := platform.MmapCodeSegment(len(src))
+	if err != nil {
+		panic(err)
+	}
+
+	copy(executable, src)
+
+	if runtime.GOARCH == "arm64" {
+		// On arm64, we cannot give all of rwx at the same time, so we change it to exec.
+		if err = platform.MprotectRX(executable); err != nil {
+			panic(err)
+		}
+	}
+	return executable
+}

internal/engine/wazevo/module_engine.go

@@ -128,6 +128,7 @@ func (m *moduleEngine) NewFunction(index wasm.Index) api.Function {
 	}
 
 	ce.execCtx.memoryGrowTrampolineAddress = &m.parent.builtinFunctions.memoryGrowExecutable[0]
+	ce.execCtx.stackGrowCallSequenceAddress = &m.parent.builtinFunctions.stackGrowExecutable[0]
 	ce.init()
 	return ce
 }

internal/engine/wazevo/wazevo_test.go

@@ -58,6 +58,7 @@ func Test_ExecutionContextOffsets(t *testing.T) {
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.stackPointerBeforeGoCall)), offsets.StackPointerBeforeGrow)
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.stackGrowRequiredSize)), offsets.StackGrowRequiredSize)
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.memoryGrowTrampolineAddress)), offsets.MemoryGrowTrampolineAddress)
+	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.stackGrowCallSequenceAddress)), offsets.StackGrowCallSequenceAddress)
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.savedRegisters))%16, wazevoapi.Offset(0),
 		"SavedRegistersBegin must be aligned to 16 bytes")
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.savedRegisters)), offsets.SavedRegistersBegin)

internal/engine/wazevo/wazevoapi/offsetdata.go

@@ -24,9 +24,10 @@ var ExecutionContextOffsets = ExecutionContextOffsetData{
 	StackPointerBeforeGrow:                  56,
 	StackGrowRequiredSize:                   64,
 	MemoryGrowTrampolineAddress:             72,
-	SavedRegistersBegin:                     80,
-	GoFunctionCallCalleeModuleContextOpaque: 1104,
-	GoFunctionCallStackBegin:                1112,
+	StackGrowCallSequenceAddress:            80,
+	SavedRegistersBegin:                     96,
+	GoFunctionCallCalleeModuleContextOpaque: 1120,
+	GoFunctionCallStackBegin:                1128,
 }
 
 // ExecutionContextOffsetData allows the compilers to get the information about offsets to the fields of wazevo.executionContext,
@@ -52,6 +53,8 @@ type ExecutionContextOffsetData struct {
 	StackGrowRequiredSize Offset
 	// MemoryGrowTrampolineAddress is an offset of `memoryGrowTrampolineAddress` field in wazevo.executionContext
 	MemoryGrowTrampolineAddress Offset
+	// StackGrowCallSequenceAddress is an offset of `stackGrowCallSequenceAddress` field in wazevo.executionContext
+	StackGrowCallSequenceAddress Offset
 	// GoCallReturnAddress is an offset of the first element of `savedRegisters` field in wazevo.executionContext
 	SavedRegistersBegin Offset
 	// GoFunctionCallCalleeModuleContextOpaque is an offset of `goFunctionCallCalleeModuleContextOpaque` field in wazevo.executionContext