Repository: tetratelabs/wazero
Pull request #211 (2022-02-07 07:26:54)

Fix float32-type parameter not working for host function
This PR fixes that `float32` parameters were not correctly passed to the host function.

mathetake: I was kind of under the impression that this is tested in spectests, but I realized that all the host funcs used there is just a mocks :D https://github.com/tetratelabs/wazero/blob/main/tests/spectest/spec_test.go#L189-L197
nullpo-head: I've updated so that 1. it does `require.Equals` on the result of the callFunction, 2. test `flot64` as well. PTAL @codefromthecrypt . 
The latter change introduced a bit duplicated code, let me know if there's more sophisticated way to do this.
nullpo-head: You're right, it can be a table-driven one. I had wrongly assumed that it couldn't because I had lost track of the fact that `reflect.value` could be used and that the parameters and return value were always `uint64` after all. I've improved it. Thank you!

Patch:
diff --git a/wasm/interpreter/interpreter.go b/wasm/interpreter/interpreter.go
index 184b4085..3f196271 100644
--- a/wasm/interpreter/interpreter.go
+++ b/wasm/interpreter/interpreter.go
@@ -501,7 +501,9 @@ func (it *interpreter) callHostFunc(ctx context.Context, f *interpreterFunction)
 		raw := it.pop()
 		kind := tp.In(i).Kind()
 		switch kind {
-		case reflect.Float64, reflect.Float32:
+		case reflect.Float32:
+			val.SetFloat(float64(math.Float32frombits(uint32(raw))))
+		case reflect.Float64:
 			val.SetFloat(math.Float64frombits(raw))
 		case reflect.Uint32, reflect.Uint64:
 			val.SetUint(raw)
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index c7abbae8..0213cb8f 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -471,7 +471,9 @@ func (e *engine) execHostFunction(f *reflect.Value, ctx *wasm.HostFunctionCallCo
 		raw := e.popValue()
 		kind := tp.In(i).Kind()
 		switch kind {
-		case reflect.Float64, reflect.Float32:
+		case reflect.Float32:
+			val.SetFloat(float64(math.Float32frombits(uint32(raw))))
+		case reflect.Float64:
 			val.SetFloat(math.Float64frombits(raw))
 		case reflect.Uint32, reflect.Uint64:
 			val.SetUint(raw)
diff --git a/wasm/jit/jit_arm64.go b/wasm/jit/jit_arm64.go
index 1e3c195b..26193f6c 100644
--- a/wasm/jit/jit_arm64.go
+++ b/wasm/jit/jit_arm64.go
@@ -15,6 +15,7 @@ package jit
 import (
 	"errors"
 	"fmt"
+	"math"
 
 	asm "github.com/twitchyliquid64/golang-asm"
 	"github.com/twitchyliquid64/golang-asm/obj"
@@ -87,92 +88,95 @@ func (c *arm64Compiler) addInstruction(inst *obj.Prog) {
 	c.builder.AddInstruction(inst)
 }
 
+func (c *arm64Compiler) applyConstToRegisterInstruction(insturction obj.As, constValue int64, destinationRegister int16) {
+	applyConst := c.newProg()
+	applyConst.As = insturction
+	applyConst.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)
+}
+
+func (c *arm64Compiler) addLoadToRegisterInstruction(insturction obj.As, baseRegister int16, offset int64, destinationRegister int16) {
+	load := c.newProg()
+	load.As = insturction
+	load.From.Type = obj.TYPE_MEM
+	load.From.Reg = baseRegister
+	// 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.
+	load.From.Offset = offset
+	load.To.Type = obj.TYPE_REG
+	load.To.Reg = destinationRegister
+	c.addInstruction(load)
+}
+
+func (c *arm64Compiler) addStoreToMemoryInstruction(insturction obj.As, baseRegister int16, offset int64, source int16) {
+	store := c.newProg()
+	store.As = insturction
+	store.To.Type = obj.TYPE_MEM
+	store.To.Reg = baseRegister
+	// 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.
+	store.To.Offset = offset
+	store.From.Type = obj.TYPE_REG
+	store.From.Reg = source
+	c.addInstruction(store)
+}
+
 func (c *arm64Compiler) String() (ret string) { return }
 
+// emitPreamble implements compiler.emitPreamble for the arm64 architecture.
 func (c *arm64Compiler) emitPreamble() error {
 	// The assembler skips the first instruction so we intentionally add NOP here.
 	nop := c.newProg()
 	nop.As = obj.ANOP
 	c.addInstruction(nop)
+
+	// Before excuting function body, we must initialize the stack base pointer register
+	// so that we can manipulate the memory stack properly.
+	c.initializeReservedStackBasePointerRegister()
 	return nil
 }
 
 func (c *arm64Compiler) returnFunction() {
-
 	// TODO: we don't support function calls yet.
 	// For now the following code just simply returns to Go code.
 
-	// Since we return from the function, we need to decrement the callframe stack pointer.
+	// Since we return from the function, we need to decrement the callframe stack pointer, and write it back.
 	callFramePointerReg, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
-	loadCurrentCallFramPointer := c.newProg()
-	loadCurrentCallFramPointer.As = arm64.AMOVD
-	loadCurrentCallFramPointer.To.Type = obj.TYPE_REG
-	loadCurrentCallFramPointer.To.Reg = callFramePointerReg
-	loadCurrentCallFramPointer.From.Type = obj.TYPE_MEM
-	loadCurrentCallFramPointer.From.Reg = reservedRegisterForEngine
-	loadCurrentCallFramPointer.From.Offset = engineGlobalContextCallFrameStackPointerOffset
-	c.addInstruction(loadCurrentCallFramPointer)
-
-	decCallFrameStackPointer := c.newProg()
-	decCallFrameStackPointer.As = arm64.ASUBS
-	decCallFrameStackPointer.To.Type = obj.TYPE_REG
-	decCallFrameStackPointer.To.Reg = callFramePointerReg
-	decCallFrameStackPointer.From.Type = obj.TYPE_CONST
-	decCallFrameStackPointer.From.Offset = 1
-	c.addInstruction(decCallFrameStackPointer)
-
-	writeDecrementedCallFrameStackPoitner := c.newProg()
-	writeDecrementedCallFrameStackPoitner.As = arm64.AMOVD
-	writeDecrementedCallFrameStackPoitner.To = loadCurrentCallFramPointer.From
-	writeDecrementedCallFrameStackPoitner.From = loadCurrentCallFramPointer.To
-	c.addInstruction(writeDecrementedCallFrameStackPoitner)
+	c.addLoadToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg)
+	c.applyConstToRegisterInstruction(arm64.ASUBS, 1, callFramePointerReg)
+	c.addStoreToMemoryInstruction(arm64.AMOVD, reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg)
 
 	c.exit(jitCallStatusCodeReturned)
 }
 
+// exit adds instructions to give the control back to engine.exec with the given status code.
 func (c *arm64Compiler) exit(status jitCallStatusCode) {
 	tmp, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
 
+	// Write the current stack pointer to the engine.stackPointer.
+	c.applyConstToRegisterInstruction(arm64.AMOVW, int64(c.locationStack.sp), tmp)
+	c.addStoreToMemoryInstruction(arm64.AMOVW, reservedRegisterForEngine, engineValueStackContextStackPointerOffset, tmp)
+
 	if status != 0 {
-		loadStatusConst := c.newProg()
-		loadStatusConst.As = arm64.AMOVW
-		loadStatusConst.To.Type = obj.TYPE_REG
-		loadStatusConst.To.Reg = tmp
-		loadStatusConst.From.Type = obj.TYPE_CONST
-		loadStatusConst.From.Offset = int64(status)
-		c.addInstruction(loadStatusConst)
-
-		setJitStatus := c.newProg()
-		setJitStatus.As = arm64.AMOVWU
-		setJitStatus.From.Type = obj.TYPE_REG
-		setJitStatus.From.Reg = tmp
-		setJitStatus.To.Type = obj.TYPE_MEM
-		setJitStatus.To.Reg = reservedRegisterForEngine
-		setJitStatus.To.Offset = engineExitContextJITCallStatusCodeOffset
-		c.addInstruction(setJitStatus)
+		c.applyConstToRegisterInstruction(arm64.AMOVW, int64(status), tmp)
+		c.addStoreToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine, engineExitContextJITCallStatusCodeOffset, tmp)
 	} else {
 		// If the status == 0, we simply use zero register to store zero.
-		setJitStatus := c.newProg()
-		setJitStatus.As = arm64.AMOVWU
-		setJitStatus.From.Type = obj.TYPE_REG
-		setJitStatus.From.Reg = arm64.REGZERO
-		setJitStatus.To.Type = obj.TYPE_MEM
-		setJitStatus.To.Reg = reservedRegisterForEngine
-		setJitStatus.To.Offset = engineExitContextJITCallStatusCodeOffset
-		c.addInstruction(setJitStatus)
+		c.addStoreToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine, engineExitContextJITCallStatusCodeOffset, zeroRegister)
 	}
 
 	// The return address to the Go code is stored in archContext.jitReturnAddress which
 	// is embedded in engine. We load the value to the tmpRegister, and then
 	// invoke RET with that register.
-	loadReturnAddress := c.newProg()
-	loadReturnAddress.As = arm64.AMOVD
-	loadReturnAddress.To.Type = obj.TYPE_REG
-	loadReturnAddress.To.Reg = tmp
-	loadReturnAddress.From.Type = obj.TYPE_MEM
-	loadReturnAddress.From.Reg = reservedRegisterForEngine
-	loadReturnAddress.From.Offset = engineArchContextJITCallReturnAddressOffset
-	c.addInstruction(loadReturnAddress)
+	c.addLoadToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine, engineArchContextJITCallReturnAddressOffset, tmp)
 
 	ret := c.newProg()
 	ret.As = obj.ARET
