Repository: tetratelabs/wazero
Pull request #251 (2022-02-16 05:40:50)

jit(arm64): support for memory size/grow, host functions
part of #187 

mathetake: ok not time to polish 
mathetake: now

Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index c18a51d0..cfa09562 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -175,23 +175,6 @@ func (c *amd64Compiler) compileHostFunction(address wasm.FunctionAddress) error
 	if err := c.callGoFunction(jitCallStatusCodeCallHostFunction, address); err != nil {
 		return err
 	}
-
-	// We consumed the function parameters from the stack after call.
-	for i := 0; i < len(c.f.FunctionType.Type.Params); i++ {
-		c.locationStack.pop()
-	}
-
-	// Also, the function results were pushed by the call.
-	for _, t := range c.f.FunctionType.Type.Results {
-		loc := c.locationStack.pushValueLocationOnStack()
-		switch t {
-		case wasm.ValueTypeI32, wasm.ValueTypeI64:
-			loc.setRegisterType(generalPurposeRegisterTypeInt)
-		case wasm.ValueTypeF32, wasm.ValueTypeF64:
-			loc.setRegisterType(generalPurposeRegisterTypeFloat)
-		}
-	}
-
 	return c.returnFunction()
 }
 
@@ -4241,7 +4224,7 @@ func (c *amd64Compiler) compileMemorySize() error {
 	getMemorySizeInPageUnit.To.Type = obj.TYPE_REG
 	getMemorySizeInPageUnit.To.Reg = loc.register
 	getMemorySizeInPageUnit.From.Type = obj.TYPE_CONST
-	getMemorySizeInPageUnit.From.Offset = 16
+	getMemorySizeInPageUnit.From.Offset = wasm.MemoryPageSizeInBits
 	c.addInstruction(getMemorySizeInPageUnit)
 	return nil
 }
diff --git a/wasm/jit/jit_arm64.go b/wasm/jit/jit_arm64.go
index 45783603..d00a0d15 100644
--- a/wasm/jit/jit_arm64.go
+++ b/wasm/jit/jit_arm64.go
@@ -13,7 +13,6 @@
 package jit
 
 import (
-	"errors"
 	"fmt"
 	"math"
 	"unsafe"
@@ -80,6 +79,8 @@ type arm64Compiler struct {
 	stackPointerCeil uint64
 	// afterAssembleCallback hold the callbacks which are called after assembling native code.
 	afterAssembleCallback []func(code []byte) error
+	// onStackPointerCeilDeterminedCallBack hold a callback which are called when the ceil of stack pointer is determined before generating native code.
+	onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64)
 }
 
 // compile implements compiler.compile for the arm64 architecture.
@@ -92,6 +93,12 @@ func (c *arm64Compiler) compile() (code []byte, staticData compiledFunctionStati
 		stackPointerCeil = c.locationStack.stackPointerCeil
 	}
 
