Repository: tetratelabs/wazero
Pull request #201 (2022-02-04 07:21:51)

jit(arm64): support for function return
This PR adds support for br instruction targeting outermost label
which means return from the function. For example, the following
function bodies are supported now:

* {End}
* {Br, 0x0, End}


part of #187 


Patch:
diff --git a/.github/workflows/commit.yaml b/.github/workflows/commit.yaml
index 57e0a1c2..9e12502a 100644
--- a/.github/workflows/commit.yaml
+++ b/.github/workflows/commit.yaml
@@ -30,12 +30,10 @@ jobs:
           path: |
             ~/.cache/go-build
             ~/go/pkg/mod
-          key: check-${{ runner.os }}-go-${{ env.GO_VERSION }}-${{ hashFiles('**/go.sum') }}
+          key: check-${{ runner.os }}-go-${{ env.GO_VERSION }}-${{ hashFiles('**/go.sum', 'Makefile') }}
 
       - run: make lint
 
-      - run: make build.lib
-
       - run: make build.spectest
 
       - run: make check
@@ -44,6 +42,7 @@ jobs:
     name: amd64, ${{ matrix.os }}, Go-${{ matrix.go-version }}
     runs-on: ${{ matrix.os }}
     strategy:
+      fail-fast: false  # don't fail fast as sometimes failures are arch/OS specific
       matrix:
         os: 
         - "ubuntu-latest"
@@ -76,6 +75,7 @@ jobs:
     name: ${{ matrix.target.arch }}, ubuntu, Go-${{ matrix.go-version }}
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false  # don't fail fast as sometimes failures are arch/OS specific
       matrix:
         go-version:
         - "1.17"  # == ${{ env.GO_VERSION }} because matrix cannot expand env variables
diff --git a/Makefile b/Makefile
index f2a49bbd..5e201d80 100644
--- a/Makefile
+++ b/Makefile
@@ -5,16 +5,6 @@ golangci_lint := github.com/golangci/golangci-lint/cmd/golangci-lint@v1.42.0
 bench:
 	@go test -run=NONE -benchmem -bench=. ./tests/...
 
-.PHONY: build.lib
-build.lib:
-	@echo "Ensuring that the wazero library can be compiled on primary platforms..."
-	@GOOS=linux GOARCH=arm64 go build ./...
-	@GOOS=linux GOARCH=amd64 go build ./...
-	@GOOS=darwin GOARCH=arm64 go build ./...
-	@GOOS=darwin GOARCH=amd64 go build ./...
-	@GOOS=windows GOARCH=arm64 go build ./...
-	@GOOS=windows GOARCH=amd64 go build ./...
-
 bench_testdata_dir := tests/bench/testdata
 
 .PHONY: build.bench
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index 09e31f88..84d9da04 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -25,6 +25,7 @@ type (
 		moduleContext
 		valueStackContext
 		exitContext
+		archContext
 
 		// The following fields are NOT acceessed by JITed code directly.
 
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 59fa5cda..6403056a 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -82,6 +82,10 @@ func init() {
 // engine is the pointer to the "*engine" as uintptr.
 func jitcall(codeSegment, engine uintptr)
 
+// archContext is embedded in Engine in order to store architecture-specific data.
+// For amd64, this is empty.
+type archContext struct{}
+
 // newCompiler returns a new compiler interface which can be used to compile the given function instance.
 // Note: ir param can be nil for host functions.
 func newCompiler(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (compiler, error) {
diff --git a/wasm/jit/jit_arm64.go b/wasm/jit/jit_arm64.go
index 1e427625..3dd7cab3 100644
--- a/wasm/jit/jit_arm64.go
+++ b/wasm/jit/jit_arm64.go
@@ -1,15 +1,41 @@
 //go:build arm64
 // +build arm64
 
+// This file implements the compiler for arm64 target.
+// Please refer to https://developer.arm.com/documentation/102374/latest/
+// if unfamiliar with arm64 instructions and semantics.
+//
+// Note: we use arm64 pkg as the assembler (github.com/twitchyliquid64/golang-asm/obj/arm64)
+// which has different notation from the original arm64 assembly. For example,
+// 64-bit variant ldr, str, stur are all corresponding to arm64.AMOVD.
+// Please refer to https://pkg.go.dev/cmd/internal/obj/arm64.
+
 package jit
 
 import (
+	"errors"
 	"fmt"
 
+	asm "github.com/twitchyliquid64/golang-asm"
+	"github.com/twitchyliquid64/golang-asm/obj"
+	"github.com/twitchyliquid64/golang-asm/obj/arm64"
+
 	"github.com/tetratelabs/wazero/wasm"
 	"github.com/tetratelabs/wazero/wasm/internal/wazeroir"
 )
 
+// archContext is embedded in Engine in order to store architecture-specific data.
+type archContext struct {
+	// jitCallReturnAddress holds the absolute return address for jitcall.
+	// The value is set whenever jitcall is executed and done in jit_arm64.s
+	// Native code can return back to the engine.exec's main loop back by
+	// executing "ret" instruction with this value. See arm64Compiler.exit.
+	jitCallReturnAddress uint64
+}
+
+// engineArchContextJITCallReturnAddressOffset is the offset of archContext.jitCallReturnAddress
+const engineArchContextJITCallReturnAddressOffset = 136
+
 // jitcall is implemented in jit_arm64.s as a Go Assembler function.
 // This is used by engine.exec and the entrypoint to enter the JITed native code.
 // codeSegment is the pointer to the initial instruction of the compiled native code.
@@ -19,297 +45,423 @@ func jitcall(codeSegment, engine uintptr)
 // newCompiler returns a new compiler interface which can be used to compile the given function instance.
 // Note: ir param can be nil for host functions.
 func newCompiler(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (compiler, error) {
-	return &arm64Copmiler{}, nil
-}
-
-type arm64Copmiler struct{}
-
-func (c *arm64Copmiler) String() (ret string) { return }
-
-func (c *arm64Copmiler) emitPreamble() error {
-	return fmt.Errorf("TODO: unsupported on arm64")
-}
-
-func (c *arm64Copmiler) generate() (code []byte, staticData compiledFunctionStaticData, maxStackPointer uint64, err error) {
-	return nil, nil, 0, fmt.Errorf("TODO: unsupported on arm64")
-}
-
-func (c *arm64Copmiler) compileHostFunction(address wasm.FunctionAddress) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
-}
-
-func (c *arm64Copmiler) compileLabel(o *wazeroir.OperationLabel) (skipThisLabel bool) {
+	// We can choose arbitrary number instead of 1024 which indicates the cache size in the compiler.
+	// TODO: optimize the number.
+	b, err := asm.NewBuilder("arm64", 1024)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create a new assembly builder: %w", err)
+	}
+
+	compiler := &arm64Compiler{
+		f:             f,
+		builder:       b,
+		locationStack: newValueLocationStack(),
+	}
+
+	return compiler, nil
+}
+
+type arm64Compiler struct {
+	builder *asm.Builder
+	f       *wasm.FunctionInstance
+	// locationStack holds the state of wazeroir virtual stack.
+	// and each item is either placed in register or the actual memory stack.
+	locationStack *valueLocationStack
+}
+
+func (c *arm64Compiler) generate() (code []byte, staticData compiledFunctionStaticData, maxStackPointer uint64, err error) {
+	code, err = mmapCodeSegment(c.builder.Assemble())
+	if err != nil {
+		return
+	}
+	return
+}
+
+func (c *arm64Compiler) newProg() (inst *obj.Prog) {
+	inst = c.builder.NewProg()
+	return
+}
+
+func (c *arm64Compiler) addInstruction(inst *obj.Prog) {
+	c.builder.AddInstruction(inst)
+}
+
+func (c *arm64Compiler) String() (ret string) { return }
+
+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)
+	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.
+	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.exit(jitCallStatusCodeReturned)
+	return
+}
+
+func (c *arm64Compiler) exit(status jitCallStatusCode) {
+	tmp, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
+
+	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)
+	} 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)
+	}
+
+	// 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)
+
+	ret := c.newProg()
+	ret.As = obj.ARET
+	ret.To.Type = obj.TYPE_REG
+	ret.To.Reg = tmp
+	c.addInstruction(ret)
+}
+
+func (c *arm64Compiler) compileHostFunction(address wasm.FunctionAddress) error {
+	return errors.New("TODO: implement compileHostFunction on arm64")
+}
+
+func (c *arm64Compiler) compileLabel(o *wazeroir.OperationLabel) (skipThisLabel bool) {
 	return false
 }
 