@@ -451,18 +455,177 @@ func (c *arm64Compiler) compileMemorySize() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
+// compileConstI32 implements compiler.compileConstI32 for the arm64 architecture.
 func (c *arm64Compiler) compileConstI32(o *wazeroir.OperationConstI32) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	return c.emitIntConstant(true, uint64(o.Value))
 }
 
+// compileConstI64 implements compiler.compileConstI64 for the arm64 architecture.
 func (c *arm64Compiler) compileConstI64(o *wazeroir.OperationConstI64) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	return c.emitIntConstant(false, o.Value)
+}
+
+// emitIntConstant adds instructions to load an integer constant.
+// is32bit is true if the target value is originally 32-bit const, false otherwise.
+// value holds the (zero-extended for 32-bit case) load target constant.
+func (c *arm64Compiler) emitIntConstant(is32bit bool, value uint64) error {
+	if value == 0 {
+		c.pushZeroValue()
+	} else {
+		// Take a register to load the value.
+		reg, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+		if err != nil {
+			return err
+		}
+
+		var inst obj.As
+		if is32bit {
+			inst = arm64.AMOVW
+		} else {
+			inst = arm64.AMOVD
+		}
+		c.applyConstToRegisterInstruction(inst, int64(value), reg)
+
+		loc := c.locationStack.pushValueOnRegister(reg)
+		loc.setRegisterType(generalPurposeRegisterTypeInt)
+	}
+	return nil
 }
 
