Repository: tetratelabs/wazero
Pull request #221 (2022-02-10 02:44:41)

jit(arm64): support call instruction
part of #187 

mathetake: ok cool passed tests

Patch:
diff --git a/wasm/binary/decoder.go b/wasm/binary/decoder.go
index bc7aee5a..d1af58a1 100644
--- a/wasm/binary/decoder.go
+++ b/wasm/binary/decoder.go
@@ -27,8 +27,10 @@ func DecodeModule(binary []byte) (*wasm.Module, error) {
 
 	m := &wasm.Module{}
 	for {
-		sectionID := make([]byte, 1)
-		if _, err := io.ReadFull(r, sectionID); err == io.EOF {
+		// TODO: except custom sections, all others are required to be in order, but we aren't checking yet.
+		// See https://www.w3.org/TR/wasm-core-1/#modules%E2%91%A0%E2%93%AA
+		sectionID, err := r.ReadByte()
+		if err == io.EOF {
 			break
 		} else if err != nil {
 			return nil, fmt.Errorf("read section id: %w", err)
@@ -36,11 +38,11 @@ func DecodeModule(binary []byte) (*wasm.Module, error) {
 
 		sectionSize, _, err := leb128.DecodeUint32(r)
 		if err != nil {
-			return nil, fmt.Errorf("get size of section for id=%d: %v", sectionID[0], err)
+			return nil, fmt.Errorf("get size of section %s: %v", wasm.SectionIDName(sectionID), err)
 		}
 
 		sectionContentStart := r.Len()
-		switch sectionID[0] {
+		switch sectionID {
 		case wasm.SectionIDCustom:
 			// First, validate the section and determine if the section for this name has already been set
 			name, nameSize, decodeErr := decodeUTF8(r, "custom section name")
@@ -53,25 +55,19 @@ func DecodeModule(binary []byte) (*wasm.Module, error) {
 			} else if name == "name" && m.NameSection != nil {
 				err = fmt.Errorf("redundant custom section %s", name)
 				break
-			} else if _, ok := m.CustomSections[name]; ok {
-				err = fmt.Errorf("redundant custom section %s", name)
-				break
 			}
 
-			// Now, either decode the NameSection or store an unsupported one
-			// TODO: Do we care to store something we don't use? We could also skip it!
-			data, dataErr := readCustomSectionData(r, sectionSize-nameSize)
-			if dataErr != nil {
-				err = dataErr
-			} else if name == "name" {
-				m.NameSection, err = decodeNameSection(data)
+			// Now, either decode the NameSection or skip an unsupported one
+			limit := sectionSize - nameSize
+			if name == "name" {
+				m.NameSection, err = decodeNameSection(r, uint64(limit))
 			} else {
-				if m.CustomSections == nil {
-					m.CustomSections = map[string][]byte{name: data}
-				} else {
-					m.CustomSections[name] = data
+				// Note: Not Seek because it doesn't err when given an offset past EOF. Rather, it leads to undefined state.
+				if _, err = io.CopyN(io.Discard, r, int64(limit)); err != nil {
+					return nil, fmt.Errorf("failed to skip name[%s]: %w", name, err)
 				}
 			}
+
 		case wasm.SectionIDType:
 			m.TypeSection, err = decodeTypeSection(r)
 		case wasm.SectionIDImport:
@@ -104,11 +100,11 @@ func DecodeModule(binary []byte) (*wasm.Module, error) {
 		}
 
 		if err != nil {
-			return nil, fmt.Errorf("section ID %d: %v", sectionID[0], err)
+			return nil, fmt.Errorf("section %s: %v", wasm.SectionIDName(sectionID), err)
 		}
 	}
 
-	functionCount, codeCount := len(m.FunctionSection), len(m.CodeSection)
+	functionCount, codeCount := m.SectionElementCount(wasm.SectionIDFunction), m.SectionElementCount(wasm.SectionIDCode)
 	if functionCount != codeCount {
 		return nil, fmt.Errorf("function and code section have inconsistent lengths: %d != %d", functionCount, codeCount)
 	}
diff --git a/wasm/binary/encoder.go b/wasm/binary/encoder.go
index 7f755ba3..3db533cb 100644
--- a/wasm/binary/encoder.go
+++ b/wasm/binary/encoder.go
@@ -11,47 +11,46 @@ var sizePrefixedName = []byte{4, 'n', 'a', 'm', 'e'}
 // See https://www.w3.org/TR/wasm-core-1/#binary-format%E2%91%A0
 func EncodeModule(m *wasm.Module) (bytes []byte) {
 	bytes = append(magic, version...)
-	for name, data := range m.CustomSections {
-		bytes = append(bytes, encodeCustomSection(name, data)...)
-	}
-	if len(m.TypeSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDType) > 0 {
 		bytes = append(bytes, encodeTypeSection(m.TypeSection)...)
 	}
-	if len(m.ImportSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDImport) > 0 {
 		bytes = append(bytes, encodeImportSection(m.ImportSection)...)
 	}
-	if len(m.FunctionSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDFunction) > 0 {
 		bytes = append(bytes, encodeFunctionSection(m.FunctionSection)...)
 	}
-	if len(m.TableSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDTable) > 0 {
 		panic("TODO: TableSection")
 	}
-	if len(m.MemorySection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDMemory) > 0 {
 		bytes = append(bytes, encodeMemorySection(m.MemorySection)...)
 	}
-	if len(m.GlobalSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDGlobal) > 0 {
 		panic("TODO: GlobalSection")
 	}
-	if len(m.ExportSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDExport) > 0 {
 		bytes = append(bytes, encodeExportSection(m.ExportSection)...)
 	}
-	if m.StartSection != nil {
+	if m.SectionElementCount(wasm.SectionIDStart) > 0 {
 		bytes = append(bytes, encodeStartSection(*m.StartSection)...)
 	}
-	if len(m.ElementSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDElement) > 0 {
 		panic("TODO: ElementSection")
 	}
-	if len(m.CodeSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDCode) > 0 {
 		bytes = append(bytes, encodeCodeSection(m.CodeSection)...)
 	}
-	if len(m.DataSection) > 0 {
+	if m.SectionElementCount(wasm.SectionIDData) > 0 {
 		panic("TODO: DataSection")
 	}
-	// >> The name section should appear only once in a module, and only after the data section.
-	// See https://www.w3.org/TR/wasm-core-1/#binary-namesec
-	if m.NameSection != nil {
-		nameSection := append(sizePrefixedName, encodeNameSectionData(m.NameSection)...)
-		bytes = append(bytes, encodeSection(wasm.SectionIDCustom, nameSection)...)
+	if m.SectionElementCount(wasm.SectionIDCustom) > 0 {
+		// >> The name section should appear only once in a module, and only after the data section.
+		// See https://www.w3.org/TR/wasm-core-1/#binary-namesec
+		if m.NameSection != nil {
+			nameSection := append(sizePrefixedName, encodeNameSectionData(m.NameSection)...)
+			bytes = append(bytes, encodeSection(wasm.SectionIDCustom, nameSection)...)
+		}
 	}
 	return
 }
diff --git a/wasm/binary/names.go b/wasm/binary/names.go
index b9590522..dc64cdfd 100644
--- a/wasm/binary/names.go
+++ b/wasm/binary/names.go
@@ -27,17 +27,17 @@ const (
 // * LocalNames decode from subsection 2
 //
 // See https://www.w3.org/TR/wasm-core-1/#binary-namesec
-func decodeNameSection(data []byte) (result *wasm.NameSection, err error) {
+func decodeNameSection(r *bytes.Reader, limit uint64) (result *wasm.NameSection, err error) {
 	// TODO: add leb128 functions that work on []byte and offset. While using a reader allows us to reuse reader-based
 	// leb128 functions, it is less efficient, causes untestable code and in some cases more complex vs plain []byte.
-	r := bytes.NewReader(data)
 	result = &wasm.NameSection{}
 
 	// subsectionID is decoded if known, and skipped if not
 	var subsectionID uint8
 	// subsectionSize is the length to skip when the subsectionID is unknown
 	var subsectionSize uint32
-	for {
+	var bytesRead uint64
+	for limit > 0 {
 		if subsectionID, err = r.ReadByte(); err != nil {
 			if err == io.EOF {
 				return result, nil
@@ -45,11 +45,12 @@ func decodeNameSection(data []byte) (result *wasm.NameSection, err error) {
 			// TODO: untestable as this can't fail for a reason beside EOF reading a byte from a buffer
 			return nil, fmt.Errorf("failed to read a subsection ID: %w", err)
 		}
+		limit--
 
-		// TODO: unused except when skipping. This means we can pass on a corrupt length of a known subsection
-		if subsectionSize, _, err = leb128.DecodeUint32(r); err != nil {
+		if subsectionSize, bytesRead, err = leb128.DecodeUint32(r); err != nil {
 			return nil, fmt.Errorf("failed to read the size of subsection[%d]: %w", subsectionID, err)
 		}
+		limit -= bytesRead
 
 		switch subsectionID {
 		case subsectionIDModuleName:
@@ -66,11 +67,13 @@ func decodeNameSection(data []byte) (result *wasm.NameSection, err error) {
 			}
 		default: // Skip other subsections.
 			// Note: Not Seek because it doesn't err when given an offset past EOF. Rather, it leads to undefined state.
-			if _, err := io.CopyN(io.Discard, r, int64(subsectionSize)); err != nil {
+			if _, err = io.CopyN(io.Discard, r, int64(subsectionSize)); err != nil {
 				return nil, fmt.Errorf("failed to skip subsection[%d]: %w", subsectionID, err)
 			}
 		}
+		limit -= uint64(subsectionSize)
 	}
+	return
 }
 
 func decodeFunctionNames(r *bytes.Reader) (wasm.NameMap, error) {
diff --git a/wasm/binary/section.go b/wasm/binary/section.go
index ddc00f3f..7c587f97 100644
--- a/wasm/binary/section.go
+++ b/wasm/binary/section.go
@@ -9,14 +9,6 @@ import (
 	"github.com/tetratelabs/wazero/wasm/internal/leb128"
 )
 
-func readCustomSectionData(r *bytes.Reader, dataSize uint32) ([]byte, error) {
-	data := make([]byte, dataSize)
-	if _, err := io.ReadFull(r, data); err != nil {
-		return nil, fmt.Errorf("cannot read custom section data: %w", err)
-	}
-	return data, nil
-}
-
 func decodeTypeSection(r io.Reader) ([]*wasm.FunctionType, error) {
 	vs, _, err := leb128.DecodeUint32(r)
 	if err != nil {
@@ -165,7 +157,7 @@ func decodeExportSection(r *bytes.Reader) (map[string]*wasm.Export, error) {
 	return exportSection, nil
 }
 
-func decodeStartSection(r *bytes.Reader) (*uint32, error) {
+func decodeStartSection(r *bytes.Reader) (*wasm.Index, error) {
 	vs, _, err := leb128.DecodeUint32(r)
 	if err != nil {
 		return nil, fmt.Errorf("get size of vector: %w", err)
@@ -218,14 +210,6 @@ func decodeDataSection(r *bytes.Reader) ([]*wasm.DataSegment, error) {
 	return result, nil
 }
 
-// encodeCustomSection encodes the opaque bytes for the given name as a SectionIDCustom
-// See https://www.w3.org/TR/wasm-core-1/#binary-customsec
-func encodeCustomSection(name string, data []byte) []byte {
-	// The contents of a custom section is the non-empty name followed by potentially empty opaque data
-	contents := append(encodeSizePrefixed([]byte(name)), data...)
-	return encodeSection(wasm.SectionIDCustom, contents)
-}
-
 // encodeSection encodes the sectionID, the size of its contents in bytes, followed by the contents.
 // See https://www.w3.org/TR/wasm-core-1/#sections%E2%91%A0
 func encodeSection(sectionID wasm.SectionID, contents []byte) []byte {
@@ -260,10 +244,10 @@ func encodeImportSection(imports []*wasm.Import) []byte {
 // WebAssembly 1.0 (MVP) Binary Format.
 //
 // See https://www.w3.org/TR/wasm-core-1/#function-section%E2%91%A0
-func encodeFunctionSection(functions []wasm.Index) []byte {
-	contents := leb128.EncodeUint32(uint32(len(functions)))
-	for _, typeIndex := range functions {
-		contents = append(contents, leb128.EncodeUint32(typeIndex)...)
+func encodeFunctionSection(typeIndices []wasm.Index) []byte {
+	contents := leb128.EncodeUint32(uint32(len(typeIndices)))
+	for _, index := range typeIndices {
+		contents = append(contents, leb128.EncodeUint32(index)...)
 	}
 	return encodeSection(wasm.SectionIDFunction, contents)
 }
@@ -307,6 +291,6 @@ func encodeExportSection(exports map[string]*wasm.Export) []byte {
 // encodeStartSection encodes a SectionIDStart for the given function index in WebAssembly 1.0 (MVP) Binary Format.
 //
 // See https://www.w3.org/TR/wasm-core-1/#start-section%E2%91%A0
-func encodeStartSection(funcidx uint32) []byte {
+func encodeStartSection(funcidx wasm.Index) []byte {
 	return encodeSection(wasm.SectionIDStart, leb128.EncodeUint32(funcidx))
 }
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 308aa8bb..2998534f 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -4729,8 +4729,8 @@ func (c *amd64Compiler) callFunction(addr wasm.FunctionAddress, addrReg int16, f
 
 	// 2) Set engine.valueStackContext.stackBasePointer for the next function.
 	{
-		// At this point, tmpRegister holds the OLD stack base pointer. We could get the new frame's
-		// stack base pointer by "OLD stack base pointer + OLD stack pointer - # of function params"
+		// At this point, tmpRegister holds the old stack base pointer. We could get the new frame's
+		// stack base pointer by "old stack base pointer + old stack pointer - # of function params"
 		// See the comments in engine.pushCallFrame which does exactly the same calculation in Go.
 		calculateNextStackBasePointer := c.newProg()
 		calculateNextStackBasePointer.As = x86.AADDQ
@@ -4870,8 +4870,7 @@ func (c *amd64Compiler) returnFunction() error {
 	// Obtain the temporary registers to be used in the followings.
 	regs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 3)
 	if !found {
-		// This in theory never happen as all the registers must be free except addrReg.
-		return fmt.Errorf("could not find enough free registers")
+		return fmt.Errorf("BUG: all the registers should be free at this point")
 	}
 	c.locationStack.markRegisterUsed(regs...)
 
diff --git a/wasm/jit/jit_arm64.go b/wasm/jit/jit_arm64.go
index 0c1cd0fc..578209f1 100644
--- a/wasm/jit/jit_arm64.go
+++ b/wasm/jit/jit_arm64.go
@@ -77,6 +77,8 @@ type arm64Compiler struct {
 	labels map[string]*labelInfo
 	// stackPointerCeil is the greatest stack pointer value (from valueLocationStack) seen during compilation.
 	stackPointerCeil uint64
+	// afterAssembleCallback hold the callbacks which are called after assembling native code.
+	afterAssembleCallback []func(code []byte) error
 }
 
 // compile implements compiler.compile for the arm64 architecture.
@@ -89,10 +91,19 @@ func (c *arm64Compiler) compile() (code []byte, staticData compiledFunctionStati
 		stackPointerCeil = c.locationStack.stackPointerCeil
 	}
 
-	code, err = mmapCodeSegment(c.builder.Assemble())
+	original := c.builder.Assemble()
+
+	for _, cb := range c.afterAssembleCallback {
+		if err = cb(original); err != nil {
+			return
+		}
+	}
+
+	code, err = mmapCodeSegment(original)
 	if err != nil {
 		return
 	}
+
 	return
 }
 
@@ -136,9 +147,11 @@ func (c *arm64Compiler) markRegisterUsed(reg int16) {
 	c.locationStack.markRegisterUsed(reg)
 }
 
-func (c *arm64Compiler) markRegisterUnused(reg int16) {
-	if !isZeroRegister(reg) {
-		c.locationStack.markRegisterUnused(reg)
+func (c *arm64Compiler) markRegisterUnused(regs ...int16) {
+	for _, reg := range regs {
+		if !isZeroRegister(reg) {
+			c.locationStack.markRegisterUnused(reg)
+		}
 	}
 }
 
@@ -158,74 +171,63 @@ func (c *arm64Compiler) applyConstToRegisterInstruction(instruction obj.As, cons
 
 // applyMemoryToRegisterInstruction adds an instruction where source operand is a memory location and destination is a register.
 // baseRegister is the base absolute address in the memory, and offset is the offset from the absolute address in baseRegister.
-func (c *arm64Compiler) applyMemoryToRegisterInstruction(instruction obj.As, baseRegister int16, offset int64, destinationRegister int16) (err error) {
-	if offset > math.MaxInt16 {
-		// This is a bug in JIT copmiler: caller must check the offset at compilation time, and avoid such a large offset
-		// by loading the const to the register beforehand and then using applyRegisterOffsetMemoryToRegisterInstruction instead.
-		err = fmt.Errorf("memory offset must be smaller than or equal %d, but got %d", math.MaxInt16, offset)
-		return
+func (c *arm64Compiler) applyMemoryToRegisterInstruction(instruction obj.As, sourceBaseRegister int16, sourceOffsetConst int64, destinationRegister int16) {
+	if sourceOffsetConst > math.MaxInt16 {
+		// The assembler can take care of offsets larger than 2^15-1 by emitting additional instructions to load such large offset,
+		// but it uses "its" temporary register which we cannot track. Therefore, we avoid directly emitting memory load with large offsets,
+		// but instead load the constant manually to "our" temporary register, then emit the load with it.
+		c.applyConstToRegisterInstruction(arm64.AMOVD, sourceOffsetConst, reservedRegisterForTemporary)
+		inst := c.newProg()
+		inst.As = instruction
+		inst.From.Type = obj.TYPE_MEM
+		inst.From.Reg = sourceBaseRegister
+		inst.From.Index = reservedRegisterForTemporary
+		inst.From.Scale = 1
+		inst.To.Type = obj.TYPE_REG
+		inst.To.Reg = destinationRegister
+		c.addInstruction(inst)
+	} else {
+		inst := c.newProg()
+		inst.As = instruction
+		inst.From.Type = obj.TYPE_MEM
+		inst.From.Reg = sourceBaseRegister
+		inst.From.Offset = sourceOffsetConst
+		inst.To.Type = obj.TYPE_REG
+		inst.To.Reg = destinationRegister
+		c.addInstruction(inst)
 	}
-	inst := c.newProg()
-	inst.As = instruction
-	inst.From.Type = obj.TYPE_MEM
-	inst.From.Reg = baseRegister
-	inst.From.Offset = offset
-	inst.To.Type = obj.TYPE_REG
-	inst.To.Reg = destinationRegister
-	c.addInstruction(inst)
-	return
-}
-
-// applyRegisterOffsetMemoryToRegisterInstruction adds an instruction where source operand is a memory location and destination is a register.
-// The difference from applyMemoryToRegisterInstruction is that here we specify the offset by a register rather than offset constant.
-func (c *arm64Compiler) applyRegisterOffsetMemoryToRegisterInstruction(instruction obj.As, baseRegister, offsetRegister, destinationRegister int16) (err error) {
-	inst := c.newProg()
-	inst.As = instruction
-	inst.From.Type = obj.TYPE_MEM
-	inst.From.Reg = baseRegister
-	inst.From.Index = offsetRegister
-	inst.From.Scale = 1
-	inst.To.Type = obj.TYPE_REG
-	inst.To.Reg = destinationRegister
-	c.addInstruction(inst)
-	return nil
 }
 
 // applyRegisterToMemoryInstruction adds an instruction where destination operand is a memory location and source is a register.
 // This is the opposite of applyMemoryToRegisterInstruction.
-func (c *arm64Compiler) applyRegisterToMemoryInstruction(instruction obj.As, baseRegister int16, offset int64, source int16) (err error) {
-	if offset > math.MaxInt16 {
-		// This is a bug in JIT copmiler: caller must check the offset at compilation time, and avoid such a large offset
-		// by loading the const to the register beforehand and then using applyRegisterToRegisterOffsetMemoryInstruction instead.
-		err = fmt.Errorf("memory offset must be smaller than or equal %d, but got %d", math.MaxInt16, offset)
-		return
+func (c *arm64Compiler) applyRegisterToMemoryInstruction(instruction obj.As, sourceRegister int16, destinationBaseRegister int16, destinationOffsetConst int64) {
+	if destinationOffsetConst > math.MaxInt16 {
+		// The assembler can take care of offsets larger than 2^15-1 by emitting additional instructions to load such large offset,
+		// but we cannot track its temporary register. Therefore, we avoid directly emitting memory load with large offsets:
+		// load the constant manually to "our" temporary register, then emit the load with it.
+		c.applyConstToRegisterInstruction(arm64.AMOVD, destinationOffsetConst, reservedRegisterForTemporary)
+		inst := c.newProg()
+		inst.As = instruction
+		inst.To.Type = obj.TYPE_MEM
+		inst.To.Reg = destinationBaseRegister
+		inst.To.Index = reservedRegisterForTemporary
+		inst.To.Scale = 1
+		inst.From.Type = obj.TYPE_REG
+		inst.From.Reg = sourceRegister
+		c.addInstruction(inst)
+	} else {
+		inst := c.newProg()
+		inst.As = instruction
+		inst.To.Type = obj.TYPE_MEM
+		inst.To.Reg = destinationBaseRegister
+		inst.To.Offset = destinationOffsetConst
+		inst.From.Type = obj.TYPE_REG
+		inst.From.Reg = sourceRegister
+		c.addInstruction(inst)
 	}
-	inst := c.newProg()
-	inst.As = instruction
-	inst.To.Type = obj.TYPE_MEM
-	inst.To.Reg = baseRegister
-	inst.To.Offset = offset
-	inst.From.Type = obj.TYPE_REG
-	inst.From.Reg = source
-	c.addInstruction(inst)
-	return
 }
 
-// applyRegisterToRegisterOffsetMemoryInstruction adds an instruction where destination operand is a memory location and source is a register.
-// The difference from applyRegisterToMemoryInstruction is that here we specify the offset by a register rather than offset constant.
-func (c *arm64Compiler) applyRegisterToRegisterOffsetMemoryInstruction(instruction obj.As, baseRegister, offsetRegister, source int16) {
-	inst := c.newProg()
-	inst.As = instruction
-	inst.To.Type = obj.TYPE_MEM
-	inst.To.Reg = baseRegister
-	inst.To.Index = offsetRegister
-	inst.To.Scale = 1
-	inst.From.Type = obj.TYPE_REG
-	inst.From.Reg = source
-	c.addInstruction(inst)
-}
-
-// applyRegisterToRegisterOffsetMemoryInstruction adds an instruction where both destination and source operands are registers.
+// applyRegisterToRegisterInstruction adds an instruction where both destination and source operands are registers.
 func (c *arm64Compiler) applyRegisterToRegisterInstruction(instruction obj.As, from, to int16) {
 	inst := c.newProg()
 	inst.As = instruction
@@ -236,7 +238,7 @@ func (c *arm64Compiler) applyRegisterToRegisterInstruction(instruction obj.As, f
 	c.addInstruction(inst)
 }
 
-// applyRegisterToRegisterOffsetMemoryInstruction adds an instruction which takes two source operands on registers and one destination register operand.
+// applyTwoRegistersToRegisterInstruction adds an instruction which takes two source operands on registers and one destination register operand.
 func (c *arm64Compiler) applyTwoRegistersToRegisterInstruction(instruction obj.As, src1, src2, destination int16) {
 	inst := c.newProg()
 	inst.As = instruction
@@ -261,14 +263,22 @@ func (c *arm64Compiler) applyTwoRegistersToNoneInstruction(instruction obj.As, s
 	c.addInstruction(inst)
 }
 
-func (c *arm64Compiler) emitUnconditionalBRInstruction(targetType obj.AddrType) (jmp *obj.Prog) {
-	jmp = c.newProg()
-	jmp.As = obj.AJMP
-	jmp.To.Type = targetType
-	c.addInstruction(jmp)
+func (c *arm64Compiler) emitUnconditionalBranchInstruction() (br *obj.Prog) {
+	br = c.newProg()
+	br.As = obj.AJMP
+	br.To.Type = obj.TYPE_BRANCH
+	c.addInstruction(br)
 	return
 }
 
+func (c *arm64Compiler) emitUnconditionalBranchToAddressOnRegister(addressRegister int16) {
+	br := c.newProg()
+	br.As = obj.AJMP
+	br.To.Type = obj.TYPE_MEM
+	br.To.Reg = addressRegister
+	c.addInstruction(br)
+}
+
 func (c *arm64Compiler) String() (ret string) { return }
 
 // pushFunctionParams pushes any function parameters onto the stack, setting appropriate register types.
@@ -303,56 +313,114 @@ func (c *arm64Compiler) emitPreamble() error {
 
 // 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 (TODO).
+// Otherwise, we branch into the caller's return address.
 func (c *arm64Compiler) returnFunction() error {
-	// TODO: we don't support function calls yet.
-	// For now the following code just returns to Go code.
+	// Release all the registers as our calling convention requires the caller-save.
+	if err := c.releaseAllRegistersToStack(); err != nil {
+		return err
+	}
 
 	// Since we return from the function, we need to decrement the callframe stack pointer, and write it back.
-	callFramePointerReg, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
-	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine,
-		engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg); err != nil {
-		return err
+	tmpRegs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 3)
+	if !found {
+		return fmt.Errorf("BUG: all the registers should be free at this point")
 	}
+
+	// Alias for readability.
+	callFramePointerReg, callFrameStackTopAddressRegister, tmpReg := tmpRegs[0], tmpRegs[1], tmpRegs[2]
+
+	// First we decrement the callframe stack pointer.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg)
 	c.applyConstToRegisterInstruction(arm64.ASUBS, 1, callFramePointerReg)
-	if err := c.applyRegisterToMemoryInstruction(arm64.AMOVD, reservedRegisterForEngine,
-		engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg); err != nil {
-		return err
-	}
+	c.applyRegisterToMemoryInstruction(arm64.AMOVD, callFramePointerReg, reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset)
+
+	// Next we compare the decremented call frame stack pointer with the engine.precviousCallFrameStackPointer.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextPreviouscallFrameStackPointer,
+		tmpReg,
+	)
+	c.applyTwoRegistersToNoneInstruction(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)
+	c.exit(jitCallStatusCodeReturned)
+
+	// Otherwise, we have to jump to the caller's return address.
+	c.setBRTargetOnNext(brIfNotEqual)
+
+	// First, we have to calculate the caller callFrame's absolute address to aquire the return address.
+	//
+	// "tmpReg = &engine.callFrameStack[0]"
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackElement0AddressOffset,
+		tmpReg,
+	)
+	// "callFrameStackTopAddressRegister = tmpReg + callFramePointerReg << ${callFrameDataSizeMostSignificantSetBit}"
+	c.emitAddInstructionWithLeftShiftedRegister(
+		callFramePointerReg, callFrameDataSizeMostSignificantSetBit,
+		tmpReg,
+		callFrameStackTopAddressRegister,
+	)
 
-	return c.exit(jitCallStatusCodeReturned)
+	// At this point, we have
+	//
+	//      [......., ra.caller, rb.caller, rc.caller, _, ra.current, rb.current, rc.current, _, ...]  <- call frame stack's data region (somewhere in the memory)
+	//                                                  ^
+	//                               callFrameStackTopAddressRegister
+	//                   (absolute address of &callFrameStack[engine.callFrameStackPointer])
+	//
+	// where:
+	//      ra.* = callFrame.returnAddress
+	//      rb.* = callFrame.returnStackBasePointer
+	//      rc.* = callFrame.compiledFunction
+	//      _  = callFrame's padding (see comment on callFrame._ field.)
+	//
+	// What we have to do in the following is that
+	//   1) Set engine.valueStackContext.stackBasePointer to the value on "rb.caller".
+	//   2) Jump into the address of "ra.caller".
+
+	// 1) Set engine.valueStackContext.stackBasePointer to the value on "rb.caller".
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		// "rb.caller" is below the top address.
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerOffset),
+		tmpReg)
+	c.applyRegisterToMemoryInstruction(arm64.AMOVD,
+		tmpReg,
+		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset)
+
+	// 2) Branch into the address of "ra.caller".
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		// "rb.caller" is below the top address.
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnAddressOffset),
+		tmpReg)
+	c.emitUnconditionalBranchToAddressOnRegister(tmpReg)
+
+	c.locationStack.markRegisterUnused(tmpRegs...)
+	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 {
 	// Write the current stack pointer to the engine.stackPointer.
 	c.applyConstToRegisterInstruction(arm64.AMOVW, int64(c.locationStack.sp), reservedRegisterForTemporary)
-	if err := c.applyRegisterToMemoryInstruction(arm64.AMOVW, reservedRegisterForEngine,
-		engineValueStackContextStackPointerOffset, reservedRegisterForTemporary); err != nil {
-		return err
-	}
+	c.applyRegisterToMemoryInstruction(arm64.AMOVW, reservedRegisterForTemporary, reservedRegisterForEngine,
+		engineValueStackContextStackPointerOffset)
 
 	if status != 0 {
 		c.applyConstToRegisterInstruction(arm64.AMOVW, int64(status), reservedRegisterForTemporary)
-		if err := c.applyRegisterToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine,
-			engineExitContextJITCallStatusCodeOffset, reservedRegisterForTemporary); err != nil {
-			return err
-		}
+		c.applyRegisterToMemoryInstruction(arm64.AMOVWU, reservedRegisterForTemporary, reservedRegisterForEngine, engineExitContextJITCallStatusCodeOffset)
 	} else {
 		// If the status == 0, we use zero register to store zero.
-		if err := c.applyRegisterToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine,
-			engineExitContextJITCallStatusCodeOffset, zeroRegister); err != nil {
-			return err
-		}
+		c.applyRegisterToMemoryInstruction(arm64.AMOVWU, zeroRegister, reservedRegisterForEngine, engineExitContextJITCallStatusCodeOffset)
 	}
 
 	// 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.
-	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine,
-		engineArchContextJITCallReturnAddressOffset, reservedRegisterForTemporary); err != nil {
-		return err
-	}
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine, engineArchContextJITCallReturnAddressOffset, reservedRegisterForTemporary)
 
 	ret := c.newProg()
 	ret.As = obj.ARET
@@ -530,8 +598,8 @@ func (c *arm64Compiler) branchInto(target *wazeroir.BranchTarget) error {
 			targetLabel.initialStack = c.locationStack.clone()
 		}
 
-		jmp := c.emitUnconditionalBRInstruction(obj.TYPE_BRANCH)
-		c.assignBranchTarget(labelKey, jmp)
+		br := c.emitUnconditionalBranchInstruction()
+		c.assignBranchTarget(labelKey, br)
 		return nil
 	}
 }
@@ -554,8 +622,229 @@ func (c *arm64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
+// compileCall implements compiler.compileCall for the arm64 architecture.
 func (c *arm64Compiler) compileCall(o *wazeroir.OperationCall) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	target := c.f.ModuleInstance.Functions[o.FunctionIndex]
+
+	if err := c.callFunction(target.Address, target.FunctionType.Type); err != nil {
+		return err
+	}
+	return nil
+}
+
+// compileCall implements compiler.compileCall and compiler.compileCallIndirect (TODO) for the arm64 architecture.
+func (c *arm64Compiler) callFunction(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.
+	if err := c.releaseAllRegistersToStack(); err != nil {
+		return err
+	}
+
+	// Obtain the free registers to be used in the followings.
+	freeRegisters, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 5)
+	if !found {
+		return fmt.Errorf("BUG: all registers except addrReg should be free at this point")
+	}
+	c.locationStack.markRegisterUsed(freeRegisters...)
+
+	// Alias for readability.
+	callFrameStackTopAddressRegister, compiledFunctionAddressRegister, oldStackBasePointer,
+		tmp := freeRegisters[0], freeRegisters[1], freeRegisters[2], freeRegisters[3]
+
+	// TODO: Check the callframe stack length, and if necessary, grow the call frame stack before jump into the target.
+
+	// "tmp = engine.callFrameStackPointer"
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset,
+		tmp)
+	// "callFrameStackTopAddressRegister = &engine.callFrameStack[0]"
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackElement0AddressOffset,
+		callFrameStackTopAddressRegister)
+	// "callFrameStackTopAddressRegister += tmp << $callFrameDataSizeMostSignificantSetBit"
+	c.emitAddInstructionWithLeftShiftedRegister(
+		tmp, callFrameDataSizeMostSignificantSetBit,
+		callFrameStackTopAddressRegister,
+		callFrameStackTopAddressRegister,
+	)
+
+	// At this point, we have:
+	//
+	//    [..., ra.current, rb.current, rc.current, _, ra.next, rb.next, rc.next, ...]  <- call frame stack's data region (somewhere in the memory)
+	//                                               ^
+	//                              callFrameStackTopAddressRegister
+	//                (absolute address of &callFrame[engine.callFrameStackPointer]])
+	//
+	// where:
+	//      ra.* = callFrame.returnAddress
+	//      rb.* = callFrame.returnStackBasePointer
+	//      rc.* = callFrame.compiledFunction
+	//      _  = callFrame's padding (see comment on callFrame._ field.)
+	//
+	// In the following comment, we use the notations in the above example.
+	//
+	// What we have to do in the following is that
+	//   1) Set rb.current so that we can return back to this function properly.
+	//   2) Set engine.valueStackContext.stackBasePointer for the next function.
+	//   3) Set rc.next to specify which function is executed on the current call frame (needs to make Go function calls).
+	//   4) Set ra.current so that we can return back to this function properly.
+
+	// 1) Set rb.current so that we can return back to this function properly.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset,
+		oldStackBasePointer)
+	c.applyRegisterToMemoryInstruction(arm64.AMOVD,
+		oldStackBasePointer,
+		// "rb.current" is BELOW the top address. See the above example for detail.
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerOffset))
+
+	// 2) Set engine.valueStackContext.stackBasePointer for the next function.
+	//
+	// At this point, oldStackBasePointer holds the old stack base pointer. We could get the new frame's
+	// stack base pointer by "old stack base pointer + old stack pointer - # of function params"
+	// See the comments in engine.pushCallFrame which does exactly the same calculation in Go.
+	c.applyConstToRegisterInstruction(arm64.AADD,
+		int64(c.locationStack.sp)-int64(len(functype.Params)),
+		oldStackBasePointer)
+	c.applyRegisterToMemoryInstruction(arm64.AMOVD,
+		oldStackBasePointer,
+		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset)
+
+	// 3) Set rc.next to specify which function is executed on the current call frame.
+	//
+	// First, we read the address of the first item of engine.compiledFunctions slice (= &engine.compiledFunctions[0])
+	// into tmp.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCompiledFunctionsElement0AddressOffset,
+		tmp)
+
+	// Next, read the address of the target function (= &engine.compiledFunctions[offset])
+	// into compiledFunctionAddressRegister.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		tmp, int64(addr)*8, // * 8 because the size of *compiledFunction equals 8 bytes.
+		compiledFunctionAddressRegister)
+
+	// Finally, we are ready to write the address of the target function's *compiledFunction into the new callframe.
+	c.applyRegisterToMemoryInstruction(arm64.AMOVD,
+		compiledFunctionAddressRegister,
+		callFrameStackTopAddressRegister, callFrameCompiledFunctionOffset)
+
+	// 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)
+	// Then write the address into the callframe.
+	c.applyRegisterToMemoryInstruction(arm64.AMOVD,
+		tmp,
+		// "ra.current" is BELOW the top address. See the above example for detail.
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnAddressOffset),
+	)
+
+	// Everthing is done to make function call now.
+	// We increment the callframe stack pointer.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset,
+		tmp)
+	c.applyConstToRegisterInstruction(arm64.AADD, 1, tmp)
+	c.applyRegisterToMemoryInstruction(arm64.AMOVD,
+		tmp,
+		reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset)
+
+	// Then, br into the target function's initial address.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+		compiledFunctionAddressRegister, compiledFunctionCodeInitialAddressOffset,
+		tmp)
+	c.emitUnconditionalBranchToAddressOnRegister(tmp)
+
+	// All the registers used are temporary so we mark them unused.
+	c.markRegisterUnused(freeRegisters...)
+
+	// We consumed the function parameters from the stack after call.
+	for i := 0; i < len(functype.Params); i++ {
+		c.locationStack.pop()
+	}
+
+	// Also, the function results were pushed by the call.
+	for _, t := range functype.Results {
+		loc := c.locationStack.pushValueLocationOnStack()
+		switch t {
+		case wasm.ValueTypeI32, wasm.ValueTypeI64:
+			loc.setRegisterType(generalPurposeRegisterTypeInt)
+		case wasm.ValueTypeF32, wasm.ValueTypeF64:
+			loc.setRegisterType(generalPurposeRegisterTypeFloat)
+		}
+	}
+
+	// On the function return, we initialize the state for this function.
+	c.initializeReservedStackBasePointerRegister()
+
+	// TODO: initialize module context, and memory pointer.
+	return nil
+}
+
+// readInstructionAddress 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) {
+	// 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
+	// "reading the current program counter" = "reading the absolute address of this ADR instruction".
+	// And then, after compilation phase, we directly edit the native code slice so that
+	// it can properly read the target instruction's absolute address.
+	readAddress := c.newProg()
+	readAddress.As = arm64.AADR
+	readAddress.From.Type = obj.TYPE_BRANCH
+	readAddress.To.Type = obj.TYPE_REG
+	readAddress.To.Reg = destinationRegister
+	c.addInstruction(readAddress)
+
+	// Setup the callback to modify the instruction bytes after compilation.
+	// Note: this is the closure over readAddress (*obj.Prog).
+	c.afterAssembleCallback = append(c.afterAssembleCallback, func(code []byte) error {
+		// Find the target instruction.
+		target := readAddress
+		for target != nil {
+			if target.As == beforeTargetInst {
+				// At this point, target is the instruction right before the target instruction.
+				// Thus, advance one more time to make target the target instruction.
+				target = target.Link
+				break
+			}
+			target = target.Link
+		}
+
+		if target == nil {
+			return fmt.Errorf("BUG: target instruction not found for read instruction address")
+		}
+
+		offset := target.Pc - readAddress.Pc
+		if offset > math.MaxUint8 {
+			// We could support up to 20-bit integer, but byte should be enough for our impl.
+			// If the necessity comes up, we could fix the below to support larger offsets.
+			return fmt.Errorf("BUG: too large offset for read")
+		}
+
+		// Now ready to write an offset byte.
+		v := byte(offset)
+		// arm64 has 4-bytes = 32-bit fixed-length instruction.
+		adrInstructionBytes := code[readAddress.Pc : readAddress.Pc+4]
+		// According to the binary format of ADR instruction in arm64:
+		// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ADR--Form-PC-relative-address-?lang=en
+		//
+		// The 0 to 1 bits live on 29 to 30 bits of the instruction.
+		adrInstructionBytes[3] |= (v & 0b00000011) << 5
+		// The 2 to 4 bits live on 5 to 7 bits of the instruction.
+		adrInstructionBytes[0] |= (v & 0b00011100) << 3
+		// The 5 to 7 bits live on 8 to 10 bits of the instruction.
+		adrInstructionBytes[1] |= (v & 0b11100000) >> 5
+		return nil
+	})
 }
 
 func (c *arm64Compiler) compileCallIndirect(o *wazeroir.OperationCallIndirect) error {
@@ -1398,12 +1687,18 @@ func (c *arm64Compiler) pushZeroValue() {
 // popTwoValuesOnRegisters pops two values from the location stacks, ensures
 // these two values are located on registers, and mark them unused.
 func (c *arm64Compiler) popTwoValuesOnRegisters() (x1, x2 *valueLocation, err error) {
-	x2, err = c.popValueOnRegister()
-	if err != nil {
+	x2 = c.locationStack.pop()
+	if err = c.ensureOnGeneralPurposeRegister(x2); err != nil {
 		return
 	}
 
-	x1, err = c.popValueOnRegister()
+	x1 = c.locationStack.pop()
+	if err = c.ensureOnGeneralPurposeRegister(x1); err != nil {
+		return
+	}
+
+	c.markRegisterUnused(x2.register)
+	c.markRegisterUnused(x1.register)
 	return
 }
 
@@ -1478,15 +1773,7 @@ func (c *arm64Compiler) loadValueOnStackToRegister(loc *valueLocation) (err erro
 		return
 	}
 
-	if offset := int64(loc.stackPointer) * 8; offset > math.MaxInt16 {
-		// The assembler can take care of offsets larger than 2^15-1 by emitting additional instructions to load such large offset,
-		// but it uses "its" temporary register which we cannot track. Therefore, we avoid directly emitting memory load with large offsets,
-		// but instead load the constant manually to "our" temporary register, then emit the load with it.
-		c.applyConstToRegisterInstruction(arm64.AMOVD, offset, reservedRegisterForTemporary)
-		c.applyRegisterOffsetMemoryToRegisterInstruction(inst, reservedRegisterForStackBasePointerAddress, reservedRegisterForTemporary, reg)
-	} else {
-		c.applyMemoryToRegisterInstruction(inst, reservedRegisterForStackBasePointerAddress, offset, reg)
-	}
+	c.applyMemoryToRegisterInstruction(inst, reservedRegisterForStackBasePointerAddress, int64(loc.stackPointer)*8, reg)
 
 	// Record that the value holds the register and the register is marked used.
 	loc.setRegister(reg)
@@ -1545,17 +1832,7 @@ func (c *arm64Compiler) releaseRegisterToStack(loc *valueLocation) (err error) {
 		inst = arm64.AFMOVD
 	}
 
-	if offset := int64(loc.stackPointer) * 8; offset > math.MaxInt16 {
-		// The assembler can take care of offsets larger than 2^15-1 by emitting additional instructions to load such large offset,
-		// but it uses "its" temporary register which we cannot track. Therefore, we avoid directly emitting memory load with large offsets,
-		// but instead load the constant manually to "our" temporary register, then emit the load with it.
-		c.applyConstToRegisterInstruction(arm64.AMOVD, offset, reservedRegisterForTemporary)
-		c.applyRegisterToRegisterOffsetMemoryInstruction(inst, reservedRegisterForStackBasePointerAddress, reservedRegisterForTemporary, loc.register)
-	} else {
-		if err = c.applyRegisterToMemoryInstruction(inst, reservedRegisterForStackBasePointerAddress, offset, loc.register); err != nil {
-			return
-		}
-	}
+	c.applyRegisterToMemoryInstruction(inst, loc.register, reservedRegisterForStackBasePointerAddress, int64(loc.stackPointer)*8)
 
 	// Mark the register is free.
 	c.locationStack.releaseRegister(loc)
@@ -1566,37 +1843,33 @@ func (c *arm64Compiler) releaseRegisterToStack(loc *valueLocation) (err error) {
 // so that it points to the absolute address of the stack base for this function.
 func (c *arm64Compiler) initializeReservedStackBasePointerRegister() error {
 	// First, load the address of the first element in the value stack into reservedRegisterForStackBasePointerAddress temporarily.
-	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
 		reservedRegisterForEngine, engineGlobalContextValueStackElement0AddressOffset,
-		reservedRegisterForStackBasePointerAddress); err != nil {
-		return err
-	}
+		reservedRegisterForStackBasePointerAddress)
 
-	// Next we move the base pointer (engine.stackBasePointer) to the tmp register.
-	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD,
+	// Next we move the base pointer (engine.stackBasePointer) to reservedRegisterForTemporary.
+	c.applyMemoryToRegisterInstruction(arm64.AMOVD,
 		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset,
-		reservedRegisterForTemporary); err != nil {
-		return err
-	}
-
-	// 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, reservedRegisterForTemporary, 3)
-	c.addInstruction(calcStackBasePointerAddress)
+		reservedRegisterForTemporary)
+
+	// Finally, we calculate "reservedRegisterForStackBasePointerAddress + reservedRegisterForTemporary << 3"
+	// where we shift tmpReg by 3 because stack pointer is an index in the []uint64
+	// so we must multiply the value by the size of uint64 = 8 bytes.
+	c.emitAddInstructionWithLeftShiftedRegister(
+		reservedRegisterForTemporary, 3, reservedRegisterForStackBasePointerAddress,
+		reservedRegisterForStackBasePointerAddress)
 	return nil
 }
 
-// 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) {
+// emitAddInstructionWithLeftShiftedRegister emits an ADD instruction to perform "destinationReg = srcReg + (shiftedSourceReg << shiftNum)".
+func (c *arm64Compiler) emitAddInstructionWithLeftShiftedRegister(shiftedSourceReg int16, shiftNum int64, srcReg, destinationReg int16) {
+	inst := c.newProg()
+	inst.As = arm64.AADD
+	inst.To.Type = obj.TYPE_REG
+	inst.To.Reg = destinationReg
+	// See https://github.com/twitchyliquid64/golang-asm/blob/v0.15.1/obj/link.go#L120-L131
 	inst.From.Type = obj.TYPE_SHIFT
-	inst.From.Offset = (int64(register)&31)<<16 | 0<<22 | (shiftNum&63)<<10
+	inst.From.Offset = (int64(shiftedSourceReg)&31)<<16 | 0<<22 | (shiftNum&63)<<10
+	inst.Reg = srcReg
+	c.addInstruction(inst)
 }
diff --git a/wasm/module.go b/wasm/module.go
index 6009dfe3..2965fe93 100644
--- a/wasm/module.go
+++ b/wasm/module.go
@@ -1,6 +1,8 @@
 package wasm
 
-// DecodeModule parses the configured source into a wasm.Module. This function returns when the source is exhausted or
+import "fmt"
+
+// DecodeModule parses the configured source into a Module. This function returns when the source is exhausted or
 // an error occurs. The result can be initialized for use via Store.Instantiate.
 //
 // Here's a description of the return values:
@@ -90,6 +92,7 @@ type Module struct {
 	// (Store.Instantiate).
 	//
 	// Note: there are no unique constraints relating to the two-level namespace of Import.Module and Import.Name.
+	//
 	// Note: In the Binary Format, this is SectionIDImport.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#import-section%E2%91%A0
@@ -103,6 +106,7 @@ type Module struct {
 	//
 	// Note: FunctionSection is index correlated with the CodeSection. If given the same position, ex. 2, a function
 	// type is at TypeSection[FunctionSection[2]], while its locals and body are at CodeSection[2].
+	//
 	// Note: In the Binary Format, this is SectionIDFunction.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#function-section%E2%91%A0
@@ -115,7 +119,8 @@ type Module struct {
 	// this module at TableSection[0].
 	//
 	// Note: Version 1.0 (MVP) of the WebAssembly spec allows at most one table definition per module, so the length of
-	// the TableSection can be zero or one.
+	// the TableSection can be zero or one, and can only be one if there is no ImportKindTable.
+	//
 	// Note: In the Binary Format, this is SectionIDTable.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#table-section%E2%91%A0
@@ -128,7 +133,8 @@ type Module struct {
 	// this module at TableSection[0].
 	//
 	// Note: Version 1.0 (MVP) of the WebAssembly spec allows at most one memory definition per module, so the length of
-	// the MemorySection can be zero or one.
+	// the MemorySection can be zero or one, and can only be one if there is no ImportKindMemory.
+	//
 	// Note: In the Binary Format, this is SectionIDMemory.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#memory-section%E2%91%A0
@@ -147,6 +153,8 @@ type Module struct {
 
 	// ExportSection contains each export defined in this module.
 	//
+	// Note: In the Binary Format, this is SectionIDExport.
+	//
 	// See https://www.w3.org/TR/wasm-core-1/#exports%E2%91%A0
 	ExportSection map[string]*Export
 
@@ -154,6 +162,7 @@ type Module struct {
 	//
 	// Note: The index here is not the position in the FunctionSection, rather in the function index namespace, which
 	// begins with imported functions.
+	//
 	// Note: In the Binary Format, this is SectionIDStart.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#start-section%E2%91%A0
@@ -174,20 +183,12 @@ type Module struct {
 
 	// NameSection is set when the SectionIDCustom "name" was successfully decoded from the binary format.
 	//
-	// Note: This is the only SectionIDCustom defined in the WebAssembly 1.0 (MVP) Binary Format. Others are in
-	// CustomSections
+	// Note: This is the only SectionIDCustom defined in the WebAssembly 1.0 (MVP) Binary Format.
+	// Others are skipped as they are not used in wazero.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#name-section%E2%91%A0
-	NameSection *NameSection
-
-	// CustomSections is set when at least one non-standard, or otherwise unsupported custom section was found in the
-	// binary format.
-	//
-	// Note: This never contains a "name" because that is standard and parsed into the NameSection.
-	// Note: In the Binary Format, this is SectionIDCode.
-	//
 	// See https://www.w3.org/TR/wasm-core-1/#custom-section%E2%91%A0
-	CustomSections map[string][]byte
+	NameSection *NameSection
 }
 
 // Index is the offset in an index namespace, not necessarily an absolute position in a Module section. This is because
@@ -226,7 +227,7 @@ const (
 	ValueTypeF64 ValueType = 0x7c
 )
 
-// ValuTypeName returns the type name of the given ValueType as a string.
+// ValueTypeName 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 {
@@ -461,3 +462,46 @@ func (m *Module) allDeclarations() (functions []Index, globals []*GlobalType, me
 	tables = append(tables, m.TableSection...)
 	return
 }
+
+// SectionElementCount returns the count of elements in a given section ID
+//
+// For example...
+// * SectionIDType returns the count of FunctionType
+// * SectionIDCustom returns one if the NameSection is present
+// * SectionIDExport returns the count of unique export names
+func (m *Module) SectionElementCount(sectionID SectionID) uint32 { // element as in vector elements!
+	switch sectionID {
+	case SectionIDCustom:
+		if m.NameSection != nil {
+			return 1
+		}
+		return 0
+	case SectionIDType:
+		return uint32(len(m.TypeSection))
+	case SectionIDImport:
+		return uint32(len(m.ImportSection))
+	case SectionIDFunction:
+		return uint32(len(m.FunctionSection))
+	case SectionIDTable:
+		return uint32(len(m.TableSection))
+	case SectionIDMemory:
+		return uint32(len(m.MemorySection))
+	case SectionIDGlobal:
+		return uint32(len(m.GlobalSection))
+	case SectionIDExport:
+		return uint32(len(m.ExportSection))
+	case SectionIDStart:
+		if m.StartSection != nil {
+			return 1
+		}
+		return 0
+	case SectionIDElement:
+		return uint32(len(m.ElementSection))
+	case SectionIDCode:
+		return uint32(len(m.CodeSection))
+	case SectionIDData:
+		return uint32(len(m.DataSection))
+	default:
+		panic(fmt.Errorf("BUG: unknown section: %d", sectionID))
+	}
+}
diff --git a/wasm/store.go b/wasm/store.go
index d23a7cde..3771303c 100644
--- a/wasm/store.go
+++ b/wasm/store.go
@@ -288,17 +288,21 @@ func (s *Store) Instantiate(module *Module, name string) error {
 
 	for i, f := range instance.Functions {
 		if err := s.engine.Compile(f); err != nil {
-			return fmt.Errorf("compilation failed at index %d/%d: %v", i, len(module.FunctionSection)-1, err)
+			idx := module.SectionElementCount(SectionIDFunction) - 1
+			return fmt.Errorf("compilation failed at index %d/%d: %v", i, idx, err)
 		}
 	}
 
 	// Check the start function is valid.
-	if startIndex := module.StartSection; startIndex != nil {
-		index := *startIndex
-		if int(index) >= len(instance.Functions) {
-			return fmt.Errorf("invalid start function index: %d", index)
+	// TODO: this should be verified during decode so that errors have the correct source positions
+	var startFunction *FunctionInstance
+	if module.StartSection != nil {
+		startIndex := *module.StartSection
+		if startIndex >= uint32(len(instance.Functions)) {
+			return fmt.Errorf("invalid start function index: %d", startIndex)
 		}
-		ft := instance.Functions[index].FunctionType
+		startFunction = instance.Functions[startIndex]
+		ft := startFunction.FunctionType
 		if len(ft.Type.Params) != 0 || len(ft.Type.Results) != 0 {
 			return fmt.Errorf("start function must have the empty function type")
 		}
@@ -308,10 +312,8 @@ func (s *Store) Instantiate(module *Module, name string) error {
 	rollbackFuncs = nil
 
 	// Execute the start function.
-	ctx := context.Background()
-	if startIndex := module.StartSection; startIndex != nil {
-		f := instance.Functions[*startIndex]
-		if _, err := s.engine.Call(ctx, f); err != nil {
+	if startFunction != nil {
+		if _, err = s.engine.Call(context.Background(), startFunction); err != nil {
 			return fmt.Errorf("calling start function failed: %v", err)
 		}
 	}
@@ -596,9 +598,9 @@ func (s *Store) buildFunctionInstances(module *Module, target *ModuleInstance) (
 	n, nLen := 0, len(functionNames)
 
 	for codeIndex, typeIndex := range module.FunctionSection {
-		if typeIndex >= uint32(len(module.TypeSection)) {
+		if typeIndex >= module.SectionElementCount(SectionIDType) {
 			return rollbackFuncs, fmt.Errorf("function type index out of range")
-		} else if codeIndex >= len(module.CodeSection) {
+		} else if uint32(codeIndex) >= module.SectionElementCount(SectionIDCode) {
 			return rollbackFuncs, fmt.Errorf("code index out of range")
 		}
 
@@ -629,7 +631,8 @@ func (s *Store) buildFunctionInstances(module *Module, target *ModuleInstance) (
 		}
 
 		if err := validateFunctionInstance(f, funcs, globals, mems, tables, module.TypeSection, maximumValuesOnStack); err != nil {
-			return rollbackFuncs, fmt.Errorf("invalid function '%s' (%d/%d): %v", f.Name, codeIndex, len(module.FunctionSection)-1, err)
+			idx := module.SectionElementCount(SectionIDFunction) - 1
+			return rollbackFuncs, fmt.Errorf("invalid function '%s' (%d/%d): %v", f.Name, codeIndex, idx, err)
 		}
 
 		err = s.addFunctionInstance(f)
@@ -680,8 +683,8 @@ func (s *Store) buildMemoryInstances(module *Module, target *ModuleInstance) (ro
 		}
 
 		size := uint64(offset) + uint64(len(d.Init))
-		maxPage := uint32(MemoryMaxPages)
-		if int(d.MemoryIndex) < len(module.MemorySection) && module.MemorySection[d.MemoryIndex].Max != nil {
+		maxPage := MemoryMaxPages
+		if d.MemoryIndex < module.SectionElementCount(SectionIDMemory) && module.MemorySection[d.MemoryIndex].Max != nil {
 			maxPage = *module.MemorySection[d.MemoryIndex].Max
 		}
 		if size > memoryPagesToBytesNum(maxPage) {
@@ -734,7 +737,7 @@ func (s *Store) buildTableInstances(module *Module, target *ModuleInstance) (rol
 		size := offset + len(elem.Init)
 
 		max := uint32(math.MaxUint32)
-		if int(elem.TableIndex) < len(module.TableSection) && module.TableSection[elem.TableIndex].Limit.Max != nil {
+		if elem.TableIndex < module.SectionElementCount(SectionIDTable) && module.TableSection[elem.TableIndex].Limit.Max != nil {
 			max = *module.TableSection[elem.TableIndex].Limit.Max
 		}
 
@@ -772,7 +775,7 @@ func (s *Store) buildTableInstances(module *Module, target *ModuleInstance) (rol
 }
 
 func (s *Store) buildExportInstances(module *Module, target *ModuleInstance) (rollbackFuncs []func(), err error) {
-	target.Exports = make(map[string]*ExportInstance, len(module.ExportSection))
+	target.Exports = make(map[string]*ExportInstance, module.SectionElementCount(SectionIDExport))
 	for name, exp := range module.ExportSection {
 		index := exp.Index
 		var ei *ExportInstance
diff --git a/wasm/text/decoder.go b/wasm/text/decoder.go
index 521f0fff..1d5f89d7 100644
--- a/wasm/text/decoder.go
+++ b/wasm/text/decoder.go
@@ -18,17 +18,32 @@ const (
 	positionParam
 	positionResult
 	positionModule
-	positionModuleType
-	positionModuleImport
-	positionModuleImportFunc
-	positionModuleFunc
-	positionModuleMemory
-	positionModuleExport
-	positionModuleExportFunc
-	positionModuleExportMemory
-	positionModuleStart
+	positionImport
+	positionImportFunc
+	positionMemory
+	positionExport
+	positionExportFunc
+	positionExportMemory
+	positionStart
 )
 
+type callbackPosition byte
+
+const (
+	// callbackPositionUnhandledToken is set on a token besides a paren.
+	callbackPositionUnhandledToken callbackPosition = iota
+	// callbackPositionUnhandledField is at the field name (tokenKeyword) which isn't "type", "param" or "result"
+	callbackPositionUnhandledField
+	// callbackPositionEndField is at the end (tokenRParen) of the field enclosing the type use.
+	callbackPositionEndField
+)
+
+// moduleParser parses a single wasm.Module from WebAssembly 1.0 (MVP) Text format.
+//
+// Note: The indexNamespace of wasm.SectionIDMemory and wasm.SectionIDTable allow up-to-one item. For example, you
+// cannot define both one import and one module-defined memory, rather one or the other (or none). Even if these rules
+// are also enforced in module instantiation, they are also enforced here, to allow relevant source line/col in errors.
+// See https://www.w3.org/TR/wasm-core-1/#modules%E2%91%A3
 type moduleParser struct {
 	// source is the entire WebAssembly text format source code being parsed.
 	source []byte
@@ -153,27 +168,38 @@ func (p *moduleParser) beginModuleField(tok tokenType, tokenBytes []byte, _, _ u
 	if tok == tokenKeyword {
 		switch string(tokenBytes) {
 		case "type":
-			p.pos = positionModuleType
+			p.pos = positionType
 			return p.typeParser.begin, nil
 		case "import":
-			p.pos = positionModuleImport
+			p.pos = positionImport
 			return p.parseImportModule, nil
 		case "func":
-			p.pos = positionModuleFunc
-			p.funcParser.currentIdx = wasm.Index(len(p.module.FunctionSection))
+			p.pos = positionFunc
+			p.funcParser.currentIdx = p.module.SectionElementCount(wasm.SectionIDFunction)
 			return p.funcParser.begin, nil
+		case "table":
+			return nil, errors.New("TODO: table")
 		case "memory":
-			p.pos = positionModuleMemory
+			if p.memoryNamespace.count > 0 {
+				return nil, moreThanOneInvalidInSection(wasm.SectionIDMemory)
+			}
+			p.pos = positionMemory
 			return p.memoryParser.begin, nil
+		case "global":
+			return nil, errors.New("TODO: global")
 		case "export":
-			p.pos = positionModuleExport
+			p.pos = positionExport
 			return p.parseExportName, nil
 		case "start":
-			if p.module.StartSection != nil {
-				return nil, errors.New("redundant start")
+			if p.module.SectionElementCount(wasm.SectionIDStart) > 0 {
+				return nil, moreThanOneInvalid("start")
 			}
-			p.pos = positionModuleStart
+			p.pos = positionStart
 			return p.parseStart, nil
+		case "elem":
+			return nil, errors.New("TODO: elem")
+		case "data":
+			return nil, errors.New("TODO: data")
 		default:
 			return nil, unexpectedFieldName(tokenBytes)
 		}
@@ -243,9 +269,9 @@ func (p *moduleParser) parseImportModule(tok tokenType, tokenBytes []byte, _, _
 // parseImportName returns parseImport after recording the import name, or errs if it couldn't be read.
 //
 // Ex. Import name is present `(import "Math" "PI" (func (result f32)))`
-//                                         starts here --^    ^
-//                                           records PI --^   |
-//                                 parseImport resumes here --+
+//                              starts here --^^   ^
+//                                records PI --+   |
+//                      parseImport resumes here --+
 //
 // Ex. Imported function name is absent `(import "Math" (func (result f32)))`
 //                                          errs here --+
@@ -284,7 +310,10 @@ func (p *moduleParser) beginImportDesc(tok tokenType, tokenBytes []byte, _, _ ui
 
 	switch string(tokenBytes) {
 	case "func":
-		p.pos = positionModuleImportFunc
+		if p.module.SectionElementCount(wasm.SectionIDFunction) > 0 {
+			return nil, importAfterModuleDefined(wasm.SectionIDFunction)
+		}
+		p.pos = positionImportFunc
 		return p.parseImportFuncID, nil
 	case "table", "memory", "global":
 		return nil, fmt.Errorf("TODO: %s", tokenBytes)
@@ -303,17 +332,18 @@ 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
-		id, err := p.funcNamespace.setID(tokenBytes)
-		if err != nil {
+		if name, err := p.funcNamespace.setID(tokenBytes); err != nil {
 			return nil, err
+		} else {
+			p.addFunctionName(name)
 		}
-		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
+// addFunctionName appends the current imported or module-defined function name to the wasm.NameSection iff it is not
+// empty.
 func (p *moduleParser) addFunctionName(name string) {
 	if name == "" {
 		return // there's no value in an empty name
@@ -331,7 +361,7 @@ func (p *moduleParser) addFunctionName(name string) {
 // Ex. If there is no signature `(import "" "main" (func))`
 //                     calls onImportFunc here ---^
 func (p *moduleParser) parseImportFunc(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	idx := wasm.Index(len(p.module.ImportSection))
+	idx := p.module.SectionElementCount(wasm.SectionIDImport)
 	if tok == tokenID { // Ex. (func $main $main)
 		return nil, fmt.Errorf("redundant ID %s", tokenBytes)
 	}
@@ -342,7 +372,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, _, _ uint32) (tokenParser, error) {
+func (p *moduleParser) onImportFunc(typeIdx wasm.Index, paramNames wasm.NameMap, pos callbackPosition, tok tokenType, tokenBytes []byte, _, _ uint32) (tokenParser, error) {
 	i := p.currentModuleField.(*wasm.Import)
 	i.Kind = wasm.ImportKindFunc
 	i.DescFunc = typeIdx
@@ -351,12 +381,12 @@ func (p *moduleParser) onImportFunc(typeIdx wasm.Index, paramNames wasm.NameMap,
 	p.funcNamespace.count++
 
 	switch pos {
-	case onTypeUseUnhandledToken:
+	case callbackPositionUnhandledToken:
 		return nil, unexpectedToken(tok, tokenBytes)
-	case onTypeUseUnhandledField:
+	case callbackPositionUnhandledField:
 		return nil, unexpectedFieldName(tokenBytes)
-	case onTypeUseEndField:
-		p.pos = positionModuleImport
+	case callbackPositionEndField:
+		p.pos = positionImport
 		return p.parseImportEnd, nil
 	}
 	return p.parseImportFuncEnd, nil
@@ -368,7 +398,7 @@ func (p *moduleParser) parseImportFuncEnd(tok tokenType, tokenBytes []byte, _, _
 		return nil, unexpectedToken(tok, tokenBytes)
 	}
 
-	p.pos = positionModuleImport
+	p.pos = positionImport
 	return p.parseImportEnd, nil
 }
 
@@ -431,7 +461,7 @@ func (p *moduleParser) parseExportName(tok tokenType, tokenBytes []byte, _, _ ui
 	case tokenString: // Ex. "" or "PI"
 		name := string(tokenBytes[1 : len(tokenBytes)-1]) // strip quotes
 		if _, ok := p.module.ExportSection[name]; ok {
-			return nil, fmt.Errorf("duplicate name %q", name)
+			return nil, fmt.Errorf("%q already exported", name)
 		}
 		p.currentModuleField = &wasm.Export{Name: name}
 		return p.parseExport, nil
@@ -464,10 +494,10 @@ func (p *moduleParser) beginExportDesc(tok tokenType, tokenBytes []byte, _, _ ui
 
 	switch string(tokenBytes) {
 	case "func":
-		p.pos = positionModuleExportFunc
+		p.pos = positionExportFunc
 		return p.parseExportDesc, nil
 	case "memory":
-		p.pos = positionModuleExportMemory
+		p.pos = positionExportMemory
 		return p.parseExportDesc, nil
 	case "table", "global":
 		return nil, fmt.Errorf("TODO: %s", tokenBytes)
@@ -481,16 +511,16 @@ func (p *moduleParser) parseExportDesc(tok tokenType, tokenBytes []byte, line, c
 	var namespace *indexNamespace
 	e := p.currentModuleField.(*wasm.Export)
 	switch p.pos {
-	case positionModuleExportFunc:
+	case positionExportFunc:
 		e.Kind = wasm.ExportKindFunc
 		namespace = p.funcNamespace
-	case positionModuleExportMemory:
+	case positionExportMemory:
 		e.Kind = wasm.ExportKindMemory
 		namespace = p.memoryNamespace
 	default:
 		panic(fmt.Errorf("BUG: unhandled parsing state on parseExportDesc: %v", p.pos))
 	}
-	eIdx := wasm.Index(len(p.module.ExportSection))
+	eIdx := p.module.SectionElementCount(wasm.SectionIDExport)
 	typeIdx, resolved, err := namespace.parseIndex(wasm.SectionIDExport, eIdx, 0, tok, tokenBytes, line, col)
 	if err != nil {
 		return nil, err
@@ -514,7 +544,7 @@ func (p *moduleParser) parseExportDescEnd(tok tokenType, tokenBytes []byte, _, _
 	case tokenUN, tokenID:
 		return nil, errors.New("redundant index")
 	case tokenRParen:
-		p.pos = positionModuleExport
+		p.pos = positionExport
 		return p.parseExportEnd, nil
 	default:
 		return nil, unexpectedToken(tok, tokenBytes)
@@ -612,10 +642,10 @@ func (p *moduleParser) resolveFunctionIndices(module *wasm.Module) error {
 // This errs if any type index is unresolved, out of range or mismatches an inlined type use signature.
 func (p *moduleParser) resolveTypeUses(module *wasm.Module) error {
 	inlinedToRealIdx := p.addInlinedTypes()
-	return p.resolveInlined(module, inlinedToRealIdx)
+	return p.resolveInlinedTypes(module, inlinedToRealIdx)
 }
 
-func (p *moduleParser) resolveInlined(module *wasm.Module, inlinedToRealIdx map[wasm.Index]wasm.Index) error {
+func (p *moduleParser) resolveInlinedTypes(module *wasm.Module, inlinedToRealIdx map[wasm.Index]wasm.Index) error {
 	// Now look for all the uses of the inlined types and apply the mapping above
 	for _, i := range p.typeUseParser.inlinedTypeIndices {
 		switch i.section {
@@ -699,28 +729,28 @@ func (p *moduleParser) errorContext() string {
 		return ""
 	case positionModule:
 		return "module"
-	case positionModuleType:
-		idx := wasm.Index(len(p.module.TypeSection))
+	case positionType:
+		idx := p.module.SectionElementCount(wasm.SectionIDType)
 		return fmt.Sprintf("module.type[%d]%s", idx, p.typeParser.errorContext())
-	case positionModuleImport, positionModuleImportFunc: // TODO: table, memory or global
-		idx := wasm.Index(len(p.module.ImportSection))
-		if p.pos == positionModuleImport {
+	case positionImport, positionImportFunc: // TODO: table, memory or global
+		idx := p.module.SectionElementCount(wasm.SectionIDImport)
+		if p.pos == positionImport {
 			return fmt.Sprintf("module.import[%d]", idx)
 		}
 		return fmt.Sprintf("module.import[%d].func%s", idx, p.typeUseParser.errorContext())
-	case positionModuleFunc:
-		idx := wasm.Index(len(p.module.FunctionSection))
+	case positionFunc:
+		idx := p.module.SectionElementCount(wasm.SectionIDFunction)
 		return fmt.Sprintf("module.func[%d]%s", idx, p.typeUseParser.errorContext())
-	case positionModuleMemory:
-		idx := wasm.Index(len(p.module.MemorySection))
+	case positionMemory:
+		idx := p.module.SectionElementCount(wasm.SectionIDMemory)
 		return fmt.Sprintf("module.memory[%d]", idx)
-	case positionModuleExport, positionModuleExportFunc: // TODO: table, memory or global
-		idx := wasm.Index(len(p.module.ExportSection))
-		if p.pos == positionModuleExport {
+	case positionExport, positionExportFunc: // TODO: table, memory or global
+		idx := p.module.SectionElementCount(wasm.SectionIDExport)
+		if p.pos == positionExport {
 			return fmt.Sprintf("module.export[%d]", idx)
 		}
 		return fmt.Sprintf("module.export[%d].func", idx)
-	case positionModuleStart:
+	case positionStart:
 		return "module.start"
 	default: // parserPosition is an enum, we expect to have handled all cases above. panic if we didn't
 		panic(fmt.Errorf("BUG: unhandled parsing state on errorContext: %v", p.pos))
diff --git a/wasm/text/errors.go b/wasm/text/errors.go
index a5c34df0..987fc807 100644
--- a/wasm/text/errors.go
+++ b/wasm/text/errors.go
@@ -45,6 +45,36 @@ func unexpectedToken(tok tokenType, tokenBytes []byte) error {
 	}
 }
 
+// importAfterModuleDefined is the failure for the condition "all imports must occur before any regular definition",
+// which applies regardless of abbreviation.
+//
+// Ex. Both of these fail because an import can only be declared when SectionIDFunction is empty.
+// `(func) (import "" "" (func))` which is the same as  `(func) (import "" "" (func))`
+//
+// See https://www.w3.org/TR/wasm-core-1/#modules%E2%91%A0%E2%91%A2
+func importAfterModuleDefined(section wasm.SectionID) error {
+	return fmt.Errorf("import after module-defined %s", wasm.SectionIDName(section))
+}
+
+// moreThanOneInvalidInSection allows enforcement of section size limits.
+//
+// Ex. All of these fail because they result in two memories.
+// * `(module (memory 1) (memory 1))`
+// * `(module (memory 1) (import "" "" (memory 1)))`
+//   * Note the latter expands to the same as the former: `(import "" "" (memory 1))`
+// * `(module (import "" "" (memory 1)) (import "" "" (memory 1)))`
+//
+// See https://www.w3.org/TR/wasm-core-1/#tables%E2%91%A0
+// See https://www.w3.org/TR/wasm-core-1/#memories%E2%91%A0
+func moreThanOneInvalidInSection(section wasm.SectionID) error {
+	return moreThanOneInvalid(wasm.SectionIDName(section))
+}
+
+// moreThanOneInvalid is the failure when a declaration that can result in more than one item.
+func moreThanOneInvalid(context string) error {
+	return fmt.Errorf("at most one %s allowed", context)
+}
+
 func unhandledSection(section wasm.SectionID) error {
 	return fmt.Errorf("BUG: unhandled %s", wasm.SectionIDName(section))
 }
diff --git a/wasm/text/func_parser.go b/wasm/text/func_parser.go
index efac06ed..9e907e5c 100644
--- a/wasm/text/func_parser.go
+++ b/wasm/text/func_parser.go
@@ -17,8 +17,8 @@ type onFunc func(typeIdx wasm.Index, code *wasm.Code, name string, localNames wa
 // funcParser parses any instructions and dispatches to onFunc.
 //
 // Ex.  `(module (func (nop)))`
-//       starts here --^    ^
-//      calls onFunc here --+
+//        begin here --^    ^
+//  end calls onFunc here --+
 //
 // Note: funcParser is reusable. The caller resets via begin.
 type funcParser struct {
@@ -73,11 +73,11 @@ func (p *funcParser) begin(tok tokenType, tokenBytes []byte, line, col uint32) (
 // are read. Finally, this finishes via endFunc.
 //
 // Ex. `(module (func $math.pi (result f32))`
-//               starts here --^           ^
+//                begin here --^           ^
 //                  endFunc resumes here --+
 //
 // Ex.    `(module (func $math.pi (result f32) (local i32) )`
-//                  starts here --^            ^           ^
+//                   begin here --^            ^           ^
 //      funcParser.afterTypeUse resumes here --+           |
 //                                  endFunc resumes here --+
 //
@@ -98,11 +98,11 @@ func (p *funcParser) parseFunc(tok tokenType, tokenBytes []byte, line, col uint3
 // Ex. Given the source `(module (func nop))`
 //          afterTypeUse starts here --^  ^
 //                    calls onFunc here --+
-func (p *funcParser) afterTypeUse(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 callbackPosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 	switch pos {
-	case onTypeUseEndField:
+	case callbackPositionEndField:
 		return p.onFunc(typeIdx, codeEnd, p.currentName, paramNames)
-	case onTypeUseUnhandledField:
+	case callbackPositionUnhandledField:
 		return sExpressionsUnsupported(tok, tokenBytes, line, col)
 	}
 
@@ -120,7 +120,7 @@ func sExpressionsUnsupported(tok tokenType, tokenBytes []byte, _, _ uint32) (tok
 	case "param":
 		return nil, errors.New("param after result")
 	case "result":
-		return nil, errors.New("duplicate result")
+		return nil, moreThanOneInvalid("result")
 	case "local":
 		return nil, errors.New("TODO: local")
 	}
diff --git a/wasm/text/index_namespace.go b/wasm/text/index_namespace.go
index 3bc406ff..f6db68d6 100644
--- a/wasm/text/index_namespace.go
+++ b/wasm/text/index_namespace.go
@@ -38,12 +38,22 @@ type indexNamespace struct {
 // setID ensures the given tokenID is unique within this context and raises an error if not. The resulting mapping is
 // stripped of the leading '$' to match other tools, as described in stripDollar.
 func (i *indexNamespace) setID(idToken []byte) (string, error) {
-	id := string(stripDollar(idToken))
-	if _, ok := i.idToIdx[id]; ok {
-		return id, fmt.Errorf("duplicate ID %s", idToken)
+	name, err := i.requireNoID(idToken)
+	if err != nil {
+		return name, err
+	}
+	i.idToIdx[name] = i.count
+	return name, nil
+}
+
+// hasID checks to see if this tokenID is unique within this context and returns an error. The result string is
+// stripped of the leading '$' to match other tools, as described in stripDollar.
+func (i *indexNamespace) requireNoID(idToken []byte) (string, error) {
+	name := string(stripDollar(idToken))
+	if _, ok := i.idToIdx[name]; ok {
+		return name, fmt.Errorf("duplicate ID %s", idToken)
 	}
-	i.idToIdx[id] = i.count
-	return id, nil
+	return name, nil
 }
 
 // parseIndex is a tokenParser called in a field that can only contain a symbolic identifier or raw numeric index.
diff --git a/wasm/text/typeuse_parser.go b/wasm/text/typeuse_parser.go
index 8c2a6838..87eb4dfa 100644
--- a/wasm/text/typeuse_parser.go
+++ b/wasm/text/typeuse_parser.go
@@ -12,32 +12,21 @@ func newTypeUseParser(module *wasm.Module, typeNamespace *indexNamespace) *typeU
 	return &typeUseParser{module: module, typeNamespace: typeNamespace}
 }
 
-type onTypeUsePosition byte
-
-const (
-	// onTypeUseUnhandledToken is set on a token besides a paren.
-	onTypeUseUnhandledToken onTypeUsePosition = iota
-	// onTypeUseUnhandledField is at the field name (tokenKeyword) which isn't "type", "param" or "result"
-	onTypeUseUnhandledField
-	// onTypeUseEndField is at the end (tokenRParen) of the field enclosing the type use.
-	onTypeUseEndField
-)
-
 // onTypeUse is invoked when the grammar "(param)* (result)?" completes.
 //
-// * typeIdx if unresolved, this is replaced in typeUseParser.resolveTypeUses
+// * typeIdx if unresolved, this is replaced in moduleParser.resolveTypeUses
 // * paramNames is nil unless IDs existed on at least one "param" field.
 // * pos is the context used to determine which tokenParser to return
 //
 // Note: this is called when neither a "param" nor a "result" field are parsed, or on any field following a "param"
 // that is not a "result": pos clarifies this.
-type onTypeUse func(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error)
+type onTypeUse func(typeIdx wasm.Index, paramNames wasm.NameMap, pos callbackPosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error)
 
-// typeUseParser parses an inlined type from a field such "func" and calls to onTypeUse or onUnknownField.
+// typeUseParser parses an inlined type from a field such "func" and calls onTypeUse.
 //
-// Ex. `(import "Math" "PI" (func $math.pi (result f32))`
+// Ex. `(import "Math" "PI" (func $math.pi (result f32)))`
 //                           starts here --^           ^
-//                         onTypeUse resumes here --+
+//                            onTypeUse resumes here --+
 //
 // Note: Unlike normal parsers, this is not used for an entire field (enclosed by parens). Rather, this only handles
 // "type", "param" and "result" inner fields in the correct order.
@@ -113,7 +102,7 @@ type typeUseParser struct {
 //                                onTypeUse resumes here --+
 //
 func (p *typeUseParser) begin(section wasm.SectionID, idx wasm.Index, onTypeUse onTypeUse, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	pos := onTypeUseUnhandledToken
+	pos := callbackPositionUnhandledToken
 	p.pos = positionInitial // to ensure errorContext reports properly
 	switch tok {
 	case tokenLParen:
@@ -122,7 +111,7 @@ func (p *typeUseParser) begin(section wasm.SectionID, idx wasm.Index, onTypeUse
 		p.onTypeUse = onTypeUse
 		return p.beginTypeParamOrResult, nil
 	case tokenRParen:
-		pos = onTypeUseEndField
+		pos = callbackPositionEndField
 	}
 	return onTypeUse(p.emptyTypeIndex(section, idx), nil, pos, tok, tokenBytes, line, col)
 }
@@ -222,7 +211,7 @@ func (p *typeUseParser) beginParamOrResult(tok tokenType, tokenBytes []byte, lin
 	case "type":
 		return nil, errors.New("redundant type")
 	default:
-		return p.end(onTypeUseUnhandledField, tok, tokenBytes, line, col)
+		return p.end(callbackPositionUnhandledField, tok, tokenBytes, line, col)
 	}
 }
 
@@ -282,9 +271,10 @@ func (p *typeUseParser) setParamID(idToken []byte) error {
 //                         records i32 --^  ^
 //   parseMoreParamsOrResult resumes here --+
 //
-// Ex. One param type is present `(param i32)`
-//                         records i32 --^  ^
-//   parseMoreParamsOrResult resumes here --+
+// Ex. Multiple param types are present `(param i32 i64)`
+//                                records i32 --^   ^  ^
+//                                    records i32 --+  |
+//              parseMoreParamsOrResult resumes here --+
 //
 // Ex. type is missing `(param)`
 //                errs here --^
@@ -351,9 +341,9 @@ func (p *typeUseParser) parseResult(tok tokenType, tokenBytes []byte, _, _ uint3
 }
 
 func (p *typeUseParser) parseEnd(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	pos := onTypeUseUnhandledToken
+	pos := callbackPositionUnhandledToken
 	if tok == tokenRParen {
-		pos = onTypeUseEndField
+		pos = callbackPositionEndField
 	}
 	return p.end(pos, tok, tokenBytes, line, col)
 }
@@ -379,7 +369,7 @@ type lineCol struct {
 }
 
 // end invokes onTypeUse to continue parsing
-func (p *typeUseParser) end(pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (parser tokenParser, err error) {
+func (p *typeUseParser) end(pos callbackPosition, tok tokenType, tokenBytes []byte, line, col uint32) (parser tokenParser, err error) {
 	// Record the potentially inlined type if needed and invoke onTypeUse with the parsed index
 	var typeIdx wasm.Index
 	if p.parsedTypeField {


Test patch:
diff --git a/wasm/binary/decoder_test.go b/wasm/binary/decoder_test.go
index c7ab44ec..9ab47ce8 100644
--- a/wasm/binary/decoder_test.go
+++ b/wasm/binary/decoder_test.go
@@ -26,21 +26,6 @@ func TestDecodeModule(t *testing.T) {
 			name:  "only name section",
 			input: &wasm.Module{NameSection: &wasm.NameSection{ModuleName: "simple"}},
 		},
-		{
-			name: "only custom section",
-			input: &wasm.Module{CustomSections: map[string][]byte{
-				"meme": {1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
-			}},
-		},
-		{
-			name: "name section and a custom section",
-			input: &wasm.Module{
-				NameSection: &wasm.NameSection{ModuleName: "simple"},
-				CustomSections: map[string][]byte{
-					"meme": {1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
-				},
-			},
-		},
 		{
 			name: "type section",
 			input: &wasm.Module{
@@ -107,6 +92,29 @@ func TestDecodeModule(t *testing.T) {
 			require.Equal(t, tc.input, m)
 		})
 	}
+	t.Run("skips custom section", func(t *testing.T) {
+		input := append(append(magic, version...),
+			wasm.SectionIDCustom, 0xf, // 15 bytes in this section
+			0x04, 'm', 'e', 'm', 'e',
+			1, 2, 3, 4, 5, 6, 7, 8, 9, 0)
+		m, e := DecodeModule(input)
+		require.NoError(t, e)
+		require.Equal(t, &wasm.Module{}, m)
+	})
+	t.Run("skips custom section, but not name", func(t *testing.T) {
+		input := append(append(magic, version...),
+			wasm.SectionIDCustom, 0xf, // 15 bytes in this section
+			0x04, 'm', 'e', 'm', 'e',
+			1, 2, 3, 4, 5, 6, 7, 8, 9, 0,
+			wasm.SectionIDCustom, 0x0e, // 14 bytes in this section
+			0x04, 'n', 'a', 'm', 'e',
+			subsectionIDModuleName, 0x07, // 7 bytes in this subsection
+			0x06, // the Module name simple is 6 bytes long
+			's', 'i', 'm', 'p', 'l', 'e')
+		m, e := DecodeModule(input)
+		require.NoError(t, e)
+		require.Equal(t, &wasm.Module{NameSection: &wasm.NameSection{ModuleName: "simple"}}, m)
+	})
 }
 
 func TestDecodeModule_Errors(t *testing.T) {
@@ -125,27 +133,16 @@ func TestDecodeModule_Errors(t *testing.T) {
 			input:       []byte("\x00asm\x01\x00\x00\x01"),
 			expectedErr: "invalid version header",
 		},
-		{
-			name: "redundant custom section",
-			input: append(append(magic, version...),
-				wasm.SectionIDCustom, 0x09, // 9 bytes in this section
-				0x04, 'm', 'e', 'm', 'e',
-				subsectionIDModuleName, 0x03, 0x01, 'x',
-				wasm.SectionIDCustom, 0x09, // 9 bytes in this section
-				0x04, 'm', 'e', 'm', 'e',
-				subsectionIDModuleName, 0x03, 0x01, 'y'),
-			expectedErr: "section ID 0: redundant custom section meme",
-		},
 		{
 			name: "redundant name section",
 			input: append(append(magic, version...),
 				wasm.SectionIDCustom, 0x09, // 9 bytes in this section
 				0x04, 'n', 'a', 'm', 'e',
-				subsectionIDModuleName, 0x03, 0x01, 'x',
+				subsectionIDModuleName, 0x02, 0x01, 'x',
 				wasm.SectionIDCustom, 0x09, // 9 bytes in this section
 				0x04, 'n', 'a', 'm', 'e',
-				subsectionIDModuleName, 0x03, 0x01, 'x'),
-			expectedErr: "section ID 0: redundant custom section name",
+				subsectionIDModuleName, 0x02, 0x01, 'x'),
+			expectedErr: "section custom: redundant custom section name",
 		},
 	}
 
diff --git a/wasm/binary/encoder_test.go b/wasm/binary/encoder_test.go
index 903466b9..6dc938b5 100644
--- a/wasm/binary/encoder_test.go
+++ b/wasm/binary/encoder_test.go
@@ -32,34 +32,6 @@ func TestModule_Encode(t *testing.T) {
 				0x06, // the Module name simple is 6 bytes long
 				's', 'i', 'm', 'p', 'l', 'e'),
 		},
-		{
-			name: "only custom section",
-			input: &wasm.Module{CustomSections: map[string][]byte{
-				"meme": {1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
-			}},
-			expected: append(append(magic, version...),
-				wasm.SectionIDCustom, 0xf, // 15 bytes in this section
-				0x04, 'm', 'e', 'm', 'e',
-				1, 2, 3, 4, 5, 6, 7, 8, 9, 0),
-		},
-		{
-			name: "name section and a custom section", // name should encode last
-			input: &wasm.Module{
-				NameSection: &wasm.NameSection{ModuleName: "simple"},
-				CustomSections: map[string][]byte{
-					"meme": {1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
-				},
-			},
-			expected: append(append(magic, version...),
-				wasm.SectionIDCustom, 0xf, // 15 bytes in this section
-				0x04, 'm', 'e', 'm', 'e',
-				1, 2, 3, 4, 5, 6, 7, 8, 9, 0,
-				wasm.SectionIDCustom, 0x0e, // 14 bytes in this section
-				0x04, 'n', 'a', 'm', 'e',
-				subsectionIDModuleName, 0x07, // 7 bytes in this subsection
-				0x06, // the Module name simple is 6 bytes long
-				's', 'i', 'm', 'p', 'l', 'e'),
-		},
 		{
 			name: "type section",
 			input: &wasm.Module{
diff --git a/wasm/binary/names_test.go b/wasm/binary/names_test.go
index d148b7cb..0cf34505 100644
--- a/wasm/binary/names_test.go
+++ b/wasm/binary/names_test.go
@@ -1,6 +1,7 @@
 package binary
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -201,7 +202,8 @@ func TestDecodeNameSection(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			ns, err := decodeNameSection(encodeNameSectionData(tc.input))
+			data := encodeNameSectionData(tc.input)
+			ns, err := decodeNameSection(bytes.NewReader(data), uint64(len(data)))
 			require.NoError(t, err)
 			require.Equal(t, tc.input, ns)
 		})
@@ -302,7 +304,7 @@ func TestDecodeNameSection_Errors(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			_, err := decodeNameSection(tc.input)
+			_, err := decodeNameSection(bytes.NewReader(tc.input), uint64(len(tc.input)))
 			require.EqualError(t, err, tc.expectedErr)
 		})
 	}
diff --git a/wasm/binary/section_test.go b/wasm/binary/section_test.go
index ee77a53d..78b480e4 100644
--- a/wasm/binary/section_test.go
+++ b/wasm/binary/section_test.go
@@ -111,6 +111,11 @@ func TestDecodeExportSection_Errors(t *testing.T) {
 	}
 }
 
+func TestEncodeFunctionSection(t *testing.T) {
+	require.Equal(t, []byte{wasm.SectionIDFunction, 0x2, 0x01, 0x05}, encodeFunctionSection([]wasm.Index{5}))
+}
+
+// TestEncodeStartSection uses the same index as TestEncodeFunctionSection to highlight the encoding is different.
 func TestEncodeStartSection(t *testing.T) {
 	require.Equal(t, []byte{wasm.SectionIDStart, 0x01, 0x05}, encodeStartSection(5))
 }
diff --git a/wasm/jit/jit_arm64_test.go b/wasm/jit/jit_arm64_test.go
index fa291f3e..965badc8 100644
--- a/wasm/jit/jit_arm64_test.go
+++ b/wasm/jit/jit_arm64_test.go
@@ -4,7 +4,6 @@
 package jit
 
 import (
-	"context"
 	"fmt"
 	"math"
 	"math/bits"
@@ -50,58 +49,6 @@ func (j *jitEnv) requireNewCompiler(t *testing.T) *arm64Compiler {
 	return ret
 }
 
-// TODO: delete this as this could be a duplication from other tests especially spectests.
-// Use this until we could run spectests on arm64.
-func TestArm64CompilerEndToEnd(t *testing.T) {
-	ctx := context.Background()
-	for _, tc := range []struct {
-		name string
-		body []byte
-		sig  *wasm.FunctionType
-	}{
-		{name: "empty", body: []byte{wasm.OpcodeEnd}, sig: &wasm.FunctionType{}},
-		{name: "br .return", body: []byte{wasm.OpcodeBr, 0, wasm.OpcodeEnd}, sig: &wasm.FunctionType{}},
-		{
-			name: "consts",
-			body: []byte{
-				wasm.OpcodeI32Const, 1, wasm.OpcodeI64Const, 1,
-				wasm.OpcodeF32Const, 1, 1, 1, 1, wasm.OpcodeF64Const, 1, 2, 3, 4, 5, 6, 7, 8,
-				wasm.OpcodeEnd,
-			},
-			// We push four constants.
-			sig: &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI32, wasm.ValueTypeI64, wasm.ValueTypeF32, wasm.ValueTypeF64}},
-		},
-		{
-			name: "add",
-			body: []byte{wasm.OpcodeI32Const, 1, wasm.OpcodeI32Const, 1, wasm.OpcodeI32Add, wasm.OpcodeEnd},
-			sig:  &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI32}},
-		},
-		{
-			name: "sub",
-			body: []byte{wasm.OpcodeI64Const, 1, wasm.OpcodeI64Const, 1, wasm.OpcodeI64Sub, wasm.OpcodeEnd},
-			sig:  &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI64}},
-		},
-		{
-			name: "mul",
-			body: []byte{wasm.OpcodeI64Const, 1, wasm.OpcodeI64Const, 1, wasm.OpcodeI64Mul, wasm.OpcodeEnd},
-			sig:  &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI64}},
-		},
-	} {
-		tc := tc
-		t.Run(tc.name, func(t *testing.T) {
-			engine := newEngine()
-			f := &wasm.FunctionInstance{
-				FunctionType: &wasm.TypeInstance{Type: tc.sig},
-				Body:         tc.body,
-			}
-			err := engine.Compile(f)
-			require.NoError(t, err)
-			_, err = engine.Call(ctx, 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.
@@ -109,25 +56,77 @@ func TestArchContextOffsetInEngine(t *testing.T) {
 }
 
 func TestArm64Compiler_returnFunction(t *testing.T) {
-	env := newJITEnvironment()
+	t.Run("exit", func(t *testing.T) {
+		env := newJITEnvironment()
 
-	// Build code.
-	compiler := env.requireNewCompiler(t)
-	err := compiler.emitPreamble()
-	require.NoError(t, err)
-	compiler.returnFunction()
+		// Build code.
+		compiler := env.requireNewCompiler(t)
+		err := compiler.emitPreamble()
+		require.NoError(t, err)
+		compiler.returnFunction()
 
-	// Generate the code under test.
-	code, _, _, err := compiler.compile()
-	require.NoError(t, err)
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+
+		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())
+	})
+	t.Run("deep call stack", func(t *testing.T) {
+		env := newJITEnvironment()
+		engine := env.engine()
+
+		// Push the call frames.
+		const callFrameNums = 10
+		stackPointerToExpectedValue := map[uint64]uint32{}
+		for funcaddr := wasm.FunctionAddress(0); funcaddr < callFrameNums; funcaddr++ {
+			//	Each function pushes its funcaddr and soon returns.
+			compiler := env.requireNewCompiler(t)
+			err := compiler.emitPreamble()
+			require.NoError(t, err)
+
+			// Push its funcaddr.
+			expValue := uint32(funcaddr)
+			err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: expValue})
+			require.NoError(t, err)
+
+			err = compiler.returnFunction()
+			require.NoError(t, err)
+
+			code, _, _, err := compiler.compile()
+			require.NoError(t, err)
+
+			// Compiles and adds to the engine.
+			compiledFunction := &compiledFunction{codeSegment: code, codeInitialAddress: uintptr(unsafe.Pointer(&code[0]))}
+			engine.addCompiledFunction(funcaddr, compiledFunction)
+
+			// Pushes the frame whose return address equals the beginning of the function just compiled.
+			frame := callFrame{
+				// Set the return address to the beginning of the function so that we can execute the constI32 above.
+				returnAddress: compiledFunction.codeInitialAddress,
+				// Note: return stack base pointer is set to funcaddr*10 and this is where the const should be pushed.
+				returnStackBasePointer: uint64(funcaddr) * 10,
+				compiledFunction:       compiledFunction,
+			}
+			engine.callFrameStack[engine.globalContext.callFrameStackPointer] = frame
+			engine.globalContext.callFrameStackPointer++
+			stackPointerToExpectedValue[frame.returnStackBasePointer] = expValue
+		}
+
+		require.Equal(t, uint64(callFrameNums), env.callFrameStackPointer())
 
-	// Run native code.
-	env.exec(code)
+		// Run code from the top frame.
+		env.exec(engine.callFrameTop().compiledFunction.codeSegment)
 
-	// 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())
+		// Check the exit status and the values on stack.
+		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+		for pos, exp := range stackPointerToExpectedValue {
+			require.Equal(t, exp, uint32(env.stack()[pos]))
+		}
+	})
 }
 
 func TestArm64Compiler_exit(t *testing.T) {
@@ -273,7 +272,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.releaseRegisterToStack(compiler.locationStack.peek())
-			compiler.returnFunction()
+			compiler.exit(jitCallStatusCodeReturned)
 
 			// Generate the code under test.
 			code, _, _, err := compiler.compile()
@@ -352,7 +351,7 @@ func TestArm64Compiler_loadValueOnStackToRegister(t *testing.T) {
 
 			// Release the value to the memory stack so that we can see the value after exiting.
 			compiler.releaseRegisterToStack(loc)
-			compiler.returnFunction()
+			compiler.exit(jitCallStatusCodeReturned)
 
 			// Generate the code under test.
 			code, _, _, err := compiler.compile()
@@ -1421,8 +1420,6 @@ func TestArm64Compiler_compileBrIf(t *testing.T) {
 				require.NoError(t, err)
 			},
 		},
-		// {name: "EQ"} TODO: after compileEq support
-		// {name: "NE"} TODO: after compileNe support
 		{
 			name: "HS",
 			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
@@ -1514,6 +1511,30 @@ func TestArm64Compiler_compileBrIf(t *testing.T) {
 				require.NoError(t, err)
 			},
 		},
+		{
+			name: "EQ",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(1), uint32(1)
+				if shoulGoElse {
+					x2++
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				err := compiler.compileEq(&wazeroir.OperationEq{Type: wazeroir.UnsignedTypeI32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "NE",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(1), uint32(2)
+				if shoulGoElse {
+					x2 = x1
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				err := compiler.compileNe(&wazeroir.OperationNe{Type: wazeroir.UnsignedTypeI32})
+				require.NoError(t, err)
+			},
+		},
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
@@ -1566,3 +1587,159 @@ func TestArm64Compiler_compileBrIf(t *testing.T) {
 		})
 	}
 }
+
+func TestArm64Compiler_readInstructionAddress(t *testing.T) {
+	t.Run("target instruction not found", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+
+		err := compiler.emitPreamble()
+		require.NoError(t, err)
+
+		// Set the acquisition target instruction to the one after JMP.
+		compiler.readInstructionAddress(obj.AJMP, reservedRegisterForTemporary)
+
+		compiler.exit(jitCallStatusCodeReturned)
+
+		// If generate the code without JMP after readInstructionAddress,
+		// the call back added must return error.
+		_, _, _, err = compiler.compile()
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "target instruction not found")
+	})
+	t.Run("too large offset", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+
+		err := compiler.emitPreamble()
+		require.NoError(t, err)
+
+		// Set the acquisition target instruction to the one after RET.
+		compiler.readInstructionAddress(obj.ARET, reservedRegisterForTemporary)
+
+		// Add many instruction between the target and readInstructionAddress.
+		for i := 0; i < 100; i++ {
+			compiler.compileConstI32(&wazeroir.OperationConstI32{Value: 10})
+		}
+
+		ret := compiler.newProg()
+		ret.As = obj.ARET
+		ret.To.Type = obj.TYPE_REG
+		ret.To.Reg = reservedRegisterForTemporary
+		compiler.returnFunction()
+
+		// If generate the code with too many instruction between ADR and
+		// the target, compile must fail.
+		_, _, _, err = compiler.compile()
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "too large offset")
+	})
+	t.Run("ok", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+
+		err := compiler.emitPreamble()
+		require.NoError(t, err)
+
+		// 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)
+
+		// Branch to the instruction after RET below via the absolute
+		// address stored in destinationRegister.
+		compiler.emitUnconditionalBranchToAddressOnRegister(addressReg)
+
+		// If we fail to branch, we reach here and exit with unreachable status,
+		// so the assertion would fail.
+		compiler.exit(jitCallStatusCodeUnreachable)
+
+		// This could be the read instruction target as this is the
+		// right after RET. Therefore, the branch instruction above
+		// must target here.
+		err = compiler.returnFunction()
+		require.NoError(t, err)
+
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+
+		env.exec(code)
+
+		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	})
+}
+
+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)
+
+		// 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}}
+
+			err := compiler.emitPreamble()
+			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.returnFunction()
+			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.emitPreamble()
+		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.callFunction(wasm.FunctionAddress(i), targetFunctionType)
+			require.NoError(t, err)
+		}
+
+		err = compiler.returnFunction()
+		require.NoError(t, err)
+
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+
+		env.exec(code)
+
+		// 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())
+	})
+}
diff --git a/wasm/module_test.go b/wasm/module_test.go
index 7789a25e..13a463ad 100644
--- a/wasm/module_test.go
+++ b/wasm/module_test.go
@@ -203,3 +203,103 @@ func TestModule_allDeclarations(t *testing.T) {
 		})
 	}
 }
+
+func TestModule_SectionSize(t *testing.T) {
+	i32, f32 := ValueTypeI32, ValueTypeF32
+	zero := uint32(0)
+	empty := &ConstantExpression{Opcode: OpcodeI32Const, Data: []byte{0x00}}
+
+	tests := []struct {
+		name     string
+		input    *Module
+		expected map[string]uint32
+	}{
+		{
+			name:     "empty",
+			input:    &Module{},
+			expected: map[string]uint32{},
+		},
+		{
+			name:     "only name section",
+			input:    &Module{NameSection: &NameSection{ModuleName: "simple"}},
+			expected: map[string]uint32{"custom": 1},
+		},
+		{
+			name: "type section",
+			input: &Module{
+				TypeSection: []*FunctionType{
+					{},
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+					{Params: []ValueType{i32, i32, i32, i32}, Results: []ValueType{i32}},
+				},
+			},
+			expected: map[string]uint32{"type": 3},
+		},
+		{
+			name: "type and import section",
+			input: &Module{
+				TypeSection: []*FunctionType{
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+					{Params: []ValueType{f32, f32}, Results: []ValueType{f32}},
+				},
+				ImportSection: []*Import{
+					{
+						Module: "Math", Name: "Mul",
+						Kind:     ImportKindFunc,
+						DescFunc: 1,
+					}, {
+						Module: "Math", Name: "Add",
+						Kind:     ImportKindFunc,
+						DescFunc: 0,
+					},
+				},
+			},
+			expected: map[string]uint32{"import": 2, "type": 2},
+		},
+		{
+			name: "type function and start section",
+			input: &Module{
+				TypeSection:     []*FunctionType{{}},
+				FunctionSection: []Index{0},
+				CodeSection: []*Code{
+					{Body: []byte{OpcodeLocalGet, 0, OpcodeLocalGet, 1, OpcodeI32Add, OpcodeEnd}},
+				},
+				ExportSection: map[string]*Export{
+					"AddInt": {Name: "AddInt", Kind: ExportKindFunc, Index: Index(0)},
+				},
+				StartSection: &zero,
+			},
+			expected: map[string]uint32{"code": 1, "export": 1, "function": 1, "start": 1, "type": 1},
+		},
+		{
+			name: "memory and data",
+			input: &Module{
+				MemorySection: []*MemoryType{{Min: 1}},
+				DataSection:   []*DataSegment{{MemoryIndex: 0, OffsetExpression: empty}},
+			},
+			expected: map[string]uint32{"data": 1, "memory": 1},
+		},
+		{
+			name: "table and element",
+			input: &Module{
+				TableSection:   []*TableType{{ElemType: 0x70, Limit: &LimitsType{Min: 1}}},
+				ElementSection: []*ElementSegment{{TableIndex: 0, OffsetExpr: empty}},
+			},
+			expected: map[string]uint32{"element": 1, "table": 1},
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			actual := map[string]uint32{}
+			for i := SectionID(0); i <= SectionIDData; i++ {
+				if size := tc.input.SectionElementCount(i); size > 0 {
+					actual[SectionIDName(i)] = size
+				}
+			}
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}
diff --git a/wasm/text/decoder_test.go b/wasm/text/decoder_test.go
index 43ab766e..5b083703 100644
--- a/wasm/text/decoder_test.go
+++ b/wasm/text/decoder_test.go
@@ -1343,34 +1343,28 @@ func TestDecodeModule(t *testing.T) {
 			},
 		},
 		{
-			name: "export different memory - numeric",
+			name: "export memory - numeric",
 			input: `(module
 	(memory 0)
-	(memory 1)
 	(export "foo" (memory 0))
-	(export "bar" (memory 1))
 )`,
 			expected: &wasm.Module{
-				MemorySection: []*wasm.MemoryType{{Min: 0}, {Min: 1}},
+				MemorySection: []*wasm.MemoryType{{Min: 0}},
 				ExportSection: map[string]*wasm.Export{
 					"foo": {Name: "foo", Kind: wasm.ExportKindMemory, Index: 0},
-					"bar": {Name: "bar", Kind: wasm.ExportKindMemory, Index: 1},
 				},
 			},
 		},
 		{
-			name: "export different memory - numeric - late",
+			name: "export memory - numeric - late",
 			input: `(module
 	(export "foo" (memory 0))
-	(export "bar" (memory 1))
 	(memory 0)
-	(memory 1)
 )`,
 			expected: &wasm.Module{
-				MemorySection: []*wasm.MemoryType{{Min: 0}, {Min: 1}},
+				MemorySection: []*wasm.MemoryType{{Min: 0}},
 				ExportSection: map[string]*wasm.Export{
 					"foo": {Name: "foo", Kind: wasm.ExportKindMemory, Index: 0},
-					"bar": {Name: "bar", Kind: wasm.ExportKindMemory, Index: 1},
 				},
 			},
 		},
@@ -1396,15 +1390,13 @@ func TestDecodeModule(t *testing.T) {
 		{
 			name: "export memory - ID",
 			input: `(module
-    (memory 1)
-    (memory $mem 2)
-    (memory $mam 3)
+    (memory $mem 1)
     (export "memory" (memory $mem))
 )`,
 			expected: &wasm.Module{
-				MemorySection: []*wasm.MemoryType{{Min: 1}, {Min: 2}, {Min: 3}},
+				MemorySection: []*wasm.MemoryType{{Min: 1}},
 				ExportSection: map[string]*wasm.Export{
-					"memory": {Name: "memory", Kind: wasm.ExportKindMemory, Index: 2},
+					"memory": {Name: "memory", Kind: wasm.ExportKindMemory, Index: 0},
 				},
 			},
 		},
@@ -1780,6 +1772,11 @@ func TestParseModule_Errors(t *testing.T) {
 			input:       "(module (import \"foo\" \"bar\" (func (type $v_v))))",
 			expectedErr: "1:41: unknown ID $v_v",
 		},
+		{
+			name:        "import func after func",
+			input:       "(module (func) (import \"\" \"\" (func)))",
+			expectedErr: "1:31: import after module-defined function in module.import[0]",
+		},
 		{
 			name:        "func invalid name",
 			input:       "(module (func baz)))",
@@ -1833,7 +1830,7 @@ func TestParseModule_Errors(t *testing.T) {
 		{
 			name:        "func duplicate result",
 			input:       "(module (func (param i32) (result i32) (result i32)))",
-			expectedErr: "1:41: duplicate result in module.func[0]",
+			expectedErr: "1:41: at most one result allowed in module.func[0]",
 		},
 		{
 			name:        "func double result type",
@@ -1885,11 +1882,6 @@ func TestParseModule_Errors(t *testing.T) {
 			input:       "(module (import \"\" \"\" (func $main)) (func $main)))",
 			expectedErr: "1:43: duplicate ID $main in module.func[0]",
 		},
-		{
-			name:        "import func ID clashes with func ID",
-			input:       "(module (func $main) (import \"\" \"\" (func $main)))",
-			expectedErr: "1:42: duplicate ID $main in module.import[0].func",
-		},
 		{
 			name:        "func ID after result",
 			input:       "(module (func (result i32) $main)))",
@@ -1946,6 +1938,11 @@ func TestParseModule_Errors(t *testing.T) {
 		)`,
 			expectedErr: "3:15: unknown ID $mein in module.code[1].body[1]",
 		},
+		{
+			name:        "second memory",
+			input:       "(module (memory 1) (memory 1))",
+			expectedErr: "1:21: at most one memory allowed in module",
+		},
 		{
 			name: "export duplicates empty name",
 			input: `(module
@@ -1954,7 +1951,7 @@ func TestParseModule_Errors(t *testing.T) {
     (export "" (func 0))
     (export "" (memory 1))
 )`,
-			expectedErr: "5:13: duplicate name \"\" in module.export[1]",
+			expectedErr: `5:13: "" already exported in module.export[1]`,
 		},
 		{
 			name: "export duplicates name",
@@ -1964,7 +1961,7 @@ func TestParseModule_Errors(t *testing.T) {
     (export "a" (func 0))
     (export "a" (memory 1))
 )`,
-			expectedErr: "5:13: duplicate name \"a\" in module.export[1]",
+			expectedErr: `5:13: "a" already exported in module.export[1]`,
 		},
 		{
 			name:        "export double name",
@@ -2057,7 +2054,7 @@ func TestParseModule_Errors(t *testing.T) {
 		{
 			name:        "double start",
 			input:       "(module (start $main) (start $main))",
-			expectedErr: "1:24: redundant start in module",
+			expectedErr: "1:24: at most one start allowed in module",
 		},
 		{
 			name:        "wrong start",
@@ -2098,13 +2095,13 @@ func TestModuleParser_ErrorContext(t *testing.T) {
 	}{
 		{input: "initial", pos: positionInitial, expected: ""},
 		{input: "module", pos: positionModule, expected: "module"},
-		{input: "module import", pos: positionModuleImport, expected: "module.import[0]"},
-		{input: "module import func", pos: positionModuleImportFunc, expected: "module.import[0].func"},
-		{input: "module func", pos: positionModuleFunc, expected: "module.func[0]"},
-		{input: "module memory", pos: positionModuleMemory, expected: "module.memory[0]"},
-		{input: "module export", pos: positionModuleExport, expected: "module.export[0]"},
-		{input: "module export func", pos: positionModuleExportFunc, expected: "module.export[0].func"},
-		{input: "start", pos: positionModuleStart, expected: "module.start"},
+		{input: "module import", pos: positionImport, expected: "module.import[0]"},
+		{input: "module import func", pos: positionImportFunc, expected: "module.import[0].func"},
+		{input: "module func", pos: positionFunc, expected: "module.func[0]"},
+		{input: "module memory", pos: positionMemory, expected: "module.memory[0]"},
+		{input: "module export", pos: positionExport, expected: "module.export[0]"},
+		{input: "module export func", pos: positionExportFunc, expected: "module.export[0].func"},
+		{input: "start", pos: positionStart, expected: "module.start"},
 	}
 
 	for _, tt := range tests {
diff --git a/wasm/text/func_parser_test.go b/wasm/text/func_parser_test.go
index 0c2e2e7c..2294f05a 100644
--- a/wasm/text/func_parser_test.go
+++ b/wasm/text/func_parser_test.go
@@ -213,7 +213,7 @@ func TestFuncParser_Errors(t *testing.T) {
 		{
 			name:        "duplicate result",
 			source:      "(func (result i32) (result i32))",
-			expectedErr: "1:21: duplicate result",
+			expectedErr: "1:21: at most one result allowed",
 		},
 	}
 
diff --git a/wasm/text/memory_parser_test.go b/wasm/text/memory_parser_test.go
index 534bcbb9..9498eb14 100644
--- a/wasm/text/memory_parser_test.go
+++ b/wasm/text/memory_parser_test.go
@@ -60,8 +60,8 @@ func TestMemoryParser(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			namespace := newIndexNamespace()
-			parsed, tp, err := parseMemoryType(namespace, tc.input)
+			memoryNamespace := newIndexNamespace()
+			parsed, tp, err := parseMemoryType(memoryNamespace, tc.input)
 			require.NoError(t, err)
 			require.Equal(t, tc.expected, parsed)
 			require.Equal(t, uint32(1), tp.memoryNamespace.count)
@@ -151,13 +151,13 @@ func TestMemoryParser_Errors(t *testing.T) {
 	})
 }
 
-func parseMemoryType(namespace *indexNamespace, input string) (*wasm.MemoryType, *memoryParser, error) {
+func parseMemoryType(memoryNamespace *indexNamespace, input string) (*wasm.MemoryType, *memoryParser, error) {
 	var parsed *wasm.MemoryType
 	var setFunc onMemory = func(min uint32, max *uint32) tokenParser {
 		parsed = &wasm.MemoryType{Min: min, Max: max}
 		return parseErr
 	}
-	tp := newMemoryParser(namespace, setFunc)
+	tp := newMemoryParser(memoryNamespace, setFunc)
 	// memoryParser starts after the '(memory', so we need to eat it first!
 	_, _, err := lex(skipTokens(2, tp.begin), []byte(input))
 	return parsed, tp, err
diff --git a/wasm/text/type_parser_test.go b/wasm/text/type_parser_test.go
index ec07c33c..76f72c2d 100644
--- a/wasm/text/type_parser_test.go
+++ b/wasm/text/type_parser_test.go
@@ -100,8 +100,8 @@ func TestTypeParser(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			namespace := newIndexNamespace()
-			parsed, tp, err := parseFunctionType(namespace, tc.input)
+			typeNamespace := newIndexNamespace()
+			parsed, tp, err := parseFunctionType(typeNamespace, tc.input)
 			require.NoError(t, err)
 			require.Equal(t, tc.expected, parsed)
 			require.Equal(t, uint32(1), tp.typeNamespace.count)
@@ -246,13 +246,13 @@ func TestTypeParser_Errors(t *testing.T) {
 	})
 }
 
-func parseFunctionType(namespace *indexNamespace, input string) (*wasm.FunctionType, *typeParser, error) {
+func parseFunctionType(typeNamespace *indexNamespace, input string) (*wasm.FunctionType, *typeParser, error) {
 	var parsed *wasm.FunctionType
 	var setFunc onType = func(ft *wasm.FunctionType) tokenParser {
 		parsed = ft
 		return parseErr
 	}
-	tp := newTypeParser(namespace, setFunc)
+	tp := newTypeParser(typeNamespace, setFunc)
 	// typeParser starts after the '(type', so we need to eat it first!
 	_, _, err := lex(skipTokens(2, tp.begin), []byte(input))
 	return parsed, tp, err
diff --git a/wasm/text/typeuse_parser_test.go b/wasm/text/typeuse_parser_test.go
index 74825ce7..73b3e41d 100644
--- a/wasm/text/typeuse_parser_test.go
+++ b/wasm/text/typeuse_parser_test.go
@@ -18,7 +18,7 @@ type typeUseParserTest struct {
 	expectedTypeIdx     wasm.Index
 	expectedParamNames  wasm.NameMap
 
-	expectedOnTypeUsePosition onTypeUsePosition
+	expectedOnTypeUsePosition callbackPosition
 	expectedOnTypeUseToken    tokenType
 	expectedTrailingTokens    []tokenType
 }
@@ -325,14 +325,14 @@ type typeUseTestFunc func(*typeUseParserTest) (*typeUseParser, func(t *testing.T
 func runTypeUseParserTests(t *testing.T, tests []*typeUseParserTest, tf typeUseTestFunc) {
 	moreTests := make([]*typeUseParserTest, 0, len(tests)*2)
 	for _, tt := range tests {
-		tt.expectedOnTypeUsePosition = onTypeUseEndField
+		tt.expectedOnTypeUsePosition = callbackPositionEndField
 		tt.expectedOnTypeUseToken = tokenRParen // at the end of the field ')'
 		tt.expectedTrailingTokens = nil
 
 		kt := *tt // copy
 		kt.name = fmt.Sprintf("%s - trailing keyword", tt.name)
 		kt.source = fmt.Sprintf("%s nop)", tt.source[:len(tt.source)-1])
-		kt.expectedOnTypeUsePosition = onTypeUseUnhandledToken
+		kt.expectedOnTypeUsePosition = callbackPositionUnhandledToken
 		kt.expectedOnTypeUseToken = tokenKeyword // at 'nop' and ')' remains
 		kt.expectedTrailingTokens = []tokenType{tokenRParen}
 		moreTests = append(moreTests, &kt)
@@ -342,11 +342,11 @@ func runTypeUseParserTests(t *testing.T, tests []*typeUseParserTest, tf typeUseT
 		ft.source = fmt.Sprintf("%s (nop))", tt.source[:len(tt.source)-1])
 		// Two outcomes, we've reached a field not named "type", "param" or "result" or we completed "result"
 		if strings.Contains(tt.source, "result") {
-			ft.expectedOnTypeUsePosition = onTypeUseUnhandledToken
+			ft.expectedOnTypeUsePosition = callbackPositionUnhandledToken
 			ft.expectedOnTypeUseToken = tokenLParen // at '(' and 'nop))' remain
 			ft.expectedTrailingTokens = []tokenType{tokenKeyword, tokenRParen, tokenRParen}
 		} else {
-			ft.expectedOnTypeUsePosition = onTypeUseUnhandledField
+			ft.expectedOnTypeUsePosition = callbackPositionUnhandledField
 			ft.expectedOnTypeUseToken = tokenKeyword // at 'nop' and '))' remain
 			ft.expectedTrailingTokens = []tokenType{tokenRParen, tokenRParen}
 		}
@@ -360,7 +360,7 @@ func runTypeUseParserTests(t *testing.T, tests []*typeUseParserTest, tf typeUseT
 			var parsedTypeIdx wasm.Index
 			var parsedParamNames wasm.NameMap
 			p := &collectTokenTypeParser{}
-			var setTypeUse onTypeUse = func(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+			var setTypeUse onTypeUse = func(typeIdx wasm.Index, paramNames wasm.NameMap, pos callbackPosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 				parsedTypeIdx = typeIdx
 				parsedParamNames = paramNames
 				require.Equal(t, tc.expectedOnTypeUsePosition, pos)
@@ -534,11 +534,11 @@ func TestTypeUseParser_FailsMatch(t *testing.T) {
 	}
 }
 
-var ignoreTypeUse onTypeUse = func(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+var ignoreTypeUse onTypeUse = func(typeIdx wasm.Index, paramNames wasm.NameMap, pos callbackPosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 	return parseNoop, nil
 }
 
-var failOnTypeUse onTypeUse = func(typeIdx wasm.Index, paramNames wasm.NameMap, pos onTypeUsePosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+var failOnTypeUse onTypeUse = func(typeIdx wasm.Index, paramNames wasm.NameMap, pos callbackPosition, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 	return nil, errors.New("unexpected to call onTypeUse on error")
 }