-func (c *arm64Copmiler) compileUnreachable() error {
+func (c *arm64Compiler) compileUnreachable() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileSwap(o *wazeroir.OperationSwap) error {
+func (c *arm64Compiler) compileSwap(o *wazeroir.OperationSwap) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileGlobalGet(o *wazeroir.OperationGlobalGet) error {
+func (c *arm64Compiler) compileGlobalGet(o *wazeroir.OperationGlobalGet) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileGlobalSet(o *wazeroir.OperationGlobalSet) error {
+func (c *arm64Compiler) compileGlobalSet(o *wazeroir.OperationGlobalSet) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileBr(o *wazeroir.OperationBr) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+func (c *arm64Compiler) compileBr(o *wazeroir.OperationBr) error {
+	if o.Target.IsReturnTarget() {
+		c.returnFunction()
+		return nil
+	} else {
+		return fmt.Errorf("TODO: only return target is available on arm64")
+	}
 }
 
-func (c *arm64Copmiler) compileBrIf(o *wazeroir.OperationBrIf) error {
+func (c *arm64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileBrTable(o *wazeroir.OperationBrTable) error {
+func (c *arm64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileCall(o *wazeroir.OperationCall) error {
+func (c *arm64Compiler) compileCall(o *wazeroir.OperationCall) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileCallIndirect(o *wazeroir.OperationCallIndirect) error {
+func (c *arm64Compiler) compileCallIndirect(o *wazeroir.OperationCallIndirect) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileDrop(o *wazeroir.OperationDrop) error {
+func (c *arm64Compiler) compileDrop(o *wazeroir.OperationDrop) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileSelect() error {
+func (c *arm64Compiler) compileSelect() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compilePick(o *wazeroir.OperationPick) error {
+func (c *arm64Compiler) compilePick(o *wazeroir.OperationPick) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileAdd(o *wazeroir.OperationAdd) error {
+func (c *arm64Compiler) compileAdd(o *wazeroir.OperationAdd) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileSub(o *wazeroir.OperationSub) error {
+func (c *arm64Compiler) compileSub(o *wazeroir.OperationSub) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileMul(o *wazeroir.OperationMul) error {
+func (c *arm64Compiler) compileMul(o *wazeroir.OperationMul) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileClz(o *wazeroir.OperationClz) error {
+func (c *arm64Compiler) compileClz(o *wazeroir.OperationClz) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileCtz(o *wazeroir.OperationCtz) error {
+func (c *arm64Compiler) compileCtz(o *wazeroir.OperationCtz) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compilePopcnt(o *wazeroir.OperationPopcnt) error {
+func (c *arm64Compiler) compilePopcnt(o *wazeroir.OperationPopcnt) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileDiv(o *wazeroir.OperationDiv) error {
+func (c *arm64Compiler) compileDiv(o *wazeroir.OperationDiv) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileRem(o *wazeroir.OperationRem) error {
+func (c *arm64Compiler) compileRem(o *wazeroir.OperationRem) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileAnd(o *wazeroir.OperationAnd) error {
+func (c *arm64Compiler) compileAnd(o *wazeroir.OperationAnd) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileOr(o *wazeroir.OperationOr) error {
+func (c *arm64Compiler) compileOr(o *wazeroir.OperationOr) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileXor(o *wazeroir.OperationXor) error {
+func (c *arm64Compiler) compileXor(o *wazeroir.OperationXor) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileShl(o *wazeroir.OperationShl) error {
+func (c *arm64Compiler) compileShl(o *wazeroir.OperationShl) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileShr(o *wazeroir.OperationShr) error {
+func (c *arm64Compiler) compileShr(o *wazeroir.OperationShr) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileRotl(o *wazeroir.OperationRotl) error {
+func (c *arm64Compiler) compileRotl(o *wazeroir.OperationRotl) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileRotr(o *wazeroir.OperationRotr) error {
+func (c *arm64Compiler) compileRotr(o *wazeroir.OperationRotr) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileAbs(o *wazeroir.OperationAbs) error {
+func (c *arm64Compiler) compileAbs(o *wazeroir.OperationAbs) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileNeg(o *wazeroir.OperationNeg) error {
+func (c *arm64Compiler) compileNeg(o *wazeroir.OperationNeg) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileCeil(o *wazeroir.OperationCeil) error {
+func (c *arm64Compiler) compileCeil(o *wazeroir.OperationCeil) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileFloor(o *wazeroir.OperationFloor) error {
+func (c *arm64Compiler) compileFloor(o *wazeroir.OperationFloor) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileTrunc(o *wazeroir.OperationTrunc) error {
+func (c *arm64Compiler) compileTrunc(o *wazeroir.OperationTrunc) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileNearest(o *wazeroir.OperationNearest) error {
+func (c *arm64Compiler) compileNearest(o *wazeroir.OperationNearest) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileSqrt(o *wazeroir.OperationSqrt) error {
+func (c *arm64Compiler) compileSqrt(o *wazeroir.OperationSqrt) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileMin(o *wazeroir.OperationMin) error {
+func (c *arm64Compiler) compileMin(o *wazeroir.OperationMin) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileMax(o *wazeroir.OperationMax) error {
+func (c *arm64Compiler) compileMax(o *wazeroir.OperationMax) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileCopysign(o *wazeroir.OperationCopysign) error {
+func (c *arm64Compiler) compileCopysign(o *wazeroir.OperationCopysign) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileI32WrapFromI64() error {
+func (c *arm64Compiler) compileI32WrapFromI64() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileITruncFromF(o *wazeroir.OperationITruncFromF) error {
+func (c *arm64Compiler) compileITruncFromF(o *wazeroir.OperationITruncFromF) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileFConvertFromI(o *wazeroir.OperationFConvertFromI) error {
+func (c *arm64Compiler) compileFConvertFromI(o *wazeroir.OperationFConvertFromI) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileF32DemoteFromF64() error {
+func (c *arm64Compiler) compileF32DemoteFromF64() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileF64PromoteFromF32() error {
+func (c *arm64Compiler) compileF64PromoteFromF32() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileI32ReinterpretFromF32() error {
+func (c *arm64Compiler) compileI32ReinterpretFromF32() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileI64ReinterpretFromF64() error {
+func (c *arm64Compiler) compileI64ReinterpretFromF64() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileF32ReinterpretFromI32() error {
+func (c *arm64Compiler) compileF32ReinterpretFromI32() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileF64ReinterpretFromI64() error {
+func (c *arm64Compiler) compileF64ReinterpretFromI64() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileExtend(o *wazeroir.OperationExtend) error {
+func (c *arm64Compiler) compileExtend(o *wazeroir.OperationExtend) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileEq(o *wazeroir.OperationEq) error {
+func (c *arm64Compiler) compileEq(o *wazeroir.OperationEq) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileNe(o *wazeroir.OperationNe) error {
+func (c *arm64Compiler) compileNe(o *wazeroir.OperationNe) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileEqz(o *wazeroir.OperationEqz) error {
+func (c *arm64Compiler) compileEqz(o *wazeroir.OperationEqz) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileLt(o *wazeroir.OperationLt) error {
+func (c *arm64Compiler) compileLt(o *wazeroir.OperationLt) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileGt(o *wazeroir.OperationGt) error {
+func (c *arm64Compiler) compileGt(o *wazeroir.OperationGt) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileLe(o *wazeroir.OperationLe) error {
+func (c *arm64Compiler) compileLe(o *wazeroir.OperationLe) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileGe(o *wazeroir.OperationGe) error {
+func (c *arm64Compiler) compileGe(o *wazeroir.OperationGe) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileLoad(o *wazeroir.OperationLoad) error {
+func (c *arm64Compiler) compileLoad(o *wazeroir.OperationLoad) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileLoad8(o *wazeroir.OperationLoad8) error {
+func (c *arm64Compiler) compileLoad8(o *wazeroir.OperationLoad8) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileLoad16(o *wazeroir.OperationLoad16) error {
+func (c *arm64Compiler) compileLoad16(o *wazeroir.OperationLoad16) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileLoad32(o *wazeroir.OperationLoad32) error {
+func (c *arm64Compiler) compileLoad32(o *wazeroir.OperationLoad32) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileStore(o *wazeroir.OperationStore) error {
+func (c *arm64Compiler) compileStore(o *wazeroir.OperationStore) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileStore8(o *wazeroir.OperationStore8) error {
+func (c *arm64Compiler) compileStore8(o *wazeroir.OperationStore8) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileStore16(o *wazeroir.OperationStore16) error {
+func (c *arm64Compiler) compileStore16(o *wazeroir.OperationStore16) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileStore32(o *wazeroir.OperationStore32) error {
+func (c *arm64Compiler) compileStore32(o *wazeroir.OperationStore32) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileMemoryGrow() error {
+func (c *arm64Compiler) compileMemoryGrow() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileMemorySize() error {
+func (c *arm64Compiler) compileMemorySize() error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileConstI32(o *wazeroir.OperationConstI32) error {
+func (c *arm64Compiler) compileConstI32(o *wazeroir.OperationConstI32) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileConstI64(o *wazeroir.OperationConstI64) error {
+func (c *arm64Compiler) compileConstI64(o *wazeroir.OperationConstI64) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileConstF32(o *wazeroir.OperationConstF32) error {
+func (c *arm64Compiler) compileConstF32(o *wazeroir.OperationConstF32) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
-func (c *arm64Copmiler) compileConstF64(o *wazeroir.OperationConstF64) error {
+func (c *arm64Compiler) compileConstF64(o *wazeroir.OperationConstF64) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
diff --git a/wasm/jit/jit_arm64.s b/wasm/jit/jit_arm64.s
index 4919f42d..82713bf5 100644
--- a/wasm/jit/jit_arm64.s
+++ b/wasm/jit/jit_arm64.s
@@ -6,6 +6,13 @@
 
 // jitcall(codeSegment, engine)
 TEXT ·jitcall(SB),NOSPLIT|NOFRAME,$0-16
-        MOVD codeSegment+0(FP),R1  // Load the address of native code.
-        MOVD engine+8(FP),R0       // Load the address of engine.
-        JMP (R1)                   // Jump to native code.
+        // Load the address of native code.
+        MOVD codeSegment+0(FP),R1
+        // Load the address of engine.
+        MOVD engine+8(FP),R0
+        // In arm64, return address is stored in R30 after jumping into the code.
+        // We save the return address value into archContext.jitReturnAddress in Engine.
+        // Note that the const 136 drifts after editting Engine or archContext struct. See TestArchContextOffsetInEngine.
+        MOVD R30,136(R0)          
+        // Jump to native code.
+        JMP (R1)
diff --git a/wasm/jit/jit_other.go b/wasm/jit/jit_other.go
index 4c03a117..a1a7d6c6 100644
--- a/wasm/jit/jit_other.go
+++ b/wasm/jit/jit_other.go
@@ -8,6 +8,8 @@ import (
 	"github.com/tetratelabs/wazero/wasm/internal/wazeroir"
 )
 
+type archContext struct{}
+
 func jitcall(codeSegment, engine uintptr) {
 	panic("unsupported GOARCH")
 }
diff --git a/wasm/jit/mmap.go b/wasm/jit/mmap.go
index 423ef5cc..202ece80 100644
--- a/wasm/jit/mmap.go
+++ b/wasm/jit/mmap.go
@@ -3,11 +3,24 @@
 
 package jit
 
-import "syscall"
+import (
+	"runtime"
+	"syscall"
+)
 
 // mmapCodeSegment copies the code into the executable region and returns the byte slice of the region.
 // See https://man7.org/linux/man-pages/man2/mmap.2.html for mmap API and flags.
 func mmapCodeSegment(code []byte) ([]byte, error) {
+	if runtime.GOARCH == "amd64" {
+		return mmapCodeSegmentAMD64(code)
+	} else {
+		return mmapCodeSegmentARM64(code)
+	}
+}
+
+// mmapCodeSegmentAMD64 gives all read-write-exec permission to the mmap region
+// to enter the function. Otherwise, segmentation fault exeception is raied.
+func mmapCodeSegmentAMD64(code []byte) ([]byte, error) {
 	mmapFunc, err := syscall.Mmap(
 		-1,
 		0,
@@ -24,3 +37,29 @@ func mmapCodeSegment(code []byte) ([]byte, error) {
 	copy(mmapFunc, code)
 	return mmapFunc, nil
 }
+
+// mmapCodeSegmentARM64 cannot give all read-write-exec permission to the mmap region.
+// Otherwise, the mmap systemcall would raise an error. Here we give read-write
+// to the region at first, write the native code and then change the perm to
+// read-exec so we can execute the native code.
+func mmapCodeSegmentARM64(code []byte) ([]byte, error) {
+	mmapFunc, err := syscall.Mmap(
+		-1,
+		0,
+		len(code),
+		// The region must be RW: RW for writing native codes.
+		syscall.PROT_READ|syscall.PROT_WRITE,
+		// Anonymous as this is not an actual file, but a memory,
+		// Private as this is in-process memory region.
+		syscall.MAP_ANON|syscall.MAP_PRIVATE,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	copy(mmapFunc, code)
+
+	// Then we're done with writing code, change the permission to RX.
+	err = syscall.Mprotect(mmapFunc, syscall.PROT_READ|syscall.PROT_EXEC)
+	return mmapFunc, err
+}
diff --git a/wasm/text/decoder.go b/wasm/text/decoder.go
index 88c01c34..521f0fff 100644
--- a/wasm/text/decoder.go
+++ b/wasm/text/decoder.go
@@ -121,7 +121,7 @@ func newModuleParser(module *wasm.Module) *moduleParser {
 	}
 	p.typeParser = newTypeParser(p.typeNamespace, p.onTypeEnd)
 	p.typeUseParser = newTypeUseParser(module, p.typeNamespace)
-	p.funcParser = newFuncParser(p.funcNamespace, p.endFunc)
+	p.funcParser = newFuncParser(p.typeUseParser, p.funcNamespace, p.endFunc)
 	p.memoryParser = newMemoryParser(p.memoryNamespace, p.endMemory)
 	return &p
 }
@@ -160,7 +160,8 @@ func (p *moduleParser) beginModuleField(tok tokenType, tokenBytes []byte, _, _ u
 			return p.parseImportModule, nil
 		case "func":
 			p.pos = positionModuleFunc
-			return p.parseFuncID, nil
+			p.funcParser.currentIdx = wasm.Index(len(p.module.FunctionSection))
+			return p.funcParser.begin, nil
 		case "memory":
 			p.pos = positionModuleMemory
 			return p.memoryParser.begin, nil
@@ -302,14 +303,25 @@ func (p *moduleParser) beginImportDesc(tok tokenType, tokenBytes []byte, _, _ ui
 //                  calls parseImportFunc here --^
 func (p *moduleParser) parseImportFuncID(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 	if tok == tokenID { // Ex. $main
-		if err := p.setFuncID(tokenBytes); err != nil {
+		id, err := p.funcNamespace.setID(tokenBytes)
+		if err != nil {
 			return nil, err
 		}
+		p.addFunctionName(id)
 		return p.parseImportFunc, nil
 	}
 	return p.parseImportFunc(tok, tokenBytes, line, col)
 }
 
+// addFunctionName appends the current imported or module-defined function name to the wasm.NameSection
+func (p *moduleParser) addFunctionName(name string) {
+	if name == "" {
+		return // there's no value in an empty name
+	}
+	na := &wasm.NameAssoc{Index: p.funcNamespace.count, Name: name}
+	p.module.NameSection.FunctionNames = append(p.module.NameSection.FunctionNames, na)
+}
+
 // parseImportFunc passes control to the typeUseParser until any signature is read, then returns onImportFunc.
 //
 // Ex. `(import "Math" "PI" (func $math.pi (result f32)))`
@@ -330,7 +342,7 @@ func (p *moduleParser) parseImportFunc(tok tokenType, tokenBytes []byte, line, c
 // onImportFunc records the type index and local names of the current imported function, and increments
 // funcNamespace as it is shared across imported and module-defined functions. Finally, this returns parseImportEnd to
 // the current import into the ImportSection.
-func (p *moduleParser) onImportFunc(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+func (p *moduleParser) onImportFunc(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, _, _ uint32) (tokenParser, error) {
 	i := p.currentModuleField.(*wasm.Import)
 	i.Kind = wasm.ImportKindFunc
 	i.DescFunc = typeIdx
@@ -380,61 +392,10 @@ func (p *moduleParser) parseImportEnd(tok tokenType, tokenBytes []byte, _, _ uin
 	return p.parseModule, nil
 }
 
-// parseFuncID records the ID of the current function, if present, and resumes with parseFunc.
-//
-// Ex. A function ID is present `(module (func $math.pi (result f32))`
-//                      records math.pi here --^
-//                             parseFunc resumes here --^
-//
-// Ex. No function ID `(module (func (result f32))`
-//            calls parseFunc here --^
-func (p *moduleParser) parseFuncID(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	if tok == tokenID { // Ex. $main
-		if err := p.setFuncID(tokenBytes); err != nil {
-			return nil, err
-		}
-		return p.parseFunc, nil
-	}
-	return p.parseFunc(tok, tokenBytes, line, col)
-}
-
-// setFuncID adds the normalized ('$' stripped) function ID to the funcNamespace and the wasm.NameSection.
-func (p *moduleParser) setFuncID(idToken []byte) error {
-	id, err := p.funcNamespace.setID(idToken)
-	if err != nil {
-		return err
-	}
-	na := &wasm.NameAssoc{Index: p.funcNamespace.count, Name: id}
-	p.module.NameSection.FunctionNames = append(p.module.NameSection.FunctionNames, na)
-	return nil
-}
-
-// parseFunc passes control to the typeUseParser until any signature is read, then funcParser until and locals or body
-// are read. Finally, this finishes via endFunc.
-//
-// Ex. `(module (func $math.pi (result f32))`
-//               starts here --^           ^
-//                  endFunc resumes here --+
-//
-// Ex.    `(module (func $math.pi (result f32) (local i32) )`
-//                  starts here --^            ^           ^
-//             funcParser.begin resumes here --+           |
-//                                  endFunc resumes here --+
-//
-// Ex. If there is no signature `(func)`
-//              calls endFunc here ---^
-func (p *moduleParser) parseFunc(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	idx := wasm.Index(len(p.module.FunctionSection))
-	if tok == tokenID { // Ex. (func $main $main)
-		return nil, fmt.Errorf("redundant ID %s", tokenBytes)
-	}
-
-	return p.typeUseParser.begin(wasm.SectionIDFunction, idx, p.funcParser.begin, tok, tokenBytes, line, col)
-}
-
 // endFunc adds the type index, code and local names for the current function, and increments funcNamespace as it is
 // shared across imported and module-defined functions. Finally, this returns parseModule to prepare for the next field.
-func (p *moduleParser) endFunc(typeIdx wasm.Index, code *wasm.Code, localNames wasm.NameMap) (tokenParser, error) {
+func (p *moduleParser) endFunc(typeIdx wasm.Index, code *wasm.Code, name string, localNames wasm.NameMap) (tokenParser, error) {
+	p.addFunctionName(name)
 	p.module.FunctionSection = append(p.module.FunctionSection, typeIdx)
 	p.module.CodeSection = append(p.module.CodeSection, code)
 	p.addLocalNames(localNames)
@@ -625,8 +586,6 @@ func (p *moduleParser) resolveTypeIndices(module *wasm.Module) error {
 
 // resolveFunctionIndices ensures any indices point are numeric or returns a FormatError if they cannot be bound.
 func (p *moduleParser) resolveFunctionIndices(module *wasm.Module) error {
-	// TODO: remove when funcParser points to the correct index
-	importCount := p.funcNamespace.count - uint32(len(module.CodeSection))
 	for _, unresolved := range p.funcNamespace.unresolvedIndices {
 		target, err := p.funcNamespace.resolve(unresolved)
 		if err != nil {
@@ -637,7 +596,7 @@ func (p *moduleParser) resolveFunctionIndices(module *wasm.Module) error {
 			if target > 255 {
 				return errors.New("TODO: unresolved function indexes that don't fit in a byte")
 			}
-			module.CodeSection[unresolved.idx-importCount].Body[unresolved.bodyOffset] = byte(target)
+			module.CodeSection[unresolved.idx].Body[unresolved.bodyOffset] = byte(target)
 		case wasm.SectionIDExport:
 			p.unresolvedExports[unresolved.idx].Index = target
 		case wasm.SectionIDStart:
diff --git a/wasm/text/func_parser.go b/wasm/text/func_parser.go
index 3473ba86..efac06ed 100644
--- a/wasm/text/func_parser.go
+++ b/wasm/text/func_parser.go
@@ -8,11 +8,11 @@ import (
 	"github.com/tetratelabs/wazero/wasm/internal/leb128"
 )
 
-func newFuncParser(funcNamespace *indexNamespace, onFunc onFunc) *funcParser {
-	return &funcParser{funcNamespace: funcNamespace, onFunc: onFunc}
+func newFuncParser(typeUseParser *typeUseParser, funcNamespace *indexNamespace, onFunc onFunc) *funcParser {
+	return &funcParser{typeUseParser: typeUseParser, funcNamespace: funcNamespace, onFunc: onFunc}
 }
 
-type onFunc func(typeIdx wasm.Index, code *wasm.Code, localNames wasm.NameMap) (tokenParser, error)
+type onFunc func(typeIdx wasm.Index, code *wasm.Code, name string, localNames wasm.NameMap) (tokenParser, error)
 
 // funcParser parses any instructions and dispatches to onFunc.
 //
@@ -25,9 +25,14 @@ type funcParser struct {
 	// onFunc is called when complete parsing the body. Unless testing, this should be moduleParser.onFuncEnd
 	onFunc onFunc
 
+	// typeUseParser is described by moduleParser.typeUseParser
+	typeUseParser *typeUseParser
+
 	// funcNamespace is described by moduleParser.funcNamespace
 	funcNamespace *indexNamespace
 
+	currentIdx        wasm.Index
+	currentName       string
 	currentTypeIdx    wasm.Index
 	currentParamNames wasm.NameMap
 
@@ -40,17 +45,63 @@ type funcParser struct {
 var end = []byte{wasm.OpcodeEnd}
 var codeEnd = &wasm.Code{Body: end}
 
-// begin is a tokenParser that starts after a type use.
+// begin should be called after reaching the "func" keyword in a module field. Parsing continues until onFunc or
+// error.
+//
+// This stage records the ID of the current function, if present, and resumes with onFunc.
+//
+// Ex. A func ID is present `(func $main nop)`
+//                  records main --^     ^
+//              parseFunc resumes here --+
+//
+// Ex. No func ID `(func nop)`
+//     calls parseFunc --^
+func (p *funcParser) begin(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+	if tok == tokenID { // Ex. $main
+		if id, err := p.funcNamespace.setID(tokenBytes); err != nil {
+			return nil, err
+		} else {
+			p.currentName = id
+		}
+		return p.parseFunc, nil
+	}
+	p.currentName = ""
+	return p.parseFunc(tok, tokenBytes, line, col)
+}
+
+// parseFunc passes control to the typeUseParser until any signature is read, then funcParser until and locals or body
+// are read. Finally, this finishes via endFunc.
+//
+// Ex. `(module (func $math.pi (result f32))`
+//               starts here --^           ^
+//                  endFunc resumes here --+
+//
+// Ex.    `(module (func $math.pi (result f32) (local i32) )`
+//                  starts here --^            ^           ^
+//      funcParser.afterTypeUse resumes here --+           |
+//                                  endFunc resumes here --+
+//
+// Ex. If there is no signature `(func)`
+//              calls endFunc here ---^
+func (p *funcParser) parseFunc(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+	if tok == tokenID { // Ex. (func $main $main)
+		return nil, fmt.Errorf("redundant ID %s", tokenBytes)
+	}
+
+	return p.typeUseParser.begin(wasm.SectionIDFunction, p.currentIdx, p.afterTypeUse, tok, tokenBytes, line, col)
+}
+
+// afterTypeUse is a tokenParser that starts after a type use.
 //
 // The onFunc field is invoked once any instructions are written into currentBody.
 //
 // Ex. Given the source `(module (func nop))`
-//                 begin starts here --^  ^
+//          afterTypeUse starts here --^  ^
 //                    calls onFunc here --+
-func (p *funcParser) begin(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+func (p *funcParser) afterTypeUse(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 	switch pos {
 	case onTypeUseEndField:
-		return p.onFunc(typeIdx, codeEnd, paramNames)
+		return p.onFunc(typeIdx, codeEnd, p.currentName, paramNames)
 	case onTypeUseUnhandledField:
 		return sExpressionsUnsupported(tok, tokenBytes, line, col)
 	}
@@ -66,8 +117,12 @@ func sExpressionsUnsupported(tok tokenType, tokenBytes []byte, _, _ uint32) (tok
 		return nil, unexpectedToken(tok, tokenBytes)
 	}
 	switch string(tokenBytes) {
-	case "result", "param": // TODO: local
-		return nil, fmt.Errorf("%s declared out of order", tokenBytes)
+	case "param":
+		return nil, errors.New("param after result")
+	case "result":
+		return nil, errors.New("duplicate result")
+	case "local":
+		return nil, errors.New("TODO: local")
 	}
 	return nil, fmt.Errorf("TODO: s-expressions are not yet supported: %s", tokenBytes)
 }
@@ -112,15 +167,14 @@ func (p *funcParser) end() (tokenParser, error) {
 	} else {
 		code = &wasm.Code{Body: append(p.currentBody, wasm.OpcodeEnd)}
 	}
-	return p.onFunc(p.currentTypeIdx, code, p.currentParamNames)
+	return p.onFunc(p.currentTypeIdx, code, p.currentName, p.currentParamNames)
 }
 
 // parseFuncIndex parses an index in the function namespace and appends it to the currentBody. If it was an ID, a
 // placeholder byte(0) is added instead and will be resolved later.
 func (p *funcParser) parseFuncIndex(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	funcIdx := p.funcNamespace.count // TODO: move all function logic here and use the code index, not including imports
 	bodyOffset := uint32(len(p.currentBody))
-	idx, resolved, err := p.funcNamespace.parseIndex(wasm.SectionIDCode, funcIdx, bodyOffset, tok, tokenBytes, line, col)
+	idx, resolved, err := p.funcNamespace.parseIndex(wasm.SectionIDCode, p.currentIdx, bodyOffset, tok, tokenBytes, line, col)
 	if err != nil {
 		return nil, err
 	}
diff --git a/wasm/text/memory_parser.go b/wasm/text/memory_parser.go
index 46491c11..3137e666 100644
--- a/wasm/text/memory_parser.go
+++ b/wasm/text/memory_parser.go
@@ -33,7 +33,7 @@ type memoryParser struct {
 	currentMax *uint32
 }
 
-// begin should be called after reaching the "memory" keyword in a memory field. Parsing continues until onMemory or
+// begin should be called after reaching the "memory" keyword in a module field. Parsing continues until onMemory or
 // error.
 //
 // This stage records the ID of the current memory, if present, and resumes with beginMin.
diff --git a/wasm/text/type_parser.go b/wasm/text/type_parser.go
index 5e44d42a..d0d3cc01 100644
--- a/wasm/text/type_parser.go
+++ b/wasm/text/type_parser.go
@@ -50,7 +50,7 @@ type typeParser struct {
 	parsedParamID bool
 }
 
-// begin should be called after reaching the "type" keyword in a type field. Parsing continues until onType or error.
+// begin should be called after reaching the "type" keyword in a module field. Parsing continues until onType or error.
 //
 // This stage records the ID of the current type, if present, and resumes with tryBeginFunc.
 //


Test patch:
diff --git a/tests/bench/interface_bench_test.go b/tests/bench/interface_bench_test.go
new file mode 100644
index 00000000..b4c1b827
--- /dev/null
+++ b/tests/bench/interface_bench_test.go
@@ -0,0 +1,50 @@
+package bench
+
+import (
+	"encoding/binary"
+	"testing"
+)
+
+func BenchmarkInterfaceVSUnexport(b *testing.B) {
+	const size = 100
+	const val = 55
+	var memInterface MemoryInstanceInterface = &memoryInstance{buffer: make([]byte, size)}
+	var memStruct = &memoryInstance{buffer: make([]byte, size)}
+
+	b.Run("interface", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			for addr := uint32(0); addr < size+1000; addr++ {
+				memInterface.PutUint32(addr, val)
+			}
+		}
+	})
+	b.Run("struct", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			for addr := uint32(0); addr < size+1000; addr++ {
+				memStruct.PutUint32(addr, val)
+			}
+		}
+	})
+}
+
+type MemoryInstanceInterface interface {
+	PutUint32(addr uint32, val uint32) bool
+}
+
+type memoryInstance struct {
+	buffer []byte // inexpert only this
+	Min    uint32 // PS especially below would make sense to also not export as export makes it mutable
+	Max    *uint32
+}
+
+func (m *memoryInstance) validateAddrRange(addr uint32, rangeSize uint64) bool {
+	return uint64(addr) < uint64(len(m.buffer)) && rangeSize <= uint64(len(m.buffer))-uint64(addr)
+}
+
+func (m *memoryInstance) PutUint32(addr uint32, val uint32) bool {
+	if !m.validateAddrRange(addr, uint64(4)) {
+		return false
+	}
+	binary.LittleEndian.PutUint32(m.buffer[addr:], val)
+	return true
+}
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 08b335fa..5889d474 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -23,137 +23,6 @@ import (
 	"github.com/tetratelabs/wazero/wasm/internal/wazeroir"
 )
 
-// TODO: have some utility functions to reduce loc here: https://github.com/tetratelabs/wazero/issues/100
-
-type jitEnv struct {
-	eng            *engine
-	moduleInstance *wasm.ModuleInstance
-}
-
-func (j *jitEnv) stackTopAsByte() byte {
-	return byte(j.stack()[j.stackPointer()-1])
-}
-
-func (j *jitEnv) stackTopAsUint16() uint16 {
-	return uint16(j.stack()[j.stackPointer()-1])
-}
-
-func (j *jitEnv) stackTopAsUint32() uint32 {
-	return uint32(j.stack()[j.stackPointer()-1])
-}
-
-func (j *jitEnv) stackTopAsInt32() int32 {
-	return int32(j.stack()[j.stackPointer()-1])
-}
-func (j *jitEnv) stackTopAsUint64() uint64 {
-	return j.stack()[j.stackPointer()-1]
-}
-
-func (j *jitEnv) stackTopAsInt64() int64 {
-	return int64(j.stack()[j.stackPointer()-1])
-}
-
-func (j *jitEnv) stackTopAsFloat32() float32 {
-	return math.Float32frombits(uint32(j.stack()[j.stackPointer()-1]))
-}
-
-func (j *jitEnv) stackTopAsFloat64() float64 {
-	return math.Float64frombits(j.stack()[j.stackPointer()-1])
-}
-
-func (j *jitEnv) memory() []byte {
-	return j.moduleInstance.Memory.Buffer
-}
-
-func (j *jitEnv) stack() []uint64 {
-	return j.eng.valueStack
-}
-
-func (j *jitEnv) jitStatus() jitCallStatusCode {
-	return j.eng.exitContext.statusCode
-}
-
-func (j *jitEnv) functionCallAddress() wasm.FunctionAddress {
-	return j.eng.exitContext.functionCallAddress
-}
-
-func (j *jitEnv) stackPointer() uint64 {
-	return j.eng.valueStackContext.stackPointer
-}
-
-func (j *jitEnv) stackBasePointer() uint64 {
-	return j.eng.valueStackContext.stackBasePointer
-}
-
-func (j *jitEnv) setStackPointer(sp uint64) {
-	j.eng.valueStackContext.stackPointer = sp
-}
-
-func (j *jitEnv) addGlobals(g ...*wasm.GlobalInstance) {
-	j.moduleInstance.Globals = append(j.moduleInstance.Globals, g...)
-}
-
-func (j *jitEnv) getGlobal(index uint32) uint64 {
-	return j.moduleInstance.Globals[index].Val
-}
-
-func (j *jitEnv) setTable(table []wasm.TableElement) {
-	j.moduleInstance.Tables[0] = &wasm.TableInstance{Table: table}
-}
-
-func (j *jitEnv) callFrameStackPeek() *callFrame {
-	return &j.eng.callFrameStack[j.eng.globalContext.callFrameStackPointer-1]
-}
-
-func (j *jitEnv) callFrameStackPointer() uint64 {
-	return j.eng.globalContext.callFrameStackPointer
-}
-
-func (j *jitEnv) setValueStackBasePointer(sp uint64) {
-	j.eng.valueStackContext.stackBasePointer = sp
-}
-
-func (j *jitEnv) setCallFrameStackPointer(sp uint64) {
-	j.eng.globalContext.callFrameStackPointer = sp
-}
-
-func (j *jitEnv) setPreviousCallFrameStackPointer(sp uint64) {
-	j.eng.globalContext.previousCallFrameStackPointer = sp
-}
-
-func (j *jitEnv) engine() *engine {
-	return j.eng
-}
-
-func (j *jitEnv) exec(code []byte) {
-	compiledFunction := &compiledFunction{
-		codeSegment:        code,
-		codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
-		source: &wasm.FunctionInstance{
-			FunctionType: &wasm.TypeInstance{Type: &wasm.FunctionType{}},
-		},
-	}
-	j.eng.pushCallFrame(compiledFunction)
-
-	jitcall(
-		uintptr(unsafe.Pointer(&code[0])),
-		uintptr(unsafe.Pointer(j.eng)),
-	)
-}
-
-const defaultMemoryPageNumInTest = 2
-
-func newJITEnvironment() *jitEnv {
-	return &jitEnv{
-		eng: newEngine(),
-		moduleInstance: &wasm.ModuleInstance{
-			Memory:  &wasm.MemoryInstance{Buffer: make([]byte, wasm.MemoryPageSize*defaultMemoryPageNumInTest)},
-			Tables:  []*wasm.TableInstance{{}},
-			Globals: []*wasm.GlobalInstance{},
-		},
-	}
-}
-
 func (j *jitEnv) requireNewCompiler(t *testing.T) *amd64Compiler {
 	b, err := asm.NewBuilder("amd64", 128)
 	require.NoError(t, err)
diff --git a/wasm/jit/jit_arm64_test.go b/wasm/jit/jit_arm64_test.go
new file mode 100644
index 00000000..66270619
--- /dev/null
+++ b/wasm/jit/jit_arm64_test.go
@@ -0,0 +1,102 @@
+//go:build arm64
+// +build arm64
+
+package jit
+
+import (
+	"testing"
+	"unsafe"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/tetratelabs/wazero/wasm"
+)
+
+func (j *jitEnv) requireNewCompiler(t *testing.T) *arm64Compiler {
+	cmp, err := newCompiler(&wasm.FunctionInstance{ModuleInstance: j.moduleInstance}, nil)
+	require.NoError(t, err)
+	ret, ok := cmp.(*arm64Compiler)
+	require.True(t, ok)
+	return ret
+}
+
+func TestArm64CompilerEndToEnd(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		body []byte
+	}{
+		{name: "empty", body: []byte{wasm.OpcodeEnd}},
+		{name: "br .return", body: []byte{wasm.OpcodeBr, 0, wasm.OpcodeEnd}},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			engine := newEngine()
+			f := &wasm.FunctionInstance{
+				FunctionType: &wasm.TypeInstance{Type: &wasm.FunctionType{}},
+				Body:         tc.body,
+			}
+			err := engine.Compile(f)
+			require.NoError(t, err)
+			_, err = engine.Call(f)
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestArchContextOffsetInEngine(t *testing.T) {
+	var eng engine
+	// If this fails, we have to fix jit_arm64.s as well.
+	require.Equal(t, int(unsafe.Offsetof(eng.jitCallReturnAddress)), engineArchContextJITCallReturnAddressOffset)
+}
+
+func TestArm64Compiler_returnFunction(t *testing.T) {
+	env := newJITEnvironment()
+
+	// Build codes.
+	compiler := env.requireNewCompiler(t)
+	err := compiler.emitPreamble()
+	require.NoError(t, err)
+	compiler.returnFunction()
+
+	// Generate the code under test.
+	code, _, _, err := compiler.generate()
+	require.NoError(t, err)
+
+	// Run native code.
+	env.exec(code)
+
+	// JIT status on engine must be returned.
+	require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	// Plus, the call frame stack pointer must be zero after return.
+	require.Equal(t, uint64(0), env.callFrameStackPointer())
+}
+
+func TestArm64Compiler_exit(t *testing.T) {
+	for _, s := range []jitCallStatusCode{
+		jitCallStatusCodeReturned,
+		jitCallStatusCodeCallHostFunction,
+		jitCallStatusCodeCallBuiltInFunction,
+		jitCallStatusCodeUnreachable,
+	} {
+		t.Run(s.String(), func(t *testing.T) {
+
+			env := newJITEnvironment()
+
+			// Build codes.
+			compiler := env.requireNewCompiler(t)
+			err := compiler.emitPreamble()
+			require.NoError(t, err)
+			compiler.exit(s)
+
+			// Generate the code under test.
+			code, _, _, err := compiler.generate()
+			require.NoError(t, err)
+
+			// Run codes
+			env.exec(code)
+
+			// JIT status on engine must be updated.
+			require.Equal(t, s, env.jitStatus())
+		})
+	}
+}
diff --git a/wasm/jit/jit_test.go b/wasm/jit/jit_test.go
index 5945c486..9d41a381 100644
--- a/wasm/jit/jit_test.go
+++ b/wasm/jit/jit_test.go
@@ -1,11 +1,144 @@
 package jit
 
 import (
+	"math"
 	"os"
 	"runtime"
 	"testing"
+	"unsafe"
+
+	"github.com/tetratelabs/wazero/wasm"
 )
 
+type jitEnv struct {
+	eng            *engine
+	moduleInstance *wasm.ModuleInstance
+}
+
+func (j *jitEnv) stackTopAsByte() byte {
+	return byte(j.stack()[j.stackPointer()-1])
+}
+
+func (j *jitEnv) stackTopAsUint16() uint16 {
+	return uint16(j.stack()[j.stackPointer()-1])
+}
+
+func (j *jitEnv) stackTopAsUint32() uint32 {
+	return uint32(j.stack()[j.stackPointer()-1])
+}
+
+func (j *jitEnv) stackTopAsInt32() int32 {
+	return int32(j.stack()[j.stackPointer()-1])
+}
+func (j *jitEnv) stackTopAsUint64() uint64 {
+	return j.stack()[j.stackPointer()-1]
+}
+
+func (j *jitEnv) stackTopAsInt64() int64 {
+	return int64(j.stack()[j.stackPointer()-1])
+}
+
+func (j *jitEnv) stackTopAsFloat32() float32 {
+	return math.Float32frombits(uint32(j.stack()[j.stackPointer()-1]))
+}
+
+func (j *jitEnv) stackTopAsFloat64() float64 {
+	return math.Float64frombits(j.stack()[j.stackPointer()-1])
+}
+
+func (j *jitEnv) memory() []byte {
+	return j.moduleInstance.Memory.Buffer
+}
+
+func (j *jitEnv) stack() []uint64 {
+	return j.eng.valueStack
+}
+
+func (j *jitEnv) jitStatus() jitCallStatusCode {
+	return j.eng.exitContext.statusCode
+}
+
+func (j *jitEnv) functionCallAddress() wasm.FunctionAddress {
+	return j.eng.exitContext.functionCallAddress
+}
+
+func (j *jitEnv) stackPointer() uint64 {
+	return j.eng.valueStackContext.stackPointer
+}
+
+func (j *jitEnv) stackBasePointer() uint64 {
+	return j.eng.valueStackContext.stackBasePointer
+}
+
+func (j *jitEnv) setStackPointer(sp uint64) {
+	j.eng.valueStackContext.stackPointer = sp
+}
+
+func (j *jitEnv) addGlobals(g ...*wasm.GlobalInstance) {
+	j.moduleInstance.Globals = append(j.moduleInstance.Globals, g...)
+}
+
+func (j *jitEnv) getGlobal(index uint32) uint64 {
+	return j.moduleInstance.Globals[index].Val
+}
+
+func (j *jitEnv) setTable(table []wasm.TableElement) {
+	j.moduleInstance.Tables[0] = &wasm.TableInstance{Table: table}
+}
+
+func (j *jitEnv) callFrameStackPeek() *callFrame {
+	return &j.eng.callFrameStack[j.eng.globalContext.callFrameStackPointer-1]
+}
+
+func (j *jitEnv) callFrameStackPointer() uint64 {
+	return j.eng.globalContext.callFrameStackPointer
+}
+
+func (j *jitEnv) setValueStackBasePointer(sp uint64) {
+	j.eng.valueStackContext.stackBasePointer = sp
+}
+
+func (j *jitEnv) setCallFrameStackPointer(sp uint64) {
+	j.eng.globalContext.callFrameStackPointer = sp
+}
+
+func (j *jitEnv) setPreviousCallFrameStackPointer(sp uint64) {
+	j.eng.globalContext.previousCallFrameStackPointer = sp
+}
+
+func (j *jitEnv) engine() *engine {
+	return j.eng
+}
+
+func (j *jitEnv) exec(code []byte) {
+	compiledFunction := &compiledFunction{
+		codeSegment:        code,
+		codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
+		source: &wasm.FunctionInstance{
+			FunctionType: &wasm.TypeInstance{Type: &wasm.FunctionType{}},
+		},
+	}
+	j.eng.pushCallFrame(compiledFunction)
+
+	jitcall(
+		uintptr(unsafe.Pointer(&code[0])),
+		uintptr(unsafe.Pointer(j.eng)),
+	)
+}
+
+const defaultMemoryPageNumInTest = 2
+
+func newJITEnvironment() *jitEnv {
+	return &jitEnv{
+		eng: newEngine(),
+		moduleInstance: &wasm.ModuleInstance{
+			Memory:  &wasm.MemoryInstance{Buffer: make([]byte, wasm.MemoryPageSize*defaultMemoryPageNumInTest)},
+			Tables:  []*wasm.TableInstance{{}},
+			Globals: []*wasm.GlobalInstance{},
+		},
+	}
+}
+
 func TestMain(m *testing.M) {
 	if runtime.GOARCH != "amd64" && runtime.GOARCH != "arm64" {
 		// JIT is currently implemented only for amd64 or arm64.
diff --git a/wasm/text/decoder_test.go b/wasm/text/decoder_test.go
index 22cbb892..26dbb815 100644
--- a/wasm/text/decoder_test.go
+++ b/wasm/text/decoder_test.go
@@ -1830,9 +1830,9 @@ func TestParseModule_Errors(t *testing.T) {
 			expectedErr: "1:34: unknown type: f65 in module.func[0].param[1]",
 		},
 		{
-			name:        "func double result",
+			name:        "func duplicate result",
 			input:       "(module (func (param i32) (result i32) (result i32)))",
-			expectedErr: "1:41: result declared out of order in module.func[0]",
+			expectedErr: "1:41: duplicate result in module.func[0]",
 		},
 		{
 			name:        "func double result type",
@@ -1922,7 +1922,7 @@ func TestParseModule_Errors(t *testing.T) {
 		{
 			name:        "func param after result",
 			input:       "(module (func (result i32) (param i32)))",
-			expectedErr: "1:29: param declared out of order in module.func[0]",
+			expectedErr: "1:29: param after result in module.func[0]",
 		},
 		{
 			name:        "func points nowhere",
diff --git a/wasm/text/func_parser_test.go b/wasm/text/func_parser_test.go
index 77852b18..0c2e2e7c 100644
--- a/wasm/text/func_parser_test.go
+++ b/wasm/text/func_parser_test.go
@@ -46,12 +46,12 @@ func TestFuncParser(t *testing.T) {
 
 		t.Run(tc.name, func(t *testing.T) {
 			var parsedCode *wasm.Code
-			var setFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, localNames wasm.NameMap) (tokenParser, error) {
+			var setFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, name string, localNames wasm.NameMap) (tokenParser, error) {
 				parsedCode = code
 				return parseErr, nil
 			}
 
-			require.NoError(t, parseFunc(newFuncParser(newIndexNamespace(), setFunc), tc.source))
+			require.NoError(t, parseFunc(newFuncParser(&typeUseParser{module: &wasm.Module{}}, newIndexNamespace(), setFunc), tc.source))
 			require.Equal(t, tc.expected, parsedCode)
 		})
 	}
@@ -102,12 +102,12 @@ func TestFuncParser_Call_Unresolved(t *testing.T) {
 
 		t.Run(tc.name, func(t *testing.T) {
 			var parsedCode *wasm.Code
-			var setFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, localNames wasm.NameMap) (tokenParser, error) {
+			var setFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, name string, localNames wasm.NameMap) (tokenParser, error) {
 				parsedCode = code
 				return parseErr, nil
 			}
 
-			fp := newFuncParser(newIndexNamespace(), setFunc)
+			fp := newFuncParser(&typeUseParser{module: &wasm.Module{}}, newIndexNamespace(), setFunc)
 			require.NoError(t, parseFunc(fp, tc.source))
 			require.Equal(t, tc.expectedCode, parsedCode)
 			require.Equal(t, []*unresolvedIndex{tc.expectedUnresolvedIndex}, fp.funcNamespace.unresolvedIndices)
@@ -159,12 +159,12 @@ func TestFuncParser_Call_Resolved(t *testing.T) {
 			funcNamespace.count++
 
 			var parsedCode *wasm.Code
-			var setFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, localNames wasm.NameMap) (tokenParser, error) {
+			var setFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, name string, localNames wasm.NameMap) (tokenParser, error) {
 				parsedCode = code
 				return parseErr, nil
 			}
 
-			fp := newFuncParser(funcNamespace, setFunc)
+			fp := newFuncParser(&typeUseParser{module: &wasm.Module{}}, funcNamespace, setFunc)
 			require.NoError(t, parseFunc(fp, tc.source))
 			require.Equal(t, tc.expected, parsedCode)
 		})
@@ -206,14 +206,14 @@ func TestFuncParser_Errors(t *testing.T) {
 			expectedErr: "1:8: TODO: s-expressions are not yet supported: f32.const",
 		},
 		{
-			name:        "param out of order", // because this parser is after the type use, so it must be wrong
-			source:      "(func (param i32))",
-			expectedErr: "1:8: param declared out of order",
+			name:        "param after result",
+			source:      "(func (result i32) (param i32))",
+			expectedErr: "1:21: param after result",
 		},
 		{
-			name:        "result out of order", // because this parser is after the type use, so it must be wrong
-			source:      "(func (result i32))",
-			expectedErr: "1:8: result declared out of order",
+			name:        "duplicate result",
+			source:      "(func (result i32) (result i32))",
+			expectedErr: "1:21: duplicate result",
 		},
 	}
 
@@ -221,23 +221,18 @@ func TestFuncParser_Errors(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			fp := newFuncParser(newIndexNamespace(), failOnFunc)
+			fp := newFuncParser(&typeUseParser{module: &wasm.Module{}}, newIndexNamespace(), failOnFunc)
 			require.EqualError(t, parseFunc(fp, tc.source), tc.expectedErr)
 		})
 	}
 }
 
-var failOnFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, localNames wasm.NameMap) (tokenParser, error) {
+var failOnFunc onFunc = func(typeIdx wasm.Index, code *wasm.Code, name string, localNames wasm.NameMap) (tokenParser, error) {
 	return nil, errors.New("unexpected to call onFunc on error")
 }
 
 func parseFunc(fp *funcParser, source string) error {
-	// TODO: all func hooks into func_parser.go, so that we don't have to fake the onTypeUse position
-	var parser tokenParser = func(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-		return fp.begin(0, nil, onTypeUseUnhandledToken, tok, tokenBytes, line, col)
-	}
-
-	line, col, err := lex(skipTokens(2, parser), []byte(source)) // skip the leading (func
+	line, col, err := lex(skipTokens(2, fp.begin), []byte(source)) // skip the leading (func
 	if err != nil {
 		err = &FormatError{Line: line, Col: col, cause: err}
 	}