+// compileConstF32 implements compiler.compileConstF32 for the arm64 architecture.
 func (c *arm64Compiler) compileConstF32(o *wazeroir.OperationConstF32) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	return c.emitFloatConstant(true, uint64(math.Float32bits(o.Value)))
 }
 
+// compileConstF64 implements compiler.compileConstF64 for the arm64 architecture.
 func (c *arm64Compiler) compileConstF64(o *wazeroir.OperationConstF64) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	return c.emitFloatConstant(false, math.Float64bits(o.Value))
+}
+
+// emitFloatConstant adds instructions to load a float constant.
+// is32bit is true if the target value is originally 32-bit const, false otherwise.
+// value holds the (zero-extended for 32-bit case) bit representation of load target float constant.
+func (c *arm64Compiler) emitFloatConstant(is32bit bool, value uint64) error {
+	// Take a register to load the value.
+	reg, err := c.allocateRegister(generalPurposeRegisterTypeFloat)
+	if err != nil {
+		return err
+	}
+
+	tmpReg := zeroRegister
+	if value != 0 {
+		// If the target value is not zero, we have to load the constant
+		// temporarily into the integer register since we cannot directly
+		// move the const into float register.
+		tmpReg, err = c.allocateRegister(generalPurposeRegisterTypeInt)
+		if err != nil {
+			return err
+		}
+
+		var inst obj.As
+		if is32bit {
+			inst = arm64.AMOVW
+		} else {
+			inst = arm64.AMOVD
+		}
+		c.applyConstToRegisterInstruction(inst, int64(value), tmpReg)
+	}
+
+	// Use FMOV instruction to move the value on integer register into the float one.
+	mov := c.newProg()
+	if is32bit {
+		mov.As = arm64.AFMOVS
+	} else {
+		mov.As = arm64.AFMOVD
+	}
+	mov.From.Type = obj.TYPE_REG
+	mov.From.Reg = tmpReg
+	mov.To.Type = obj.TYPE_REG
+	mov.To.Reg = reg
+	c.addInstruction(mov)
+
+	loc := c.locationStack.pushValueOnRegister(reg)
+	loc.setRegisterType(generalPurposeRegisterTypeFloat)
+	return nil
+}
+
+func (c *arm64Compiler) pushZeroValue() {
+	c.locationStack.pushValueOnRegister(zeroRegister)
+}
+
+// allocateRegister returns an unused register of the given type. The register will be taken
+// either from the free register pool or by spilling an used register. If we allocate an used register,
+// this adds an instruction to write the value on a register back to memory stack region.
+// Note: resulting registers are NOT marked as used so the call site should mark it used if necessary.
+func (c *arm64Compiler) allocateRegister(t generalPurposeRegisterType) (reg int16, err error) {
+	var ok bool
+	// Try to get the unused register.
+	reg, ok = c.locationStack.takeFreeRegister(t)
+	if ok {
+		return
+	}
+
+	// If not found, we have to steal the register.
+	stealTarget, ok := c.locationStack.takeStealTargetFromUsedRegister(t)
+	if !ok {
+		err = fmt.Errorf("cannot steal register")
+		return
+	}
+
+	// Release the steal target register value onto stack location.
+	reg = stealTarget.register
+	c.releaseRegisterToStack(stealTarget)
+	return
+}
+
+// releaseRegisterToStack adds an instruction to write the value on a register back to memory stack region.
+func (c *arm64Compiler) releaseRegisterToStack(loc *valueLocation) {
+	var inst obj.As
+	switch loc.regType {
+	case generalPurposeRegisterTypeInt:
+		inst = arm64.AMOVD
+	case generalPurposeRegisterTypeFloat:
+		inst = arm64.AFMOVD
+	}
+
+	c.addStoreToMemoryInstruction(inst, reservedRegisterForStackBasePointerAddress, int64(loc.stackPointer)*8, loc.register)
+
+	// Mark the register is free.
+	c.locationStack.releaseRegister(loc)
+}
+
+// initializeReservedStackBasePointerRegister adds intructions to initialize reservedRegisterForStackBasePointerAddress
+// so that it points to the absolute address of the stack base for this function.
+func (c *arm64Compiler) initializeReservedStackBasePointerRegister() {
+	tmpReg, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
+
+	// First, load the address of the first element in the value stack into reservedRegisterForStackBasePointerAddress temporarily.
+	c.addLoadToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextValueStackElement0AddressOffset,
+		reservedRegisterForStackBasePointerAddress)
+
+	// Next we move the base pointer (engine.stackBasePointer) to the tmp register.
+	c.addLoadToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset,
+		tmpReg)
+
+	// Finally, we calculate "reservedRegisterForStackBasePointerAddress + tmpReg * 8"
+	// where we multiply tmpReg by 8 because stack pointer is an index in the []uint64
+	// so as an bytes we must multiply the size of uint64 = 8 bytes.
+	calcStackBasePointerAddress := c.newProg()
+	calcStackBasePointerAddress.As = arm64.AADD
+	calcStackBasePointerAddress.To.Type = obj.TYPE_REG
+	calcStackBasePointerAddress.To.Reg = reservedRegisterForStackBasePointerAddress
+	// We calculate "tmpReg * 8" as "tmpReg << 3".
+	setLeftShiftedRegister(calcStackBasePointerAddress, tmpReg, 3)
+	c.addInstruction(calcStackBasePointerAddress)
+}
+
+// setShiftedRegister modifies the given *obj.Prog so that .From (source operand)
+// becomes the "left shifted register". For example, this is used to emit instruction like
+// "add  x1, x2, x3, lsl #3" which means "x1 = x2 + (x3 << 3)".
+// See https://github.com/twitchyliquid64/golang-asm/blob/v0.15.1/obj/link.go#L120-L131
+func setLeftShiftedRegister(inst *obj.Prog, register int16, shiftNum int64) {
+	inst.From.Type = obj.TYPE_SHIFT
+	inst.From.Offset = (int64(register)&31)<<16 | 0<<22 | (shiftNum&63)<<10
 }