+	// Now that the ceil of stack pointer is determined, we are invoking the callback.
+	// Note: this must be called before Assemble() befolow.
+	if c.onStackPointerCeilDeterminedCallBack != nil {
+		c.onStackPointerCeilDeterminedCallBack(stackPointerCeil)
+	}
+
 	original := c.builder.Assemble()
 
 	for _, cb := range c.afterAssembleCallback {
@@ -161,17 +168,18 @@ func (c *arm64Compiler) markRegisterUnused(regs ...int16) {
 }
 
 // compileConstToRegisterInstruction adds an instruction where source operand is a constant and destination is a register.
-func (c *arm64Compiler) compileConstToRegisterInstruction(instruction obj.As, constValue int64, destinationRegister int16) {
-	applyConst := c.newProg()
-	applyConst.As = instruction
-	applyConst.From.Type = obj.TYPE_CONST
+func (c *arm64Compiler) compileConstToRegisterInstruction(instruction obj.As, constValue int64, destinationRegister int16) (inst *obj.Prog) {
+	inst = c.newProg()
+	inst.As = instruction
+	inst.From.Type = obj.TYPE_CONST
 	// Note: in raw arm64 assembly, immediates larger than 16-bits
 	// are not supported, but the assembler takes care of this and
 	// emits corresponding (at most) 4-instructions to load such large constants.
-	applyConst.From.Offset = constValue
-	applyConst.To.Type = obj.TYPE_REG
-	applyConst.To.Reg = destinationRegister
-	c.addInstruction(applyConst)
+	inst.From.Offset = constValue
+	inst.To.Type = obj.TYPE_REG
+	inst.To.Reg = destinationRegister
+	c.addInstruction(inst)
+	return
 }
 
 // compileMemoryToRegisterInstruction adds an instruction where source operand points a memory location and destination is a register.
@@ -277,9 +285,9 @@ func (c *arm64Compiler) compileTwoRegistersToNoneInstruction(instruction obj.As,
 	c.addInstruction(inst)
 }
 
-func (c *arm64Compiler) compileUnconditionalBranchInstruction() (br *obj.Prog) {
+func (c *arm64Compiler) compilelBranchInstruction(inst obj.As) (br *obj.Prog) {
 	br = c.newProg()
-	br.As = obj.AJMP
+	br.As = inst
 	br.To.Type = obj.TYPE_BRANCH
 	c.addInstruction(br)
 	return
@@ -334,18 +342,19 @@ func (c *arm64Compiler) pushFunctionParams() {
 // compilePreamble implements compiler.compilePreamble for the arm64 architecture.
 func (c *arm64Compiler) compilePreamble() error {
 	// The assembler skips the first instruction so we intentionally add NOP here.
+	// TODO: delete after #233
 	c.compileNOP()
 
 	c.pushFunctionParams()
 
-	// TODO: Add maybeGrowValueStack() .
-
-	// Before executing function body, we must initialize the stack base pointer register
-	// so that we can manipulate the memory stack properly.
-	if err := c.compileReservedStackBasePointerRegisterInitialization(); err != nil {
+	// Check if it's necessary to grow the value stack before entering function body.
+	if err := c.compileMaybeGrowValueStack(); err != nil {
 		return err
 	}
 
+	// We must initialize the stack base pointer register so that we can manipulate the stack properly.
+	c.compileReservedStackBasePointerRegisterInitialization()
+
 	if err := c.compileModuleContextInitialization(); err != nil {
 		return err
 	}
@@ -354,6 +363,63 @@ func (c *arm64Compiler) compilePreamble() error {
 	return nil
 }
 
+// compileMaybeGrowValueStack adds instructions to check the necessity to grow the value stack,
+// and if so, make the builtin function call to do so. These instructions are called in the function's
+// preamble.
+func (c *arm64Compiler) compileMaybeGrowValueStack() error {
+	tmpRegs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 2)
+	if !found {
+		return fmt.Errorf("BUG: all the registers should be free at this point")
+	}
+	tmpX, tmpY := tmpRegs[0], tmpRegs[1]
+
+	// "tmpX = len(engine.valueStack)"
+	c.compileMemoryToRegisterInstruction(
+		arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextValueStackLenOffset,
+		tmpX,
+	)
+
+	// "tmpY = engine.stackBasePointer"
+	c.compileMemoryToRegisterInstruction(
+		arm64.AMOVD,
+		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset,
+		tmpY,
+	)
+
+	// "tmpX = tmpX - tmpY", in other words "tmpX = len(engine.valueStack) - engine.stackBasePointer"
+	c.compileRegisterToRegisterInstruction(
+		arm64.ASUB,
+		tmpY,
+		tmpX,
+	)
+
+	// "tmpY = stackPointerCeil"
+	loadStackPointerCeil := c.compileConstToRegisterInstruction(
+		arm64.AMOVD,
+		math.MaxInt32,
+		tmpY,
+	)
+	// At this point of compilation, we don't know the value of stack pointe ceil,
+	// so we layzily resolve the value later.
+	c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) { loadStackPointerCeil.From.Offset = int64(stackPointerCeil) }
+
+	// Compare tmpX (len(engine.valueStack) - engine.stackBasePointer) and tmpY (engine.stackPointerCeil)
+	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, tmpX, tmpY)
+
+	// If ceil > valueStackLen - stack base pointer, we need to grow the stack by calling builtin Go function.
+	brIfValueStackOK := c.compilelBranchInstruction(arm64.ABLS)
+	if err := c.compileCallGoFunction(jitCallStatusCodeCallBuiltInFunction, builtinFunctionAddressGrowValueStack); err != nil {
+		return err
+	}
+
+	// Otherwise, skip calling it.
+	c.setBranchTargetOnNext(brIfValueStackOK)
+
+	c.locationStack.markRegisterUnused(tmpRegs...)
+	return nil
+}
+
 // returnFunction emits instructions to return from the current function frame.
 // If the current frame is the bottom, the code goes back to the Go code with jitCallStatusCodeReturned status.
 // Otherwise, we branch into the caller's return address.
@@ -384,11 +450,8 @@ func (c *arm64Compiler) compileReturnFunction() error {
 	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, callFramePointerReg, tmpReg)
 
 	// If the values are identical, we return back to the Go code with returned status.
-	brIfNotEqual := c.newProg()
-	brIfNotEqual.As = arm64.ABNE
-	brIfNotEqual.To.Type = obj.TYPE_BRANCH
-	c.addInstruction(brIfNotEqual)
-	if err := c.exit(jitCallStatusCodeReturned); err != nil {
+	brIfNotEqual := c.compilelBranchInstruction(arm64.ABNE)
+	if err := c.compileExitFromNativeCode(jitCallStatusCodeReturned); err != nil {
 		return err
 	}
 
@@ -446,8 +509,8 @@ func (c *arm64Compiler) compileReturnFunction() error {
 	return nil
 }
 
-// exit adds instructions to give the control back to engine.exec with the given status code.
-func (c *arm64Compiler) exit(status jitCallStatusCode) error {
+// compileExitFromNativeCode adds instructions to give the control back to engine.exec with the given status code.
+func (c *arm64Compiler) compileExitFromNativeCode(status jitCallStatusCode) error {
 	// Write the current stack pointer to the engine.stackPointer.
 	c.compileConstToRegisterInstruction(arm64.AMOVD, int64(c.locationStack.sp), reservedRegisterForTemporary)
 	c.compileRegisterToMemoryInstruction(arm64.AMOVD, reservedRegisterForTemporary, reservedRegisterForEngine,
@@ -474,8 +537,15 @@ func (c *arm64Compiler) exit(status jitCallStatusCode) error {
 	return nil
 }
 
+// compileHostFunction implements compiler.compileHostFunction for the arm64 architecture.
 func (c *arm64Compiler) compileHostFunction(address wasm.FunctionAddress) error {
-	return errors.New("TODO: implement compileHostFunction on arm64")
+	// First we must update the location stack to reflect the number of host function inputs.
+	c.pushFunctionParams()
+
+	if err := c.compileCallGoFunction(jitCallStatusCodeCallHostFunction, address); err != nil {
+		return err
+	}
+	return c.compileReturnFunction()
 }
 
 // setLocationStack sets the given valueLocationStack to .locationStack field,
@@ -520,7 +590,7 @@ func (c *arm64Compiler) compileLabel(o *wazeroir.OperationLabel) (skipThisLabel
 
 // compileUnreachable implements compiler.compileUnreachable for the arm64 architecture.
 func (c *arm64Compiler) compileUnreachable() error {
-	return c.exit(jitCallStatusCodeUnreachable)
+	return c.compileExitFromNativeCode(jitCallStatusCodeUnreachable)
 }
 
 // compileSwap implements compiler.compileSwap for the arm64 architecture.
@@ -760,7 +830,7 @@ func (c *arm64Compiler) compileBranchInto(target *wazeroir.BranchTarget) error {
 			targetLabel.initialStack = c.locationStack.clone()
 		}
 
-		br := c.compileUnconditionalBranchInstruction()
+		br := c.compilelBranchInstruction(obj.AJMP)
 		c.assignBranchTarget(labelKey, br)
 		return nil
 	}
@@ -787,15 +857,11 @@ func (c *arm64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
 // compileCall implements compiler.compileCall for the arm64 architecture.
 func (c *arm64Compiler) compileCall(o *wazeroir.OperationCall) error {
 	target := c.f.ModuleInstance.Functions[o.FunctionIndex]
-
-	if err := c.compileCallFunction(target.Address, target.FunctionType.Type); err != nil {
-		return err
-	}
-	return nil
+	return c.compileCallImpl(target.Address, target.FunctionType.Type)
 }
 
-// compileCall implements compiler.compileCall and compiler.compileCallIndirect (TODO) for the arm64 architecture.
-func (c *arm64Compiler) compileCallFunction(addr wasm.FunctionAddress, functype *wasm.FunctionType) error {
+// compileCallImpl implements compiler.compileCall and compiler.compileCallIndirect (TODO) for the arm64 architecture.
+func (c *arm64Compiler) compileCallImpl(addr wasm.FunctionAddress, functype *wasm.FunctionType) error {
 	// TODO: the following code can be generalized for CallIndirect.
 
 	// Release all the registers as our calling convention requires the caller-save.
@@ -810,25 +876,39 @@ func (c *arm64Compiler) compileCallFunction(addr wasm.FunctionAddress, functype
 	c.locationStack.markRegisterUsed(freeRegisters...)
 
 	// Alias for readability.
-	callFrameStackTopAddressRegister, compiledFunctionAddressRegister, oldStackBasePointer,
-		tmp := freeRegisters[0], freeRegisters[1], freeRegisters[2], freeRegisters[3]
+	callFrameStackPointerRegister, callFrameStackTopAddressRegister, compiledFunctionAddressRegister, oldStackBasePointer,
+		tmp := freeRegisters[0], freeRegisters[1], freeRegisters[2], freeRegisters[3], freeRegisters[4]
 
-	// TODO: Check the callframe stack length, and if necessary, grow the call frame stack before jump into the target.
-
-	// "tmp = engine.callFrameStackPointer"
+	// First, we have to check if we need to grow the callFrame stack.
+	//
+	// "callFrameStackPointerRegister = engine.callFrameStackPointer"
 	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
 		reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset,
-		tmp)
-	// "callFrameStackTopAddressRegister = &engine.callFrameStack[0]"
-	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
-		reservedRegisterForEngine, engineGlobalContextCallFrameStackElement0AddressOffset,
-		callFrameStackTopAddressRegister)
-	// "callFrameStackTopAddressRegister += tmp << $callFrameDataSizeMostSignificantSetBit"
-	c.compileAddInstructionWithLeftShiftedRegister(
-		tmp, callFrameDataSizeMostSignificantSetBit,
-		callFrameStackTopAddressRegister,
-		callFrameStackTopAddressRegister,
+		callFrameStackPointerRegister)
+	// "tmp = len(engine.callFrameStack)"
+	c.compileMemoryToRegisterInstruction(
+		arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackLenOffset,
+		tmp,
 	)
+	// Compare tmp(len(engine.callFrameStack)) with callFrameStackPointerRegister(engine.callFrameStackPointer).
+	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, tmp, callFrameStackPointerRegister)
+	brIfCallFrameStackOK := c.compilelBranchInstruction(arm64.ABNE)
+
+	// If these values equal, we need to grow the callFrame stack.
+	if err := c.compileCallGoFunction(jitCallStatusCodeCallBuiltInFunction, builtinFunctionAddressGrowCallFrameStack); err != nil {
+		return err
+	}
+
+	// On the function return, we again have to set engine.callFrameStackPointer into callFrameStackPointerRegister.
+	// "callFrameStackPointerRegister = engine.callFrameStackPointer"
+	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset,
+		callFrameStackPointerRegister)
+
+	// Now that we ensured callFrameStack length is enough.
+	c.setBranchTargetOnNext(brIfCallFrameStackOK)
+	c.compileCalcCallFrameStackTopAddress(callFrameStackPointerRegister, callFrameStackTopAddressRegister)
 
 	// At this point, we have:
 	//
@@ -894,7 +974,7 @@ func (c *arm64Compiler) compileCallFunction(addr wasm.FunctionAddress, functype
 	// 4) Set ra.current so that we can return back to this function properly.
 	//
 	// First, Get the return address into the tmp.
-	c.readInstructionAddress(obj.AJMP, tmp)
+	c.compileReadInstructionAddress(obj.AJMP, tmp)
 	// Then write the address into the callframe.
 	c.compileRegisterToMemoryInstruction(arm64.AMOVD,
 		tmp,
@@ -948,14 +1028,29 @@ func (c *arm64Compiler) compileCallFunction(addr wasm.FunctionAddress, functype
 	return nil
 }
 
-// readInstructionAddress adds an ADR instruction to set the absolute address of "target instruction"
+// compileCalcCallFrameStackTopAddress adds instructions to set the absolute address of
+// engine.callFrameStack[callFrameStackPointerRegister] into destinationRegister.
+func (c *arm64Compiler) compileCalcCallFrameStackTopAddress(callFrameStackPointerRegister, destinationRegister int16) {
+	// "destinationRegister = &engine.callFrameStack[0]"
+	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackElement0AddressOffset,
+		destinationRegister)
+	// "destinationRegister += callFrameStackPointerRegister << $callFrameDataSizeMostSignificantSetBit"
+	c.compileAddInstructionWithLeftShiftedRegister(
+		callFrameStackPointerRegister, callFrameDataSizeMostSignificantSetBit,
+		destinationRegister,
+		destinationRegister,
+	)
+}
+
+// compileReadInstructionAddress adds an ADR instruction to set the absolute address of "target instruction"
 // into destinationRegister. "target instruction" is specified by beforeTargetInst argument and
 // the target is determined by "the instruction right after beforeTargetInst type".
 //
 // For example, if beforeTargetInst == RET and we have the instruction sequence like
 // ADR -> X -> Y -> ... -> RET -> MOV, then the ADR instruction emitted by this function set the absolute
 // address of MOV instruction into the destination register.
-func (c *arm64Compiler) readInstructionAddress(beforeTargetInst obj.As, destinationRegister int16) {
+func (c *arm64Compiler) compileReadInstructionAddress(beforeTargetInst obj.As, destinationRegister int16) {
 	// Emit ADR instruction to read the specified instruction's absolute address.
 	// Note: we cannot emit the "ADR REG, $(target's offset from here)" due to the
 	// incapability of the assembler. Instead, we emit "ADR REG, ." meaning that
@@ -1107,9 +1202,7 @@ func (c *arm64Compiler) compileSelect() error {
 	// At this point, x1 is non-zero register, and x2 is either general purpuse or zero register.
 
 	c.compileTwoRegistersToNoneInstruction(arm64.ACMPW, zeroRegister, cv.register)
-	brIfNotZero := c.newProg()
-	brIfNotZero.As = arm64.ABNE
-	brIfNotZero.To.Type = obj.TYPE_BRANCH
+	brIfNotZero := c.compilelBranchInstruction(arm64.ABNE)
 	c.addInstruction(brIfNotZero)
 
 	// If cv == 0, we move the value of x2 to the x1.register.
@@ -1958,7 +2051,7 @@ func (c *arm64Compiler) compileMemoryAccessOffsetSetup(offsetArg uint32, targetS
 		c.compileConstToRegisterInstruction(arm64.AADD, offsetConst, offsetRegister)
 	} else {
 		// If the offset const is too large, we exit with jitCallStatusCodeMemoryOutOfBounds.
-		err = c.exit(jitCallStatusCodeMemoryOutOfBounds)
+		err = c.compileExitFromNativeCode(jitCallStatusCodeMemoryOutOfBounds)
 		return
 	}
 
@@ -1969,13 +2062,11 @@ func (c *arm64Compiler) compileMemoryAccessOffsetSetup(offsetArg uint32, targetS
 
 	// Check if offsetRegister(= base+offsetArg+targetSizeInBytes) > len(memory.Buffer).
 	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, reservedRegisterForTemporary, offsetRegister)
-	boundsOK := c.newProg()
-	boundsOK.As = arm64.ABLS
-	boundsOK.To.Type = obj.TYPE_BRANCH
-	c.addInstruction(boundsOK)
+	boundsOK := c.compilelBranchInstruction(arm64.ABLS)
 
-	// If offsetRegister(= base+offsetArg+targetSizeInBytes) exceeds the memory length, we exit the function with jitCallStatusCodeMemoryOutOfBounds.
-	if err = c.exit(jitCallStatusCodeMemoryOutOfBounds); err != nil {
+	// If offsetRegister(= base+offsetArg+targetSizeInBytes) exceeds the memory length,
+	//  we exit the function with jitCallStatusCodeMemoryOutOfBounds.
+	if err = c.compileExitFromNativeCode(jitCallStatusCodeMemoryOutOfBounds); err != nil {
 		return
 	}
 
@@ -1985,12 +2076,105 @@ func (c *arm64Compiler) compileMemoryAccessOffsetSetup(offsetArg uint32, targetS
 	return offsetRegister, nil
 }
 
+// compileMemoryGrow implements compileMemoryGrow variants for arm64 architecture.
 func (c *arm64Compiler) compileMemoryGrow() error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	c.maybeCompileMoveTopConditionalToFreeGeneralPurposeRegister()
+
+	if err := c.compileCallGoFunction(jitCallStatusCodeCallBuiltInFunction, builtinFunctionAddressMemoryGrow); err != nil {
+		return err
+	}
+
+	// After return, we re-initialize reserved registers just like preamble of functions.
+	c.compileReservedStackBasePointerRegisterInitialization()
+	c.compileReservedMemoryRegisterInitialization()
+	return nil
 }
 
+// compileMemorySize implements compileMemorySize variants for arm64 architecture.
 func (c *arm64Compiler) compileMemorySize() error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	c.maybeCompileMoveTopConditionalToFreeGeneralPurposeRegister()
+
+	reg, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+	if err != nil {
+		return err
+	}
+
+	// "reg = len(memory.Buffer)"
+	c.compileMemoryToRegisterInstruction(
+		arm64.AMOVD,
+		reservedRegisterForEngine, engineModuleContextMemorySliceLenOffset,
+		reg,
+	)
+
+	// memory.size loads the page size of memory, so we have to divide by the page size.
+	// "reg = reg >> wasm.MemoryPageSizeInBits (== reg / wasm.MemoryPageSize) "
+	c.compileConstToRegisterInstruction(
+		arm64.ALSR,
+		wasm.MemoryPageSizeInBits,
+		reg,
+	)
+
+	c.locationStack.pushValueLocationOnRegister(reg)
+	return nil
+}
+
+// compileCallGoFunction adds instructions to call a Go function whose address equals the addr parameter.
+// jitStatus is set before making call, and it should be either jitCallStatusCodeCallBuiltInFunction or
+// jitCallStatusCodeCallHostFunction.
+func (c *arm64Compiler) compileCallGoFunction(jitStatus jitCallStatusCode, addr wasm.FunctionAddress) error {
+	// Release all the registers as our calling convention requires the caller-save.
+	if err := c.compileReleaseAllRegistersToStack(); err != nil {
+		return err
+	}
+
+	freeRegs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 4)
+	if !found {
+		return fmt.Errorf("BUG: all registers except addrReg should be free at this point")
+	}
+	c.locationStack.markRegisterUsed(freeRegs...)
+
+	// Alias these free tmp registers for readability.
+	tmp, currentCallFrameStackPointerRegister, currentCallFrameTopAddressRegister, returnAddressRegister :=
+		freeRegs[0], freeRegs[1], freeRegs[2], freeRegs[3]
+
+	// Set the target function address to engine.functionCallAddress
+	// "tmp = $addr"
+	c.compileConstToRegisterInstruction(
+		arm64.AMOVD,
+		int64(addr),
+		tmp,
+	)
+	// "[reservedRegisterForEngine + engineExitContextFunctionCallAddressOffset] = tmp"
+	// In other words, "engine.functionCallAddress = tmp (== $addr)"
+	c.compileRegisterToMemoryInstruction(
+		arm64.AMOVD,
+		tmp,
+		reservedRegisterForEngine, engineExitContextFunctionCallAddressOffset,
+	)
+
+	// Next, we have to set the return address into callFrameStack[engine.callFrameStackPointer-1].returnAddress.
+	//
+	// "currentCallFrameStackPointerRegister = engine.callFrameStackPointer"
+	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset,
+		currentCallFrameStackPointerRegister)
+
+	// Set the address of callFrameStack[engine.callFrameStackPointer] into currentCallFrameTopAddressRegister.
+	c.compileCalcCallFrameStackTopAddress(currentCallFrameStackPointerRegister, currentCallFrameTopAddressRegister)
+
+	// Set the return address (after RET in c.exit below) into returnAddressRegister.
+	c.compileReadInstructionAddress(obj.ARET, returnAddressRegister)
+
+	// Write returnAddressRegister into callFrameStack[engine.callFrameStackPointer-1].returnAddress.
+	c.compileRegisterToMemoryInstruction(
+		arm64.AMOVD,
+		returnAddressRegister,
+		// callFrameStack[engine.callFrameStackPointer-1] is below of currentCallFrameTopAddressRegister.
+		currentCallFrameTopAddressRegister, -(callFrameDataSize - callFrameReturnAddressOffset),
+	)
+
+	c.markRegisterUnused(freeRegs...)
+	return c.compileExitFromNativeCode(jitStatus)
 }
 
 // compileConstI32 implements compiler.compileConstI32 for the arm64 architecture.
@@ -2248,7 +2432,7 @@ func (c *arm64Compiler) compileReleaseRegisterToStack(loc *valueLocation) (err e
 
 // compileReservedStackBasePointerRegisterInitialization adds intructions to initialize reservedRegisterForStackBasePointerAddress
 // so that it points to the absolute address of the stack base for this function.
-func (c *arm64Compiler) compileReservedStackBasePointerRegisterInitialization() error {
+func (c *arm64Compiler) compileReservedStackBasePointerRegisterInitialization() {
 	// First, load the address of the first element in the value stack into reservedRegisterForStackBasePointerAddress temporarily.
 	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
 		reservedRegisterForEngine, engineGlobalContextValueStackElement0AddressOffset,
@@ -2265,7 +2449,6 @@ func (c *arm64Compiler) compileReservedStackBasePointerRegisterInitialization()
 	c.compileAddInstructionWithLeftShiftedRegister(
 		reservedRegisterForTemporary, 3, reservedRegisterForStackBasePointerAddress,
 		reservedRegisterForStackBasePointerAddress)
-	return nil
 }
 
 func (c *arm64Compiler) compileReservedMemoryRegisterInitialization() {
@@ -2301,10 +2484,8 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 
 	// If the module instance address stays the same, we could skip the entire code below.
 	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, moduleInstanceAddressRegister, tmpX)
-	brIfMdouleUnchanged := c.newProg()
-	brIfMdouleUnchanged.As = arm64.ABEQ
-	brIfMdouleUnchanged.To.Type = obj.TYPE_BRANCH
-	c.addInstruction(brIfMdouleUnchanged)
+	brIfModuleUnchanged := c.compilelBranchInstruction(arm64.ABEQ)
+	c.addInstruction(brIfModuleUnchanged)
 
 	// Otherwise, we have to update the following fields:
 	// * engine.moduleContext.globalElement0Address
@@ -2379,7 +2560,7 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 
 	// TODO: Update tableElement0Address and tableSliceLen.
 
-	c.setBranchTargetOnNext(brIfMdouleUnchanged)
+	c.setBranchTargetOnNext(brIfModuleUnchanged)
 	c.locationStack.markRegisterUnused(regs...)
 	return nil
 }
diff --git a/wasm/mempage.go b/wasm/mempage.go
index e0e69e0c..529fdb73 100644
--- a/wasm/mempage.go
+++ b/wasm/mempage.go
@@ -8,18 +8,18 @@ const (
 	// MemoryMaxPages is maximum number of pages defined (2^16).
 	// See https://www.w3.org/TR/wasm-core-1/#grow-mem
 	MemoryMaxPages = MemoryPageSize
-	// memoryPageSizeInBit satisfies the relation: "1 << memoryPageSizeInBit == memoryPageSize".
-	memoryPageSizeInBit = uint32(16)
+	// MemoryPageSizeInBits satisfies the relation: "1 << MemoryPageSizeInBits == MemoryPageSize".
+	MemoryPageSizeInBits = 16
 )
 
 // MemoryPagesToBytesNum converts the given pages into the number of bytes contained in these pages.
 func memoryPagesToBytesNum(pages uint32) (bytesNum uint64) {
-	return uint64(pages) << memoryPageSizeInBit
+	return uint64(pages) << MemoryPageSizeInBits
 }
 
 // MemoryPagesToBytesNum converts the given number of bytes into the number of pages.
 func memoryBytesNumToPages(bytesNum uint64) (pages uint32) {
-	return uint32(bytesNum >> memoryPageSizeInBit)
+	return uint32(bytesNum >> MemoryPageSizeInBits)
 }
 
 // Grow extends the memory buffer by "newPages" * memoryPageSize.


Test patch:
diff --git a/wasm/jit/jit_arm64_test.go b/wasm/jit/jit_arm64_test.go
index 2c8337d6..e275aea6 100644
--- a/wasm/jit/jit_arm64_test.go
+++ b/wasm/jit/jit_arm64_test.go
@@ -149,7 +149,7 @@ func TestArm64Compiler_exit(t *testing.T) {
 			expStackPointer := uint64(100)
 			compiler.locationStack.sp = expStackPointer
 			require.NoError(t, err)
-			compiler.exit(s)
+			compiler.compileExitFromNativeCode(s)
 
 			// Compile and execute the code under test.
 			code, _, _, err := compiler.compile()
@@ -274,7 +274,7 @@ func TestArm64Compiler_releaseRegisterToStack(t *testing.T) {
 
 			// Release the register allocated value to the memory stack so that we can see the value after exiting.
 			compiler.compileReleaseRegisterToStack(compiler.locationStack.peek())
-			compiler.exit(jitCallStatusCodeReturned)
+			compiler.compileExitFromNativeCode(jitCallStatusCodeReturned)
 
 			// Generate the code under test.
 			code, _, _, err := compiler.compile()
@@ -353,7 +353,7 @@ func TestArm64Compiler_compileLoadValueOnStackToRegister(t *testing.T) {
 
 			// Release the value to the memory stack so that we can see the value after exiting.
 			compiler.compileReleaseRegisterToStack(loc)
-			compiler.exit(jitCallStatusCodeReturned)
+			compiler.compileExitFromNativeCode(jitCallStatusCodeReturned)
 
 			// Generate the code under test.
 			code, _, _, err := compiler.compile()
@@ -1013,7 +1013,7 @@ func TestArm64Compiler_compileShr(t *testing.T) {
 	})
 }
 
-func TestArm64Compiler_compielePick(t *testing.T) {
+func TestArm64Compiler_compilePick(t *testing.T) {
 	const pickTargetValue uint64 = 12345
 	op := &wazeroir.OperationPick{Depth: 1}
 
@@ -1114,7 +1114,7 @@ func TestArm64Compiler_compielePick(t *testing.T) {
 	}
 }
 
-func TestArm64Compiler_compieleDrop(t *testing.T) {
+func TestArm64Compiler_compileDrop(t *testing.T) {
 	t.Run("range nil", func(t *testing.T) {
 		env := newJITEnvironment()
 		compiler := env.requireNewCompiler(t)
@@ -1308,21 +1308,16 @@ func TestArm64Compiler_compileBr(t *testing.T) {
 		// Emit code for the backward label.
 		backwardLabel := &wazeroir.Label{Kind: wazeroir.LabelKindHeader, FrameID: 0}
 		requireAddLabel(t, compiler, backwardLabel)
-		compiler.exit(jitCallStatusCodeReturned)
+		compiler.compileExitFromNativeCode(jitCallStatusCodeReturned)
 
 		// Now emit the body. First we add NOP so that we can execute code after the target label.
-		nop := compiler.newProg()
-		nop.As = obj.ANOP
-		compiler.addInstruction(nop)
+		nop := compiler.compileNOP()
 
-		err := compiler.compilePreamble()
-		require.NoError(t, err)
-
-		err = compiler.compileBr(&wazeroir.OperationBr{Target: &wazeroir.BranchTarget{Label: backwardLabel}})
+		err := compiler.compileBr(&wazeroir.OperationBr{Target: &wazeroir.BranchTarget{Label: backwardLabel}})
 		require.NoError(t, err)
 
 		// We must not reach the code after Br, so emit the code exiting with unreachable status.
-		compiler.exit(jitCallStatusCodeUnreachable)
+		compiler.compileExitFromNativeCode(jitCallStatusCodeUnreachable)
 
 		code, _, _, err := compiler.compile()
 		require.NoError(t, err)
@@ -1352,11 +1347,11 @@ func TestArm64Compiler_compileBr(t *testing.T) {
 		require.NoError(t, err)
 
 		// We must not reach the code after Br, so emit the code exiting with Unreachable status.
-		compiler.exit(jitCallStatusCodeUnreachable)
+		compiler.compileExitFromNativeCode(jitCallStatusCodeUnreachable)
 
 		// Emit code for the forward label where we emit the expectedValue and then exit.
 		requireAddLabel(t, compiler, forwardLabel)
-		compiler.exit(jitCallStatusCodeReturned)
+		compiler.compileExitFromNativeCode(jitCallStatusCodeReturned)
 
 		code, _, _, err := compiler.compile()
 		require.NoError(t, err)
@@ -1553,15 +1548,15 @@ func TestArm64Compiler_compileBrIf(t *testing.T) {
 
 					err = compiler.compileBrIf(&wazeroir.OperationBrIf{Then: thenBranchTarget, Else: elseBranchTarget})
 					require.NoError(t, err)
-					compiler.exit(unreachableStatus)
+					compiler.compileExitFromNativeCode(unreachableStatus)
 
 					// Emit code for .then label.
 					requireAddLabel(t, compiler, thenBranchTarget.Target.Label)
-					compiler.exit(thenLabelExitStatus)
+					compiler.compileExitFromNativeCode(thenLabelExitStatus)
 
 					// Emit code for .else label.
 					requireAddLabel(t, compiler, elseBranchTarget.Target.Label)
-					compiler.exit(elseLabelExitStatus)
+					compiler.compileExitFromNativeCode(elseLabelExitStatus)
 
 					code, _, _, err := compiler.compile()
 					require.NoError(t, err)
@@ -1600,11 +1595,11 @@ func TestArm64Compiler_readInstructionAddress(t *testing.T) {
 		require.NoError(t, err)
 
 		// Set the acquisition target instruction to the one after JMP.
-		compiler.readInstructionAddress(obj.AJMP, reservedRegisterForTemporary)
+		compiler.compileReadInstructionAddress(obj.AJMP, reservedRegisterForTemporary)
 
-		compiler.exit(jitCallStatusCodeReturned)
+		compiler.compileExitFromNativeCode(jitCallStatusCodeReturned)
 
-		// If generate the code without JMP after readInstructionAddress,
+		// If generate the code without JMP after compileReadInstructionAddress,
 		// the call back added must return error.
 		_, _, _, err = compiler.compile()
 		require.Error(t, err)
@@ -1618,9 +1613,9 @@ func TestArm64Compiler_readInstructionAddress(t *testing.T) {
 		require.NoError(t, err)
 
 		// Set the acquisition target instruction to the one after RET.
-		compiler.readInstructionAddress(obj.ARET, reservedRegisterForTemporary)
+		compiler.compileReadInstructionAddress(obj.ARET, reservedRegisterForTemporary)
 
-		// Add many instruction between the target and readInstructionAddress.
+		// Add many instruction between the target and compileReadInstructionAddress.
 		for i := 0; i < 100; i++ {
 			compiler.compileConstI32(&wazeroir.OperationConstI32{Value: 10})
 		}
@@ -1647,7 +1642,7 @@ func TestArm64Compiler_readInstructionAddress(t *testing.T) {
 		// Set the acquisition target instruction to the one after RET,
 		// and read the absolute address into destinationRegister.
 		const addressReg = reservedRegisterForTemporary
-		compiler.readInstructionAddress(obj.ARET, addressReg)
+		compiler.compileReadInstructionAddress(obj.ARET, addressReg)
 
 		// Branch to the instruction after RET below via the absolute
 		// address stored in destinationRegister.
@@ -1655,7 +1650,7 @@ func TestArm64Compiler_readInstructionAddress(t *testing.T) {
 
 		// If we fail to branch, we reach here and exit with unreachable status,
 		// so the assertion would fail.
-		compiler.exit(jitCallStatusCodeUnreachable)
+		compiler.compileExitFromNativeCode(jitCallStatusCodeUnreachable)
 
 		// This could be the read instruction target as this is the
 		// right after RET. Therefore, the branch instruction above
@@ -1672,83 +1667,99 @@ func TestArm64Compiler_readInstructionAddress(t *testing.T) {
 	})
 }
 
-func TestArm64Compiler_compieleCall(t *testing.T) {
-	t.Run("need to grow call frame stack", func(t *testing.T) {
-		t.Skip("TODO")
-	})
-	t.Run("callframe stack ok", func(t *testing.T) {
-		env := newJITEnvironment()
-		engine := env.engine()
-		expectedValue := uint32(0)
+func TestArm64Compiler_compileCall(t *testing.T) {
+	for _, growCallFrameStack := range []bool{false, true} {
+		growCallFrameStack := growCallFrameStack
+		t.Run(fmt.Sprintf("grow=%v", growCallFrameStack), func(t *testing.T) {
+			env := newJITEnvironment()
+			engine := env.engine()
+			expectedValue := uint32(0)
 
-		// Emit the call target function.
-		const numCalls = 10
-		targetFunctionType := &wasm.FunctionType{
-			Params:  []wasm.ValueType{wasm.ValueTypeI32},
-			Results: []wasm.ValueType{wasm.ValueTypeI32},
-		}
-		for i := 0; i < numCalls; i++ {
-			// Each function takes one arguments, adds the value with 100 + i and returns the result.
-			addTargetValue := uint32(100 + i)
-			expectedValue += addTargetValue
+			if growCallFrameStack {
+				env.setCallFrameStackPointer(engine.globalContext.callFrameStackLen - 1)
+				env.setPreviousCallFrameStackPointer(engine.globalContext.callFrameStackLen - 1)
+			}
 
-			compiler := env.requireNewCompiler(t)
-			compiler.f = &wasm.FunctionInstance{FunctionType: &wasm.TypeInstance{Type: targetFunctionType},
-				ModuleInstance: &wasm.ModuleInstance{}}
+			// Emit the call target function.
+			const numCalls = 10
+			targetFunctionType := &wasm.FunctionType{
+				Params:  []wasm.ValueType{wasm.ValueTypeI32},
+				Results: []wasm.ValueType{wasm.ValueTypeI32},
+			}
+			for i := 0; i < numCalls; i++ {
+				// Each function takes one arguments, adds the value with 100 + i and returns the result.
+				addTargetValue := uint32(100 + i)
+				expectedValue += addTargetValue
+
+				compiler := env.requireNewCompiler(t)
+				compiler.f = &wasm.FunctionInstance{FunctionType: &wasm.TypeInstance{Type: targetFunctionType},
+					ModuleInstance: &wasm.ModuleInstance{}}
+
+				err := compiler.compilePreamble()
+				require.NoError(t, err)
+
+				err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(addTargetValue)})
+				require.NoError(t, err)
+				err = compiler.compileAdd(&wazeroir.OperationAdd{Type: wazeroir.UnsignedTypeI32})
+				require.NoError(t, err)
+				err = compiler.compileReturnFunction()
+				require.NoError(t, err)
 
+				code, _, _, err := compiler.compile()
+				require.NoError(t, err)
+				engine.addCompiledFunction(wasm.FunctionAddress(i), &compiledFunction{
+					codeSegment:        code,
+					codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
+				})
+			}
+
+			// Now we start building the caller's code.
+			compiler := env.requireNewCompiler(t)
 			err := compiler.compilePreamble()
 			require.NoError(t, err)
 
-			err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(addTargetValue)})
+			const initialValue = 100
+			expectedValue += initialValue
+			err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: 0}) // Dummy value so the base pointer would be non-trivial for callees.
 			require.NoError(t, err)
-			err = compiler.compileAdd(&wazeroir.OperationAdd{Type: wazeroir.UnsignedTypeI32})
+			err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: initialValue})
 			require.NoError(t, err)
+
+			// Call all the built functions.
+			for i := 0; i < numCalls; i++ {
+				err = compiler.compileCallImpl(wasm.FunctionAddress(i), targetFunctionType)
+				require.NoError(t, err)
+			}
+
 			err = compiler.compileReturnFunction()
 			require.NoError(t, err)
 
 			code, _, _, err := compiler.compile()
 			require.NoError(t, err)
-			engine.addCompiledFunction(wasm.FunctionAddress(i), &compiledFunction{
-				codeSegment:        code,
-				codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
-			})
-		}
-
-		// Now we start building the caller's code.
-		compiler := env.requireNewCompiler(t)
-		err := compiler.compilePreamble()
-		require.NoError(t, err)
-
-		const initialValue = 100
-		expectedValue += initialValue
-		err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: 0}) // Dummy value so the base pointer would be non-trivial for callees.
-		require.NoError(t, err)
-		err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: initialValue})
-		require.NoError(t, err)
-
-		// Call all the built functions.
-		for i := 0; i < numCalls; i++ {
-			err = compiler.compileCallFunction(wasm.FunctionAddress(i), targetFunctionType)
-			require.NoError(t, err)
-		}
 
-		err = compiler.compileReturnFunction()
-		require.NoError(t, err)
+			env.exec(code)
 
-		code, _, _, err := compiler.compile()
-		require.NoError(t, err)
+			if growCallFrameStack {
+				// If the call frame stack pointer equals the length of call frame stack length,
+				// we have to call the builtin function to grow the slice.
+				require.Equal(t, jitCallStatusCodeCallBuiltInFunction, env.jitStatus())
+				require.Equal(t, builtinFunctionAddressGrowCallFrameStack, env.functionCallAddress(), env.functionCallAddress())
 
-		env.exec(code)
+				// Grow the callFrame stack, and exec again from the return address.
+				env.engine().builtinFunctionGrowCallFrameStack()
+				jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(env.engine())))
+			}
 
-		// Check status and returned values.
-		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
-		require.Equal(t, uint64(2), env.stackPointer()) // Must be 2 (dummy value + the calculation results)
-		require.Equal(t, uint64(0), env.stackBasePointer())
-		require.Equal(t, expectedValue, env.stackTopAsUint32())
-	})
+			// Check status and returned values.
+			require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+			require.Equal(t, uint64(2), env.stackPointer()) // Must be 2 (dummy value + the calculation results)
+			require.Equal(t, uint64(0), env.stackBasePointer())
+			require.Equal(t, expectedValue, env.stackTopAsUint32())
+		})
+	}
 }
 
-func TestArm64Compiler_compieleSelect(t *testing.T) {
+func TestArm64Compiler_compileSelect(t *testing.T) {
 	for _, isFloat := range []bool{false, true} {
 		isFloat := isFloat
 		t.Run(fmt.Sprintf("float=%v", isFloat), func(t *testing.T) {
@@ -1812,7 +1823,7 @@ func TestArm64Compiler_compieleSelect(t *testing.T) {
 	}
 }
 
-func TestArm64Compiler_compieleSwap(t *testing.T) {
+func TestArm64Compiler_compileSwap(t *testing.T) {
 	const x, y uint64 = 100, 200
 	op := &wazeroir.OperationSwap{Depth: 10}
 
@@ -1901,13 +1912,14 @@ func TestAmd64Compiler_compileModuleContextInitialization(t *testing.T) {
 			compiler.f.ModuleInstance = tc.moduleInstance
 
 			// The assembler skips the first instruction so we intentionally add NOP here.
+			// TODO: delete after #233
 			compiler.compileNOP()
 
 			err := compiler.compileModuleContextInitialization()
 			require.NoError(t, err)
 			require.Empty(t, compiler.locationStack.usedRegisters)
 
-			compiler.exit(jitCallStatusCodeReturned)
+			compiler.compileExitFromNativeCode(jitCallStatusCodeReturned)
 
 			// Generate the code under test.
 			code, _, _, err := compiler.compile()
@@ -2424,13 +2436,11 @@ func TestArm64Compiler_compileLoad(t *testing.T) {
 			} else {
 				require.Equal(t, generalPurposeRegisterTypeInt, loadedLocation.registerType())
 			}
-
-			// Generate the code under test.
 			compiler.compileReturnFunction()
+
+			// Generate and run the code under test.
 			code, _, _, err := compiler.compile()
 			require.NoError(t, err)
-
-			// Run code.
 			env.exec(code)
 
 			// Verify the loaded value.
@@ -2439,3 +2449,164 @@ func TestArm64Compiler_compileLoad(t *testing.T) {
 		})
 	}
 }
+
+func TestArm64Compiler_compileMemoryGrow(t *testing.T) {
+	env := newJITEnvironment()
+	compiler := env.requireNewCompiler(t)
+	err := compiler.compilePreamble()
+	require.NoError(t, err)
+
+	err = compiler.compileMemoryGrow()
+	require.NoError(t, err)
+
+	// Emit arbitrary code after MemoryGrow returned so that we can verify
+	// that the code can set the return address properly.
+	const expValue uint32 = 100
+	err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: expValue})
+	require.NoError(t, err)
+	compiler.compileReturnFunction()
+	require.NoError(t, err)
+
+	// Generate and run the code under test.
+	code, _, _, err := compiler.compile()
+	require.NoError(t, err)
+	env.exec(code)
+
+	// After the initial exec, the code must exit with builtin function call status and funcaddress for memory grow.
+	require.Equal(t, jitCallStatusCodeCallBuiltInFunction, env.jitStatus())
+	require.Equal(t, builtinFunctionAddressMemoryGrow, env.functionCallAddress())
+
+	// Reenter from the return address.
+	jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(env.engine())))
+
+	// Check if the code successfully executed the code after builtin function call.
+	require.Equal(t, expValue, env.stackTopAsUint32())
+	require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+}
+
+func TestAmd64Compiler_compileMemorySize(t *testing.T) {
+	env := newJITEnvironment()
+	compiler := env.requireNewCompiler(t)
+	compiler.f.ModuleInstance = env.moduleInstance
+
+	err := compiler.compilePreamble()
+	require.NoError(t, err)
+
+	// Emit memory.size instructions.
+	err = compiler.compileMemorySize()
+	require.NoError(t, err)
+	// At this point, the size of memory should be pushed onto the stack.
+	require.Equal(t, uint64(1), compiler.locationStack.sp)
+	require.Equal(t, generalPurposeRegisterTypeInt, compiler.locationStack.peek().registerType())
+
+	compiler.compileReturnFunction()
+	require.NoError(t, err)
+
+	// Generate and run the code under test.
+	code, _, _, err := compiler.compile()
+	require.NoError(t, err)
+	env.exec(code)
+
+	require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	require.Equal(t, uint32(defaultMemoryPageNumInTest), env.stackTopAsUint32())
+}
+
+func TestAmd64Compiler_compileMaybeGrowValueStack(t *testing.T) {
+	t.Run("not grow", func(t *testing.T) {
+		const stackPointerCeil = 5
+		for _, baseOffset := range []uint64{5, 10, 20} {
+			t.Run(fmt.Sprintf("%d", baseOffset), func(t *testing.T) {
+				env := newJITEnvironment()
+				compiler := env.requireNewCompiler(t)
+
+				// The assembler skips the first instruction so we intentionally add NOP here.
+				// TODO: delete after #233
+				compiler.compileNOP()
+
+				err := compiler.compileMaybeGrowValueStack()
+				require.NoError(t, err)
+				require.NotNil(t, compiler.onStackPointerCeilDeterminedCallBack)
+
+				valueStackLen := uint64(len(env.stack()))
+				stackBasePointer := valueStackLen - baseOffset // Ceil <= valueStackLen - stackBasePointer = no need to grow!
+				compiler.onStackPointerCeilDeterminedCallBack(stackPointerCeil)
+				compiler.onStackPointerCeilDeterminedCallBack = nil
+				env.setValueStackBasePointer(stackBasePointer)
+
+				compiler.compileExitFromNativeCode(jitCallStatusCodeReturned)
+
+				// Generate and run the code under test.
+				code, _, _, err := compiler.compile()
+				require.NoError(t, err)
+				env.exec(code)
+
+				// The status code must be "Returned", not "BuiltinFunctionCall".
+				require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+			})
+		}
+	})
+	t.Run("grow", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+
+		// The assembler skips the first instruction so we intentionally add NOP here.
+		// TODO: delete after #233
+		compiler.compileNOP()
+
+		err := compiler.compileMaybeGrowValueStack()
+		require.NoError(t, err)
+
+		// On the return from grow value stack, we simply return.
+		compiler.compileReturnFunction()
+
+		stackPointerCeil := uint64(6)
+		compiler.stackPointerCeil = stackPointerCeil
+		valueStackLen := uint64(len(env.stack()))
+		stackBasePointer := valueStackLen - 5 // Ceil > valueStackLen - stackBasePointer = need to grow!
+		env.setValueStackBasePointer(stackBasePointer)
+
+		// Generate and run the code under test.
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+		env.exec(code)
+
+		// Check if the call exits with builtin function call status.
+		require.Equal(t, jitCallStatusCodeCallBuiltInFunction, env.jitStatus())
+
+		// Reenter from the return address.
+		returnAddress := env.callFrameStackPeek().returnAddress
+		require.NotZero(t, returnAddress)
+		jitcall(returnAddress, uintptr(unsafe.Pointer(env.engine())))
+
+		// Check the result. This should be "Returned".
+		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	})
+}
+
+func TestArm64Compiler_compileHostFunction(t *testing.T) {
+	env := newJITEnvironment()
+	compiler := env.requireNewCompiler(t)
+
+	// The assembler skips the first instruction so we intentionally add NOP here.
+	// TODO: delete after #233
+	compiler.compileNOP()
+
+	addr := wasm.FunctionAddress(100)
+	err := compiler.compileHostFunction(addr)
+	require.NoError(t, err)
+
+	// Generate and run the code under test.
+	code, _, _, err := compiler.compile()
+	require.NoError(t, err)
+	env.exec(code)
+
+	// On the return, the code must exit with the host call status and the specified call address.
+	require.Equal(t, jitCallStatusCodeCallHostFunction, env.jitStatus())
+	require.Equal(t, addr, env.functionCallAddress())
+
+	// Re-enter the return address.
+	jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(env.engine())))
+
+	// After that, the code must exit with returned status.
+	require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+}
diff --git a/wasm/mempage_test.go b/wasm/mempage_test.go
index 8fa2bfce..7063e7fa 100644
--- a/wasm/mempage_test.go
+++ b/wasm/mempage_test.go
@@ -7,7 +7,7 @@ import (
 )
 
 func TestMemoryPageConsts(t *testing.T) {
-	require.Equal(t, MemoryPageSize, uint32(1)<<memoryPageSizeInBit)
+	require.Equal(t, MemoryPageSize, uint32(1)<<MemoryPageSizeInBits)
 	require.Equal(t, MemoryPageSize, MemoryMaxPages)
 	require.Equal(t, MemoryPageSize, uint32(1<<16))
 }
@@ -24,7 +24,7 @@ func Test_MemoryBytesNumToPages(t *testing.T) {
 	}
 }
 
-func TestMemryInstance_Grow_Size(t *testing.T) {
+func TestMemoryInstance_Grow_Size(t *testing.T) {
 	t.Run("with max", func(t *testing.T) {
 		max := uint32(10)
 		m := &MemoryInstance{Max: &max, Buffer: make([]byte, 0)}