diff --git a/wasm/jit/jit_value_location_arm64.go b/wasm/jit/jit_value_location_arm64.go
index bd042727..4f983a91 100644
--- a/wasm/jit/jit_value_location_arm64.go
+++ b/wasm/jit/jit_value_location_arm64.go
@@ -15,14 +15,18 @@ const (
 	reservedRegisterForMemory = arm64.REG_R2
 )
 
+// zeroRegister is the alias of the arm64-specific zero register for readability.
+const zeroRegister int16 = arm64.REGZERO
+
 var (
 	generalPurposeFloatRegisters = []int16{
-		arm64.REG_V0, arm64.REG_V1, arm64.REG_V2, arm64.REG_V3, arm64.REG_V4, arm64.REG_V5,
-		arm64.REG_V6, arm64.REG_V7, arm64.REG_V8, arm64.REG_V9, arm64.REG_V10, arm64.REG_V11,
-		arm64.REG_V12, arm64.REG_V13, arm64.REG_V14, arm64.REG_V15, arm64.REG_V16, arm64.REG_V17,
-		arm64.REG_V18, arm64.REG_V19, arm64.REG_V20, arm64.REG_V21, arm64.REG_V22, arm64.REG_V23,
-		arm64.REG_V24, arm64.REG_V25, arm64.REG_V26, arm64.REG_V27, arm64.REG_V28, arm64.REG_V29,
-		arm64.REG_V30, arm64.REG_V31,
+		arm64.REG_F0, arm64.REG_F1, arm64.REG_F2, arm64.REG_F3,
+		arm64.REG_F4, arm64.REG_F5, arm64.REG_F6, arm64.REG_F7, arm64.REG_F8,
+		arm64.REG_F9, arm64.REG_F10, arm64.REG_F11, arm64.REG_F12, arm64.REG_F13,
+		arm64.REG_F14, arm64.REG_F15, arm64.REG_F16, arm64.REG_F17, arm64.REG_F18,
+		arm64.REG_F19, arm64.REG_F20, arm64.REG_F21, arm64.REG_F22, arm64.REG_F23,
+		arm64.REG_F24, arm64.REG_F25, arm64.REG_F26, arm64.REG_F27, arm64.REG_F28,
+		arm64.REG_F29, arm64.REG_F30, arm64.REG_F31,
 	}
 	unreservedGeneralPurposeIntRegisters = []int16{
 		arm64.REG_R3, arm64.REG_R4, arm64.REG_R5, arm64.REG_R6, arm64.REG_R7, arm64.REG_R8,
diff --git a/wasm/module.go b/wasm/module.go
index d08b42d6..6009dfe3 100644
--- a/wasm/module.go
+++ b/wasm/module.go
@@ -226,7 +226,10 @@ const (
 	ValueTypeF64 ValueType = 0x7c
 )
 
-func valueTypeName(t ValueType) string {
+// ValuTypeName returns the type name of the given ValueType as a string.
+// These type names match the names used in the WebAssembly text format.
+// Note that ValueTypeName returns "unknown", if an undefined ValueType value is passed.
+func ValueTypeName(t ValueType) string {
 	switch t {
 	case ValueTypeI32:
 		return "i32"
@@ -242,14 +245,14 @@ func valueTypeName(t ValueType) string {
 
 func (t *FunctionType) String() (ret string) {
 	for _, b := range t.Params {
-		ret += valueTypeName(b)
+		ret += ValueTypeName(b)
 	}
 	if len(t.Params) == 0 {
 		ret += "null"
 	}
 	ret += "_"
 	for _, b := range t.Results {
-		ret += valueTypeName(b)
+		ret += ValueTypeName(b)
 	}
 	if len(t.Results) == 0 {
 		ret += "null"


Test patch:
diff --git a/tests/engine/adhoc_test.go b/tests/engine/adhoc_test.go
index ed1253d3..d998ae79 100644
--- a/tests/engine/adhoc_test.go
+++ b/tests/engine/adhoc_test.go
@@ -1,12 +1,15 @@
 package adhoc
 
 import (
+	"bytes"
 	"context"
 	_ "embed"
+	"math"
 	"reflect"
 	"runtime"
 	"sync"
 	"testing"
+	"text/template"
 
 	"github.com/stretchr/testify/require"
 
@@ -60,6 +63,9 @@ func runTests(t *testing.T, newEngine func() wasm.Engine) {
 	t.Run("imported-and-exported func", func(t *testing.T) {
 		importedAndExportedFunc(t, newEngine)
 	})
+	t.Run("host function with float32 type", func(t *testing.T) {
+		hostFuncWithFloatParam(t, newEngine)
+	})
 }
 
 func fibonacci(t *testing.T, newEngine func() wasm.Engine) {
@@ -194,9 +200,9 @@ func recursiveEntry(t *testing.T, newEngine func() wasm.Engine) {
 	require.NoError(t, err)
 }
 
+// importedAndExportedFunc fails if the engine cannot call an "imported-and-then-exported-back" function
 func importedAndExportedFunc(t *testing.T, newEngine func() wasm.Engine) {
 	ctx := context.Background()
-	// Test that the engine can call "imported-and-then-exported-back" function correctly
 	mod, err := text.DecodeModule([]byte(`(module
 		;; arbitrary function with params
 		(import "env" "add_int" (func $add_int (param i32 i32) (result i32)))
@@ -221,3 +227,80 @@ func importedAndExportedFunc(t *testing.T, newEngine func() wasm.Engine) {
 	require.NoError(t, err)
 	require.Equal(t, uint64(42), results[0])
 }
+
+//  hostFuncWithFloatParam fails if a float parameter corrupts a host function value
+func hostFuncWithFloatParam(t *testing.T, newEngine func() wasm.Engine) {
+	ctx := context.Background()
+	watTemplate, err := template.New("").Parse(`(module
+		;; '{{ . }}' will be replaced by the type name, f32 or f64.
+
+		;; 'identity_float' just returns the given float value as is.
+		(import "test" "identity_float" (func $test.identity_float (param {{ . }}) (result {{ . }})))
+
+		;; 'call->test.identity_float' proxies 'test.identity_float' via call in order to test floats aren't corrupted 
+		;; when the float values are passed through OpCodeCall
+		(func $call->test.identity_float (param {{ . }}) (result {{ . }})
+			local.get 0
+			call $test.identity_float
+		)
+		(export "call->test.identity_float" (func $call->test.identity_float))
+		)`)
+	require.NoError(t, err)
+
+	tests := []struct {
+		testName          string
+		floatType         byte          // wasm.ValueTypeF32 or wasm.ValueTypeF64
+		identityFloatFunc reflect.Value // imported as 'identity_float' to the wasm module
+		floatParam        uint64        // passed to identityFloatFunc as a parameter
+		expectedFloatVal  float64
+	}{
+		{
+			testName:  "host function with f32 param",
+			floatType: wasm.ValueTypeF32,
+			identityFloatFunc: reflect.ValueOf(func(ctx *wasm.HostFunctionCallContext, value float32) float32 {
+				return value
+			}),
+			floatParam:       uint64(math.Float32bits(math.MaxFloat32)), // float bits as a uint32 value, but casted to uint64 to be passed to CallFunction
+			expectedFloatVal: float64(math.MaxFloat32),                  // arbitrary f32 value
+		},
+		{
+			testName:  "host function with f64 param",
+			floatType: wasm.ValueTypeF64,
+			identityFloatFunc: reflect.ValueOf(func(ctx *wasm.HostFunctionCallContext, value float64) float64 {
+				return value
+			}),
+			floatParam:       math.Float64bits(math.MaxFloat64),
+			expectedFloatVal: float64(math.MaxFloat64), // arbitrary f64 value that doesn' fit in f32
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+		t.Run(tc.testName, func(t *testing.T) {
+			// Build and instantiate the wat for `f32` or `f64`
+			var wat bytes.Buffer
+			err = watTemplate.Execute(&wat, wasm.ValueTypeName(tc.floatType))
+			require.NoError(t, err)
+
+			mod, err := text.DecodeModule(wat.Bytes())
+			require.NoError(t, err)
+
+			store := wasm.NewStore(newEngine())
+
+			err = store.AddHostFunction("test", "identity_float", tc.identityFloatFunc)
+			require.NoError(t, err)
+
+			err = store.Instantiate(mod, "mod")
+			require.NoError(t, err)
+
+			// This call should return the `expectedFloatVal`.
+			// That ensures that the float32 values are not corrupted when they are routed through OpCodeCall.
+			results, resultTypes, err := store.CallFunction(ctx, "mod", "call->test.identity_float", tc.floatParam)
+			require.NoError(t, err)
+			require.Len(t, results, len(resultTypes))
+			require.Equal(t, tc.floatType, resultTypes[0])
+			// Note that the both f32 and f64 result value are expressed as the float64 bits value in wazero.
+			require.Equal(t, tc.expectedFloatVal, math.Float64frombits(results[0]))
+		})
+	}
+}
diff --git a/wasm/jit/jit_arm64_test.go b/wasm/jit/jit_arm64_test.go
index 816f2c0d..c79e3037 100644
--- a/wasm/jit/jit_arm64_test.go
+++ b/wasm/jit/jit_arm64_test.go
@@ -5,12 +5,15 @@ package jit
 
 import (
 	"context"
+	"fmt"
+	"math"
 	"testing"
 	"unsafe"
 
 	"github.com/stretchr/testify/require"
 
 	"github.com/tetratelabs/wazero/wasm"
+	"github.com/tetratelabs/wazero/wasm/internal/wazeroir"
 )
 
 func (j *jitEnv) requireNewCompiler(t *testing.T) *arm64Compiler {
@@ -87,6 +90,9 @@ func TestArm64Compiler_exit(t *testing.T) {
 			// Build codes.
 			compiler := env.requireNewCompiler(t)
 			err := compiler.emitPreamble()
+
+			expStackPointer := uint64(100)
+			compiler.locationStack.sp = expStackPointer
 			require.NoError(t, err)
 			compiler.exit(s)
 
@@ -99,6 +105,82 @@ func TestArm64Compiler_exit(t *testing.T) {
 
 			// JIT status on engine must be updated.
 			require.Equal(t, s, env.jitStatus())
+
+			// Stack pointer must be written on engine.stackPointer on return.
+			require.Equal(t, expStackPointer, env.stackPointer())
+		})
+	}
+}
+
+func TestArm64Compiler_compileConsts(t *testing.T) {
+	for _, op := range []wazeroir.OperationKind{
+		wazeroir.OperationKindConstI32,
+		wazeroir.OperationKindConstI64,
+		wazeroir.OperationKindConstF32,
+		wazeroir.OperationKindConstF64,
+	} {
+		op := op
+		t.Run(op.String(), func(t *testing.T) {
+			for _, val := range []uint64{
+				0x0, 0x1, 0x1111000, 1 << 16, 1 << 21, 1 << 27, 1 << 32, 1<<32 + 1, 1 << 53,
+				math.Float64bits(math.Inf(1)),
+				math.Float64bits(math.Inf(-1)),
+				math.Float64bits(math.NaN()),
+				math.MaxUint32,
+				math.MaxInt32,
+				math.MaxUint64,
+				math.MaxInt64,
+				uint64(math.Float32bits(float32(math.Inf(1)))),
+				uint64(math.Float32bits(float32(math.Inf(-1)))),
+				uint64(math.Float32bits(float32(math.NaN()))),
+			} {
+				t.Run(fmt.Sprintf("0x%x", val), func(t *testing.T) {
+					env := newJITEnvironment()
+
+					// Build codes.
+					compiler := env.requireNewCompiler(t)
+					err := compiler.emitPreamble()
+					require.NoError(t, err)
+
+					switch op {
+					case wazeroir.OperationKindConstI32:
+						err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(val)})
+					case wazeroir.OperationKindConstI64:
+						err = compiler.compileConstI64(&wazeroir.OperationConstI64{Value: val})
+					case wazeroir.OperationKindConstF32:
+						err = compiler.compileConstF32(&wazeroir.OperationConstF32{Value: math.Float32frombits(uint32(val))})
+					case wazeroir.OperationKindConstF64:
+						err = compiler.compileConstF64(&wazeroir.OperationConstF64{Value: math.Float64frombits(val)})
+					}
+					require.NoError(t, err)
+
+					// After compiling const operations, we must see the register allocated value on the top of value.
+					loc := compiler.locationStack.peek()
+					require.True(t, loc.onRegister())
+
+					// Release the register allocated value to the memory stack so that we can see the value after exiting.
+					compiler.releaseRegisterToStack(loc)
+					compiler.returnFunction()
+
+					// Generate the code under test.
+					code, _, _, err := compiler.compile()
+					require.NoError(t, err)
+
+					// Run native code.
+					env.exec(code)
+
+					// JIT status on engine must be returned.
+					require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+					require.Equal(t, uint64(1), env.stackPointer())
+
+					switch op {
+					case wazeroir.OperationKindConstI32, wazeroir.OperationKindConstF32:
+						require.Equal(t, uint32(val), env.stackTopAsUint32())
+					case wazeroir.OperationKindConstI64, wazeroir.OperationKindConstF64:
+						require.Equal(t, val, env.stackTopAsUint64())
+					}
+				})
+			}
 		})
 	}
 }
diff --git a/wasm/module_test.go b/wasm/module_test.go
index 9cfd7215..7789a25e 100644
--- a/wasm/module_test.go
+++ b/wasm/module_test.go
@@ -99,7 +99,7 @@ func TestValueTypeName(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			require.Equal(t, tc.expected, valueTypeName(tc.input))
+			require.Equal(t, tc.expected, ValueTypeName(tc.input))
 		})
 	}
 }
diff --git a/wasm/store_test.go b/wasm/store_test.go
index 8b8150c5..2b59ce46 100644
--- a/wasm/store_test.go
+++ b/wasm/store_test.go
@@ -439,7 +439,7 @@ func TestStore_executeConstExpression(t *testing.T) {
 	})
 	t.Run("non global expr", func(t *testing.T) {
 		for _, vt := range []ValueType{ValueTypeI32, ValueTypeI64, ValueTypeF32, ValueTypeF64} {
-			t.Run(valueTypeName(vt), func(t *testing.T) {
+			t.Run(ValueTypeName(vt), func(t *testing.T) {
 				t.Run("valid", func(t *testing.T) {
 					// Allocate bytes with enough size for all types.
 					expr := &ConstantExpression{Data: make([]byte, 8)}
@@ -526,7 +526,7 @@ func TestStore_executeConstExpression(t *testing.T) {
 				{valueType: ValueTypeF32, val: uint64(math.Float32bits(634634432.12311))},
 				{valueType: ValueTypeF64, val: math.Float64bits(1.12312311)},
 			} {
-				t.Run(valueTypeName(tc.valueType), func(t *testing.T) {
+				t.Run(ValueTypeName(tc.valueType), func(t *testing.T) {
 					// The index specified in Data equals zero.
 					expr := &ConstantExpression{Data: []byte{0}, Opcode: OpcodeGlobalGet}
 					globals := []*GlobalInstance{{Val: tc.val, Type: &GlobalType{ValType: tc.valueType}}}