Repository: tetratelabs/wazero
Pull request #265 (2022-02-21 01:36:53)

jit(arm64): support for call_indirect and br_table
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

mathetake: ok done with call_indirect
mathetake: ok passed all tests!!!!!!!!!!!!!
mathetake: after this I will enable spectest (plus do the possible bug fixes) in another PR then #187 will be closed 
mathetake: most of the logic already had been done in amd64 previously, so this was a easy game

Patch:
diff --git a/internal/wasm/jit/engine.go b/internal/wasm/jit/engine.go
index eec16ee4..ace48e76 100644
--- a/internal/wasm/jit/engine.go
+++ b/internal/wasm/jit/engine.go
@@ -225,7 +225,8 @@ const (
 	moduleInstanceTablesOffset  = 80
 
 	// Offsets for wasm.TableInstance.
-	tableInstanceTableOffset = 0
+	tableInstanceTableOffset    = 0
+	tableInstanceTableLenOffset = 8
 
 	// Offsets for wasm.MemoryInstance.
 	memoryInstanceBufferOffset    = 0
diff --git a/internal/wasm/jit/jit_amd64.go b/internal/wasm/jit/jit_amd64.go
index f18ea010..4e748fad 100644
--- a/internal/wasm/jit/jit_amd64.go
+++ b/internal/wasm/jit/jit_amd64.go
@@ -674,8 +674,7 @@ func (c *amd64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
 func (c *amd64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
 	index := c.locationStack.pop()
 
-	// If the operation doesn't have target but default,
-	// branch into the default label and return early.
+	// If the operation only consists of the default target, we branch into it and return early.
 	if len(o.Targets) == 0 {
 		c.locationStack.releaseRegister(index)
 		if err := c.emitDropRange(o.Default.ToDrop); err != nil {
@@ -818,7 +817,7 @@ func (c *amd64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
 			locationStack = saved.clone()
 		} else {
 			target = o.Default
-			// If this is the deafult branch, we just use the original one
+			// If this is the default branch, we use the original one
 			// as this is the last code in this block.
 			locationStack = saved
 		}
@@ -839,7 +838,7 @@ func (c *amd64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
 			if uint64(nop.Pc)-uint64(base) >= math.MaxUint32 {
 				// TODO: this happens when users try loading an extremely large webassembly binary
 				// which contains a br_table statement with approximately 4294967296 (2^32) targets.
-				// We would like to support that binary, but realistically speacking, that kind of binary
+				// We would like to support that binary, but realistically speaking, that kind of binary
 				// could result in more than ten giga bytes of native JITed code where we have to care about
 				// huge stacks whose height might exceed 32-bit range, and such huge stack doesn't work with the
 				// current implementation.
@@ -5416,17 +5415,6 @@ func (c *amd64Compiler) initializeModuleContext() error {
 	// is ensured by function validation at module instantiation phase, and that's
 	// why it is ok to skip the initialization if the module's table doesn't exist.
 	if len(c.f.ModuleInstance.Tables) > 0 {
-		getTablesCount := c.newProg()
-		getTablesCount.As = x86.AMOVQ
-		getTablesCount.To.Type = obj.TYPE_REG
-		getTablesCount.To.Reg = tmpRegister
-		getTablesCount.From.Type = obj.TYPE_MEM
-		getTablesCount.From.Reg = moduleInstanceAddressRegister
-		// We add "+8" to get the length of ModuleInstance.Tables
-		// since the slice header {Data uintptr, Len, int64, Cap int64} internally.
-		getTablesCount.From.Offset = moduleInstanceTablesOffset + 8
-		c.addInstruction(getTablesCount)
-
 		// First, we need to read the *wasm.TableInstance.
 		readTableInstancePointer := c.newProg()
 		readTableInstancePointer.As = x86.AMOVQ
@@ -5438,7 +5426,7 @@ func (c *amd64Compiler) initializeModuleContext() error {
 		c.addInstruction(readTableInstancePointer)
 
 		resolveTableInstanceAddressFromPointer := c.newProg()
-		resolveTableInstanceAddressFromPointer.As = x86.AMOVQ // Note this is LEA instruction.
+		resolveTableInstanceAddressFromPointer.As = x86.AMOVQ
 		resolveTableInstanceAddressFromPointer.To.Type = obj.TYPE_REG
 		resolveTableInstanceAddressFromPointer.To.Reg = tmpRegister
 		resolveTableInstanceAddressFromPointer.From.Type = obj.TYPE_MEM
@@ -5473,9 +5461,7 @@ func (c *amd64Compiler) initializeModuleContext() error {
 		readTableLength.To.Reg = tmpRegister2
 		readTableLength.From.Type = obj.TYPE_MEM
 		readTableLength.From.Reg = tmpRegister
-		// We add "+8" to get the length of Tables[0].Table
-		// since the slice header {Data uintptr, Len, int64, Cap int64} internally.
-		readTableLength.From.Offset = tableInstanceTableOffset + 8
+		readTableLength.From.Offset = tableInstanceTableLenOffset
 		c.addInstruction(readTableLength)
 
 		// And put the length into tableSliceLen.
diff --git a/internal/wasm/jit/jit_arm64.go b/internal/wasm/jit/jit_arm64.go
index 275d4e1a..71e3f6f2 100644
--- a/internal/wasm/jit/jit_arm64.go
+++ b/internal/wasm/jit/jit_arm64.go
@@ -13,6 +13,7 @@
 package jit
 
 import (
+	"encoding/binary"
 	"fmt"
 	"math"
 	"unsafe"
@@ -106,6 +107,13 @@ type arm64Compiler struct {
 	afterAssembleCallback []func(code []byte) error
 	// onStackPointerCeilDeterminedCallBack hold a callback which are called when the ceil of stack pointer is determined before generating native code.
 	onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64)
+	// compiledFunctionStaticData holds br_table offset tables.
+	// See compiledFunctionStaticData and arm64Compiler.compileBrTable.
+	staticData compiledFunctionStaticData
+}
+
+func (c *arm64Compiler) addStaticData(d []byte) {
+	c.staticData = append(c.staticData, d)
 }
 
 // compile implements compiler.compile for the arm64 architecture.
@@ -900,20 +908,166 @@ func (c *arm64Compiler) assignBranchTarget(labelKey string, br *obj.Prog) {
 	}
 }
 
+// compileBrTable implements compiler.compileBrTable for the arm64 architecture.
 func (c *arm64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	// If the operation only consists of the default target, we branch into it and return early.
+	if len(o.Targets) == 0 {
+		loc := c.locationStack.pop()
+		if loc.onRegister() {
+			c.markRegisterUnused(loc.register)
+		}
+		if err := c.compileDropRange(o.Default.ToDrop); err != nil {
+			return err
+		}
+		return c.compileBranchInto(o.Default.Target)
+	}
+
+	index := c.locationStack.pop()
+	if err := c.compileEnsureOnGeneralPurposeRegister(index); err != nil {
+		return err
+	}
+
+	if isZeroRegister(index.register) {
+		reg, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+		if err != nil {
+			return err
+		}
+		index.setRegister(reg)
+		c.markRegisterUsed(reg)
+
+		// Zero the value on a picked register.
+		c.compileRegisterToRegisterInstruction(arm64.AMOVD, zeroRegister, reg)
+	}
+
+	tmpReg, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+	if err != nil {
+		return err
+	}
+
+	// Load the branch table's length.
+	// "tmpReg = len(o.Targets)"
+	c.compileConstToRegisterInstruction(arm64.AMOVW, int64(len(o.Targets)), tmpReg)
+	// Compare the length with offest.
+	c.compileTwoRegistersToNoneInstruction(arm64.ACMPW, tmpReg, index.register)
+	// If the value exceeds the length, we will branch into the default target (corresponding to len(o.Targets) index).
+	brDefaultIndex := c.compilelBranchInstruction(arm64.ABLO)
+	c.compileRegisterToRegisterInstruction(arm64.AMOVW, tmpReg, index.register)
+	c.setBranchTargetOnNext(brDefaultIndex)
+
+	// We prepare the static data which holds the offset of
+	// each target's first instruction (incl. default)
+	// relative to the beginning of label tables.
+	//
+	// For example, if we have targets=[L0, L1] and default=L_DEFAULT,
+	// we emit the the code like this at [Emit the code for each targets and default branch] below.
+	//
+	// L0:
+	//  0x123001: XXXX, ...
+	//  .....
+	// L1:
+	//  0x123005: YYY, ...
+	//  .....
+	// L_DEFAULT:
+	//  0x123009: ZZZ, ...
+	//
+	// then offsetData becomes like [0x0, 0x5, 0x8].
+	// By using this offset list, we could jump into the label for the index by
+	// "jmp offsetData[index]+0x123001" and "0x123001" can be acquired by "LEA"
+	// instruction.
+	//
+	// Note: We store each offset of 32-bit unsigned integer as 4 consecutive bytes. So more precisely,
+	// the above example's offsetData would be [0x0, 0x0, 0x0, 0x0, 0x5, 0x0, 0x0, 0x0, 0x8, 0x0, 0x0, 0x0].
+	//
+	// Note: this is similar to how GCC implements Switch statements in C.
+	offsetData := make([]byte, 4*(len(o.Targets)+1))
+	c.addStaticData(offsetData)
+
+	// "tmpReg = &offsetData[0]"
+	c.compileConstToRegisterInstruction(
+		arm64.AMOVD,
+		int64(uintptr(unsafe.Pointer(&offsetData[0]))),
+		tmpReg,
+	)
+
+	// "index.register = tmpReg + (index.register << 2) (== &offsetData[offset])"
+	c.compileAddInstructionWithLeftShiftedRegister(index.register, 2, tmpReg, index.register)
+
+	// "index.regsetr = *index.reigier (== offsetData[offset])"
+	c.compileMemoryToRegisterInstruction(arm64.AMOVW, index.register, 0, index.register)
+
+	// Now we read the address of the beginning of the jump table.
+	// In the above example, this corresponds to reading the address of 0x123001.
+	c.compileReadInstructionAddress(obj.AJMP, tmpReg)
+
+	// Now we have the address of L0 in tmp register, and the offset to the target label in the index.register.
+	// So we could achieve the br_table jump by adding them and jump into the resulting address.
+	c.compileRegisterToRegisterInstruction(arm64.AADD, tmpReg, index.register)
+
+	c.compileUnconditionalBranchToAddressOnRegister(index.register)
+
+	// We no longer need the index's register, so mark it unused.
+	c.markRegisterUnused(index.register)
+
+	// [Emit the code for each targets and default branch]
+	labelInitialInstructions := make([]*obj.Prog, len(o.Targets)+1)
+	saved := c.locationStack
+	for i := range labelInitialInstructions {
+		// Emit the initial instruction of each target where
+		// we use NOP as we don't yet know the next instruction in each label.
+		init := c.compileNOP()
+		labelInitialInstructions[i] = init
+
+		var locationStack *valueLocationStack
+		var target *wazeroir.BranchTargetDrop
+		if i < len(o.Targets) {
+			target = o.Targets[i]
+			// Clone the location stack so the branch-specific code doesn't
+			// affect others.
+			locationStack = saved.clone()
+		} else {
+			target = o.Default
+			// If this is the default branch, we use the original one
+			// as this is the last code in this block.
+			locationStack = saved
+		}
+		c.setLocationStack(locationStack)
+		if err := c.compileDropRange(target.ToDrop); err != nil {
+			return err
+		}
+		if err := c.compileBranchInto(target.Target); err != nil {
+			return err
+		}
+	}
+
+	c.afterAssembleCallback = append(c.afterAssembleCallback, func(code []byte) error {
+		// Build the offset table for each target including default one.
+		base := labelInitialInstructions[0].Pc // This corresponds to the L0's address in the example.
+		for i, nop := range labelInitialInstructions {
+			if uint64(nop.Pc)-uint64(base) >= math.MaxUint32 {
+				// TODO: this happens when users try loading an extremely large webassembly binary
+				// which contains a br_table statement with approximately 4294967296 (2^32) targets.
+				// We would like to support that binary, but realistically speaking, that kind of binary
+				// could result in more than ten giga bytes of native JITed code where we have to care about
+				// huge stacks whose height might exceed 32-bit range, and such huge stack doesn't work with the
+				// current implementation.
+				return fmt.Errorf("too large br_table")
+			}
+			// We store the offset from the beiggning of the L0's initial instruction.
+			binary.LittleEndian.PutUint32(offsetData[i*4:(i+1)*4], uint32(nop.Pc)-uint32(base))
+		}
+		return nil
+	})
+	return nil
 }
 
 // compileCall implements compiler.compileCall for the arm64 architecture.
 func (c *arm64Compiler) compileCall(o *wazeroir.OperationCall) error {
 	target := c.f.ModuleInstance.Functions[o.FunctionIndex]
-	return c.compileCallImpl(target.Address, target.FunctionType.Type)
+	return c.compileCallImpl(target.Address, nilRegister, target.FunctionType.Type)
 }
 
-// compileCallImpl implements compiler.compileCall and compiler.compileCallIndirect (TODO) for the arm64 architecture.
-func (c *arm64Compiler) compileCallImpl(addr wasm.FunctionAddress, functype *wasm.FunctionType) error {
-	// TODO: the following code can be generalized for CallIndirect.
-
+// compileCallImpl implements compiler.compileCall and compiler.compileCallIndirect for the arm64 architecture.
+func (c *arm64Compiler) compileCallImpl(addr wasm.FunctionAddress, addrRegister int16, functype *wasm.FunctionType) error {
 	// Release all the registers as our calling convention requires the caller-save.
 	if err := c.compileReleaseAllRegistersToStack(); err != nil {
 		return err
@@ -946,10 +1100,29 @@ func (c *arm64Compiler) compileCallImpl(addr wasm.FunctionAddress, functype *was
 	brIfCallFrameStackOK := c.compilelBranchInstruction(arm64.ABNE)
 
 	// If these values equal, we need to grow the callFrame stack.
+	// For call_indirect, we need to push the value back to the register.
+	if !isNilRegister(addrRegister) {
+		// If we need to get the target funcaddr from register (call_indirect case), we must save it before growing callframe stack,
+		// as the register is not saved across function calls.
+		savedOffsetLocation := c.locationStack.pushValueLocationOnRegister(addrRegister)
+		c.compileReleaseRegisterToStack(savedOffsetLocation)
+	}
+
 	if err := c.compileCallGoFunction(jitCallStatusCodeCallBuiltInFunction, builtinFunctionAddressGrowCallFrameStack); err != nil {
 		return err
 	}
 
+	// For call_indirect, we need to push the value back to the register.
+	if !isNilRegister(addrRegister) {
+		// Since this is right after callGoFunction, we have to initialize the stack base pointer
+		// to properly load the value on memory stack.
+		c.compileReservedStackBasePointerRegisterInitialization()
+
+		savedOffsetLocation := c.locationStack.pop()
+		savedOffsetLocation.setRegister(addrRegister)
+		c.compileLoadValueOnStackToRegister(savedOffsetLocation)
+	}
+
 	// On the function return, we again have to set engine.callFrameStackPointer into callFrameStackPointerRegister.
 	// "callFrameStackPointerRegister = engine.callFrameStackPointer"
 	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
@@ -995,12 +1168,12 @@ func (c *arm64Compiler) compileCallImpl(addr wasm.FunctionAddress, functype *was
 	// 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.compileConstToRegisterInstruction(arm64.AADD,
-		int64(c.locationStack.sp)-int64(len(functype.Params)),
-		oldStackBasePointer)
-	c.compileRegisterToMemoryInstruction(arm64.AMOVD,
-		oldStackBasePointer,
-		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset)
+	if offset := int64(c.locationStack.sp) - int64(len(functype.Params)); offset > 0 {
+		c.compileConstToRegisterInstruction(arm64.AADD, offset, oldStackBasePointer)
+		c.compileRegisterToMemoryInstruction(arm64.AMOVD,
+			oldStackBasePointer,
+			reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset)
+	}
 
 	// 3) Set rc.next to specify which function is executed on the current call frame.
 	//
@@ -1012,9 +1185,20 @@ func (c *arm64Compiler) compileCallImpl(addr wasm.FunctionAddress, functype *was
 
 	// Next, read the address of the target function (= &engine.compiledFunctions[offset])
 	// into compiledFunctionAddressRegister.
-	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
-		tmp, int64(addr)*8, // * 8 because the size of *compiledFunction equals 8 bytes.
-		compiledFunctionAddressRegister)
+	if isNilRegister(addrRegister) {
+		c.compileMemoryToRegisterInstruction(
+			arm64.AMOVD,
+			tmp, int64(addr)*8, // * 8 because the size of *compiledFunction equals 8 bytes.
+			compiledFunctionAddressRegister)
+	} else {
+		// Shift addrRegister by 3 because the size of *compiledFunction equals 8 bytes.
+		c.compileConstToRegisterInstruction(arm64.ALSL, 3, addrRegister)
+		c.compileMemoryWithRegisterOffsetToRegisterInstruction(
+			arm64.AMOVD,
+			tmp, addrRegister,
+			compiledFunctionAddressRegister,
+		)
+	}
 
 	// Finally, we are ready to write the address of the target function's *compiledFunction into the new callframe.
 	c.compileRegisterToMemoryInstruction(arm64.AMOVD,
@@ -1046,6 +1230,7 @@ func (c *arm64Compiler) compileCallImpl(addr wasm.FunctionAddress, functype *was
 	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
 		compiledFunctionAddressRegister, compiledFunctionCodeInitialAddressOffset,
 		tmp)
+
 	c.compileUnconditionalBranchToAddressOnRegister(tmp)
 
 	// All the registers used are temporary so we mark them unused.
@@ -1157,8 +1342,103 @@ func (c *arm64Compiler) compileReadInstructionAddress(beforeTargetInst obj.As, d
 	})
 }
 
+// compileCallIndirect implements compiler.compileCallIndirect for the arm64 architecture.
 func (c *arm64Compiler) compileCallIndirect(o *wazeroir.OperationCallIndirect) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	offset := c.locationStack.pop()
+	if err := c.compileEnsureOnGeneralPurposeRegister(offset); err != nil {
+		return err
+	}
+
+	if isZeroRegister(offset.register) {
+		reg, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+		if err != nil {
+			return err
+		}
+		offset.setRegister(reg)
+		c.markRegisterUsed(reg)
+
+		// Zero the value on a picked register.
+		c.compileRegisterToRegisterInstruction(arm64.AMOVD, zeroRegister, reg)
+	}
+
+	tmp, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+	if err != nil {
+		return err
+	}
+
+	// First, we need to check if the offset doesn't exceed the length of table.
+	// "tmp = len(table)"
+	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
+		reservedRegisterForEngine, engineModuleContextTableSliceLenOffset,
+		tmp,
+	)
+	// "cmp tmp, offset"
+	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, tmp, offset.register)
+
+	// If it exceeds len(table), we exit the execution.
+	brIfOffsetOK := c.compilelBranchInstruction(arm64.ABLO)
+	c.compileExitFromNativeCode(jitCallStatusCodeInvalidTableAccess)
+
+	// Otherwise, we proceed to do function type check.
+	c.setBranchTargetOnNext(brIfOffsetOK)
+
+	// We need to obtains the absolute address of table element.
+	// "tmp = &table[0]"
+	c.compileMemoryToRegisterInstruction(
+		arm64.AMOVD,
+		reservedRegisterForEngine, engineModuleContextTableElement0AddressOffset,
+		tmp,
+	)
+	// "offset = tmp + (offset << 4) (== &table[offset])"
+	c.compileAddInstructionWithLeftShiftedRegister(
+		offset.register, 4,
+		tmp,
+		offset.register,
+	)
+
+	// Check if table[offset].TypeID == targetFunctionType.
+	targetFunctionType := c.f.ModuleInstance.Types[o.TypeIndex]
+	// "tmp = table[offset].TypeID"
+	c.compileMemoryToRegisterInstruction(
+		arm64.AMOVD, offset.register, tableElementFunctionTypeIDOffset,
+		tmp,
+	)
+	// "reservedRegisterForTemporary = targetFunctionType.TypeID"
+	c.compileConstToRegisterInstruction(arm64.AMOVD, int64(targetFunctionType.TypeID), reservedRegisterForTemporary)
+	// Compare these two values, and if they equal, we are ready to make function call.
+	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, tmp, reservedRegisterForTemporary)
+	brIfTypeMatched := c.compilelBranchInstruction(arm64.ABEQ)
+
+	// Otherwise, we have to exit the execution with either jitCallStatusCodeTypeMismatchOnIndirectCall or jitCallStatusCodeInvalidTableAccess.
+	{
+		// We exit with jitCallStatusCodeInvalidTableAccess if the targetFunctionType.TypeID equals the uninitialized one (wasm.UninitializedTableElementTypeID).
+		c.compileConstToRegisterInstruction(arm64.AMOVD, int64(wasm.UninitializedTableElementTypeID), reservedRegisterForTemporary)
+		c.compileTwoRegistersToNoneInstruction(arm64.ACMPW, tmp, reservedRegisterForTemporary)
+
+		brIfInitizlied := c.compilelBranchInstruction(arm64.ABNE)
+		c.compileExitFromNativeCode(jitCallStatusCodeInvalidTableAccess)
+
+		// Otherwise exit with jitCallStatusCodeTypeMismatchOnIndirectCall.
+		c.setBranchTargetOnNext(brIfInitizlied)
+		c.compileExitFromNativeCode(jitCallStatusCodeTypeMismatchOnIndirectCall)
+	}
+
+	c.setBranchTargetOnNext(brIfTypeMatched)
+
+	// Now all checks passed, so read the target's function address, and make call.
+	c.compileMemoryToRegisterInstruction(
+		arm64.AMOVW,
+		offset.register, tableElementFunctionAddressOffset,
+		offset.register,
+	)
+
+	if err := c.compileCallImpl(0, offset.register, targetFunctionType.Type); err != nil {
+		return err
+	}
+
+	// The offset register should be marked as un-used as we consumed in the function call.
+	c.markRegisterUnused(offset.register)
+	return nil
 }
 
 // compileDrop implements compiler.compileDrop for the arm64 architecture.
@@ -3036,8 +3316,8 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 	// * engine.moduleContext.globalElement0Address
 	// * engine.moduleContext.memoryElement0Address
 	// * engine.moduleContext.memorySliceLen
-	// * (TODO) engine.moduleContext.tableElement0Address
-	// * (TODO) engine.moduleContext.tableSliceLen
+	// * engine.moduleContext.tableElement0Address
+	// * engine.moduleContext.tableSliceLen
 
 	// Update globalElement0Address.
 	//
@@ -3103,7 +3383,49 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 		)
 	}
 
-	// TODO: Update tableElement0Address and tableSliceLen.
+	// Update tableElement0Address and tableSliceLen.
+	//
+	// Note: if there's table instruction in the function, the existence of the table
+	// is ensured by function validation at module instantiation phase, and that's
+	// why it is ok to skip the initialization if the module's table doesn't exist.
+	if len(c.f.ModuleInstance.Tables) > 0 {
+		// "tmpX = &tables[0] (type of **wasm.TableInstance)"
+		c.compileMemoryToRegisterInstruction(
+			arm64.AMOVD,
+			moduleInstanceAddressRegister, moduleInstanceTablesOffset,
+			tmpX,
+		)
+		// "tmpX = *tmpX (tables[0])"
+		c.compileMemoryToRegisterInstruction(arm64.AMOVD, tmpX, 0, tmpX)
+
+		// Update engine.tableElement0Address.
+		// "tmpY = &tables[0].Table[0]"
+		c.compileMemoryToRegisterInstruction(
+			arm64.AMOVD,
+			tmpX, tableInstanceTableOffset,
+			tmpY,
+		)
+		// "engine.tableElement0Address = tmpY".
+		c.compileRegisterToMemoryInstruction(
+			arm64.AMOVD,
+			tmpY,
+			reservedRegisterForEngine, engineModuleContextTableElement0AddressOffset,
+		)
+
+		// Update engine.tableSliceLen.
+		// "tmpY = len(tables[0].Table)"
+		c.compileMemoryToRegisterInstruction(
+			arm64.AMOVD,
+			tmpX, tableInstanceTableLenOffset,
+			tmpY,
+		)
+		// "engine.tableSliceLen = tmpY".
+		c.compileRegisterToMemoryInstruction(
+			arm64.AMOVD,
+			tmpY,
+			reservedRegisterForEngine, engineModuleContextTableSliceLenOffset,
+		)
+	}
 
 	c.setBranchTargetOnNext(brIfModuleUnchanged)
 	c.locationStack.markRegisterUnused(regs...)
diff --git a/internal/wasm/jit/jit_value_location.go b/internal/wasm/jit/jit_value_location.go
index b36048fe..7ff9a960 100644
--- a/internal/wasm/jit/jit_value_location.go
+++ b/internal/wasm/jit/jit_value_location.go
@@ -17,6 +17,18 @@ func isNilRegister(r int16) bool {
 	return r == nilRegister
 }
 
+func isIntRegister(r int16) bool {
+	return unreservedGeneralPurposeIntRegisters[0] <= r && r <= unreservedGeneralPurposeIntRegisters[len(unreservedGeneralPurposeIntRegisters)-1]
+}
+
+func isFloatRegister(r int16) bool {
+	return generalPurposeFloatRegisters[0] <= r && r <= generalPurposeFloatRegisters[len(generalPurposeFloatRegisters)-1]
+}
+
+func isZeroRegister(r int16) bool {
+	return r == zeroRegister
+}
+
 // conditionalRegisterState indicates a state of the conditional flag register.
 // In arm64, conditional registers are defined as arm64.COND_*.
 // In amd64, we define each flag value in value_locations_amd64.go
@@ -208,7 +220,9 @@ func (s *valueLocationStack) markRegisterUnused(regs ...int16) {
 
 func (s *valueLocationStack) markRegisterUsed(regs ...int16) {
 	for _, reg := range regs {
-		s.usedRegisters[reg] = struct{}{}
+		if !isZeroRegister(reg) {
+			s.usedRegisters[reg] = struct{}{}
+		}
 	}
 }
 
diff --git a/internal/wasm/jit/jit_value_location_amd64.go b/internal/wasm/jit/jit_value_location_amd64.go
index fd064545..c31dcf58 100644
--- a/internal/wasm/jit/jit_value_location_amd64.go
+++ b/internal/wasm/jit/jit_value_location_amd64.go
@@ -18,6 +18,8 @@ const (
 	reservedRegisterForTemporary = nilRegister
 )
 
+const zeroRegister int16 = nilRegister
+
 var (
 	generalPurposeFloatRegisters = []int16{
 		x86.REG_X0, x86.REG_X1, x86.REG_X2, x86.REG_X3,
@@ -37,14 +39,6 @@ var (
 	}
 )
 
-func isIntRegister(r int16) bool {
-	return unreservedGeneralPurposeIntRegisters[0] <= r && r <= unreservedGeneralPurposeIntRegisters[len(unreservedGeneralPurposeIntRegisters)-1]
-}
-
-func isFloatRegister(r int16) bool {
-	return generalPurposeFloatRegisters[0] <= r && r <= generalPurposeFloatRegisters[len(generalPurposeFloatRegisters)-1]
-}
-
 const (
 	conditionalRegisterStateE  = conditionalRegisterStateUnset + 1 + iota // ZF equal to zero
 	conditionalRegisterStateNE                                            //˜ZF not equal to zero
diff --git a/internal/wasm/jit/jit_value_location_arm64.go b/internal/wasm/jit/jit_value_location_arm64.go
index 20e39f09..51c3410c 100644
--- a/internal/wasm/jit/jit_value_location_arm64.go
+++ b/internal/wasm/jit/jit_value_location_arm64.go
@@ -44,15 +44,3 @@ var (
 func simdRegisterForScalarFloatRegister(freg int16) int16 {
 	return freg + (arm64.REG_F31 - arm64.REG_F0) + 1
 }
-
-func isIntRegister(r int16) bool {
-	return unreservedGeneralPurposeIntRegisters[0] <= r && r <= unreservedGeneralPurposeIntRegisters[len(unreservedGeneralPurposeIntRegisters)-1]
-}
-
-func isFloatRegister(r int16) bool {
-	return generalPurposeFloatRegisters[0] <= r && r <= generalPurposeFloatRegisters[len(generalPurposeFloatRegisters)-1]
-}
-
-func isZeroRegister(r int16) bool {
-	return r == zeroRegister
-}


Test patch:
diff --git a/internal/wasm/jit/engine_test.go b/internal/wasm/jit/engine_test.go
index 61bf8572..8a042b85 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -67,6 +67,9 @@ func TestVerifyOffsetValue(t *testing.T) {
 	// Offsets for wasm.TableInstance.
 	var tableInstance wasm.TableInstance
 	require.Equal(t, int(unsafe.Offsetof(tableInstance.Table)), tableInstanceTableOffset)
+	// We add "+8" to get the length of Tables[0].Table
+	// since the slice header is laid out as {Data uintptr, Len int64, Cap int64} on memory.
+	require.Equal(t, int(unsafe.Offsetof(tableInstance.Table)+8), tableInstanceTableLenOffset)
 
 	// Offsets for wasm.MemoryInstance
 	var memoryInstance wasm.MemoryInstance
diff --git a/internal/wasm/jit/jit_amd64_test.go b/internal/wasm/jit/jit_amd64_test.go
index 50816e8e..dfd99c57 100644
--- a/internal/wasm/jit/jit_amd64_test.go
+++ b/internal/wasm/jit/jit_amd64_test.go
@@ -6153,10 +6153,6 @@ func TestAmd64Compiler_compileCall(t *testing.T) {
 }
 
 func TestAmd64Compiler_compileCallIndirect(t *testing.T) {
-	// Ensure that the offset of wasm.TableInstance doesn't drift.
-	require.Equal(t, int(unsafe.Offsetof((&wasm.TableElement{}).FunctionAddress)), tableElementFunctionAddressOffset)
-	require.Equal(t, int(unsafe.Offsetof((&wasm.TableElement{}).FunctionTypeID)), tableElementFunctionTypeIDOffset)
-
 	t.Run("out of bounds", func(t *testing.T) {
 		env := newJITEnvironment()
 		env.setTable(make([]wasm.TableElement, 10))
diff --git a/internal/wasm/jit/jit_arm64_test.go b/internal/wasm/jit/jit_arm64_test.go
index 304439ab..78447a47 100644
--- a/internal/wasm/jit/jit_arm64_test.go
+++ b/internal/wasm/jit/jit_arm64_test.go
@@ -1719,10 +1719,13 @@ func TestArm64Compiler_compileCall(t *testing.T) {
 
 				code, _, _, err := compiler.compile()
 				require.NoError(t, err)
-				engine.addCompiledFunction(wasm.FunctionAddress(i), &compiledFunction{
+				addr := wasm.FunctionAddress(i)
+				engine.addCompiledFunction(addr, &compiledFunction{
 					codeSegment:        code,
 					codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
 				})
+				env.moduleInstance.Functions = append(env.moduleInstance.Functions,
+					&wasm.FunctionInstance{FunctionType: &wasm.TypeInstance{Type: targetFunctionType}, Address: addr})
 			}
 
 			// Now we start building the caller's code.
@@ -1739,7 +1742,7 @@ func TestArm64Compiler_compileCall(t *testing.T) {
 
 			// Call all the built functions.
 			for i := 0; i < numCalls; i++ {
-				err = compiler.compileCallImpl(wasm.FunctionAddress(i), targetFunctionType)
+				err = compiler.compileCall(&wazeroir.OperationCall{FunctionIndex: uint32(i)})
 				require.NoError(t, err)
 			}
 
@@ -1748,7 +1751,6 @@ func TestArm64Compiler_compileCall(t *testing.T) {
 
 			code, _, _, err := compiler.compile()
 			require.NoError(t, err)
-
 			env.exec(code)
 
 			if growCallFrameStack {
@@ -1771,6 +1773,211 @@ func TestArm64Compiler_compileCall(t *testing.T) {
 	}
 }
 
+func TestArm64Compiler_compileCallIndirect(t *testing.T) {
+	t.Run("out of bounds", func(t *testing.T) {
+		env := newJITEnvironment()
+		env.setTable(make([]wasm.TableElement, 10))
+		compiler := env.requireNewCompiler(t)
+		err := compiler.compilePreamble()
+		require.NoError(t, err)
+
+		targetOperation := &wazeroir.OperationCallIndirect{}
+		// Ensure that the module instance has the type information for targetOperation.TypeIndex.
+		compiler.f = &wasm.FunctionInstance{
+			FunctionKind:   wasm.FunctionKindWasm,
+			ModuleInstance: &wasm.ModuleInstance{Types: []*wasm.TypeInstance{{Type: &wasm.FunctionType{}}}},
+		}
+
+		// Place the offfset value.
+		err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: 10})
+		require.NoError(t, err)
+
+		err = compiler.compileCallIndirect(targetOperation)
+		require.NoError(t, err)
+
+		// We expect to exit from the code in callIndirect so the subsequet code must be unreachable.
+		err = compiler.compileExitFromNativeCode(jitCallStatusCodeUnreachable)
+		require.NoError(t, err)
+
+		// Generate the code under test and run.
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+		env.exec(code)
+
+		require.Equal(t, jitCallStatusCodeInvalidTableAccess, env.jitStatus())
+	})
+
+	t.Run("uninitialized", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+		err := compiler.compilePreamble()
+		require.NoError(t, err)
+
+		targetOperation := &wazeroir.OperationCallIndirect{}
+		targetOffset := &wazeroir.OperationConstI32{Value: uint32(0)}
+		// Ensure that the module instance has the type information for targetOperation.TypeIndex,
+		compiler.f = &wasm.FunctionInstance{
+			ModuleInstance: &wasm.ModuleInstance{Types: []*wasm.TypeInstance{{Type: &wasm.FunctionType{}}}},
+			FunctionKind:   wasm.FunctionKindWasm,
+		}
+		// and the typeID doesn't match the table[targetOffset]'s type ID.
+		table := make([]wasm.TableElement, 10)
+		env.setTable(table)
+		table[0] = wasm.TableElement{FunctionTypeID: wasm.UninitializedTableElementTypeID}
+
+		// Place the offset value.
+		err = compiler.compileConstI32(targetOffset)
+		require.NoError(t, err)
+		err = compiler.compileCallIndirect(targetOperation)
+		require.NoError(t, err)
+
+		// We expect to exit from the code in callIndirect so the subsequet code must be unreachable.
+		err = compiler.compileExitFromNativeCode(jitCallStatusCodeUnreachable)
+		require.NoError(t, err)
+
+		// Generate the code under test and run.
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+		env.exec(code)
+
+		require.Equal(t, jitCallStatusCodeInvalidTableAccess, env.jitStatus())
+	})
+
+	t.Run("type not match", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+		err := compiler.compilePreamble()
+		require.NoError(t, err)
+
+		targetOperation := &wazeroir.OperationCallIndirect{}
+		targetOffset := &wazeroir.OperationConstI32{Value: uint32(0)}
+		env.moduleInstance.Types = []*wasm.TypeInstance{{Type: &wasm.FunctionType{}, TypeID: 1000}}
+		// Ensure that the module instance has the type information for targetOperation.TypeIndex,
+		// and the typeID doesn't match the table[targetOffset]'s type ID.
+		table := make([]wasm.TableElement, 10)
+		env.setTable(table)
+		table[0] = wasm.TableElement{FunctionTypeID: 50}
+
+		// Place the offfset value.
+		err = compiler.compileConstI32(targetOffset)
+		require.NoError(t, err)
+
+		// Now emit the code.
+		require.NoError(t, compiler.compileCallIndirect(targetOperation))
+
+		// We expect to exit from the code in callIndirect so the subsequet code must be unreachable.
+		err = compiler.compileExitFromNativeCode(jitCallStatusCodeUnreachable)
+		require.NoError(t, err)
+
+		// Generate the code under test and run.
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+		env.exec(code)
+
+		require.Equal(t, jitCallStatusCodeTypeMismatchOnIndirectCall, env.jitStatus())
+	})
+
+	t.Run("ok", func(t *testing.T) {
+		for _, growCallFrameStack := range []bool{false, true} {
+			growCallFrameStack := growCallFrameStack
+			t.Run(fmt.Sprintf("grow=%v", growCallFrameStack), func(t *testing.T) {
+				targetType := &wasm.FunctionType{
+					Params:  []wasm.ValueType{},
+					Results: []wasm.ValueType{wasm.ValueTypeI32}}
+				targetTypeID := wasm.FunctionTypeID(10) // Arbitrary number is fine for testing.
+				operation := &wazeroir.OperationCallIndirect{TypeIndex: 0}
+
+				// Ensure that the module instance has the type information for targetOperation.TypeIndex,
+				// and the typeID  matches the table[targetOffset]'s type ID.
+				moduleInstance := &wasm.ModuleInstance{Types: make([]*wasm.TypeInstance, 100)}
+				moduleInstance.Types[operation.TableIndex] = &wasm.TypeInstance{Type: targetType, TypeID: targetTypeID}
+
+				table := make([]wasm.TableElement, 10)
+				for i := 0; i < len(table); i++ {
+					table[i] = wasm.TableElement{FunctionAddress: wasm.FunctionAddress(i), FunctionTypeID: targetTypeID}
+				}
+
+				for i := 0; i < len(table); i++ {
+					t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+						env := newJITEnvironment()
+						env.setTable(table)
+						engine := env.engine()
+
+						// First we create the call target function with function address = i,
+						// and it returns one value.
+						expectedReturnValue := uint32(i * 1000)
+						{
+							compiler := env.requireNewCompiler(t)
+							err := compiler.compilePreamble()
+							require.NoError(t, err)
+							err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: expectedReturnValue})
+							require.NoError(t, err)
+							err = compiler.compileReturnFunction()
+							require.NoError(t, err)
+
+							code, _, _, err := compiler.compile()
+							require.NoError(t, err)
+
+							cf := &compiledFunction{
+								codeSegment:        code,
+								codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
+							}
+							engine.addCompiledFunction(table[i].FunctionAddress, cf)
+						}
+
+						if growCallFrameStack {
+							env.setCallFrameStackPointer(engine.globalContext.callFrameStackLen - 1)
+							env.setPreviousCallFrameStackPointer(engine.globalContext.callFrameStackLen - 1)
+						}
+
+						compiler := env.requireNewCompiler(t)
+						err := compiler.compilePreamble()
+						require.NoError(t, err)
+
+						compiler.f = &wasm.FunctionInstance{ModuleInstance: moduleInstance, FunctionKind: wasm.FunctionKindWasm}
+
+						// Place the offfset value. Here we try calling a function of functionaddr == table[i].FunctionAddress.
+						err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(i)})
+						require.NoError(t, err)
+
+						// At this point, we should have one item (offset value) on the stack.
+						require.Equal(t, uint64(1), compiler.locationStack.sp)
+
+						require.NoError(t, compiler.compileCallIndirect(operation))
+
+						// At this point, we consumed the offset value, but the function returns one value,
+						// so the stack pointer results in the same.
+						require.Equal(t, uint64(1), compiler.locationStack.sp)
+
+						err = compiler.compileReturnFunction()
+						require.NoError(t, err)
+
+						// Generate the code under test and run.
+						code, _, _, err := compiler.compile()
+						require.NoError(t, err)
+						env.exec(code)
+
+						if growCallFrameStack {
+							// If the call frame stack pointer equals the length of call frame stack length,
+							// we have to call the builtin function to grow the slice.
+							require.Equal(t, jitCallStatusCodeCallBuiltInFunction, env.jitStatus())
+							require.Equal(t, builtinFunctionAddressGrowCallFrameStack, env.functionCallAddress(), env.functionCallAddress())
+
+							// Grow the callFrame stack, and exec again from the return address.
+							env.engine().builtinFunctionGrowCallFrameStack()
+							jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(env.engine())))
+						}
+
+						require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+						require.Equal(t, uint64(1), env.stackPointer())
+						require.Equal(t, expectedReturnValue, env.stackTopAsUint32())
+					})
+				}
+			})
+		}
+	})
+}
+
 func TestArm64Compiler_compileSelect(t *testing.T) {
 	for _, isFloat := range []bool{false, true} {
 		isFloat := isFloat
@@ -1888,27 +2095,44 @@ func TestArm64Compiler_compileModuleContextInitialization(t *testing.T) {
 			moduleInstance: &wasm.ModuleInstance{
 				Globals: []*wasm.GlobalInstance{{Val: 100}},
 				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				// TODO: Add  table
+				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
 			},
 		},
 		{
 			name: "globals nil",
 			moduleInstance: &wasm.ModuleInstance{
 				Memory: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				// TODO: Add  table
+				Tables: []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
 			},
 		},
 		{
 			name: "memory nil",
 			moduleInstance: &wasm.ModuleInstance{
 				Globals: []*wasm.GlobalInstance{{Val: 100}},
-				// TODO: Add  table
+				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
+			},
+		},
+		{
+			name: "table nil",
+			moduleInstance: &wasm.ModuleInstance{
+				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:  []*wasm.TableInstance{{Table: nil}},
+				Globals: []*wasm.GlobalInstance{{Val: 100}},
+			},
+		},
+		{
+			name: "table empty",
+			moduleInstance: &wasm.ModuleInstance{
+				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
+				Globals: []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
 			name: "memory zero length",
 			moduleInstance: &wasm.ModuleInstance{
 				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
 				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 0)},
 			},
 		},
@@ -3833,3 +4057,167 @@ func TestArm64Compiler_compileFConvertFromI(t *testing.T) {
 		})
 	}
 }
+
+func TestAmd64Compiler_compileBrTable(t *testing.T) {
+	requireRunAndExpectedValueReturned := func(t *testing.T, c *arm64Compiler, expValue uint32) {
+		// Emit code for each label which returns the frame ID.
+		for returnValue := uint32(0); returnValue < 7; returnValue++ {
+			label := &wazeroir.Label{Kind: wazeroir.LabelKindHeader, FrameID: returnValue}
+			c.ir.LabelCallers[label.String()] = 1
+			_ = c.compileLabel(&wazeroir.OperationLabel{Label: label})
+			_ = c.compileConstI32(&wazeroir.OperationConstI32{Value: label.FrameID})
+			err := c.compileReturnFunction()
+			require.NoError(t, err)
+		}
+
+		// Generate the code under test and run.
+		env := newJITEnvironment()
+		code, _, _, err := c.compile()
+		require.NoError(t, err)
+		// fmt.Println(hex.EncodeToString(code))
+		env.exec(code)
+
+		// Check the returned value.
+		require.Equal(t, uint64(1), env.stackPointer())
+		require.Equal(t, expValue, env.stackTopAsUint32())
+	}
+
+	getBranchTargetDropFromFrameID := func(frameid uint32) *wazeroir.BranchTargetDrop {
+		return &wazeroir.BranchTargetDrop{Target: &wazeroir.BranchTarget{
+			Label: &wazeroir.Label{FrameID: frameid, Kind: wazeroir.LabelKindHeader}},
+		}
+	}
+
+	for _, tc := range []struct {
+		name          string
+		index         int64
+		o             *wazeroir.OperationBrTable
+		expectedValue uint32
+	}{
+		{
+			name:          "only default with index 0",
+			o:             &wazeroir.OperationBrTable{Default: getBranchTargetDropFromFrameID(6)},
+			index:         0,
+			expectedValue: 6,
+		},
+		{
+			name:          "only default with index 100",
+			o:             &wazeroir.OperationBrTable{Default: getBranchTargetDropFromFrameID(6)},
+			index:         100,
+			expectedValue: 6,
+		},
+		{
+			name: "select default with targets and good index",
+			o: &wazeroir.OperationBrTable{
+				Targets: []*wazeroir.BranchTargetDrop{
+					getBranchTargetDropFromFrameID(1),
+					getBranchTargetDropFromFrameID(2),
+				},
+				Default: getBranchTargetDropFromFrameID(6),
+			},
+			index:         3,
+			expectedValue: 6,
+		},
+		{
+			name: "select default with targets and huge index",
+			o: &wazeroir.OperationBrTable{
+				Targets: []*wazeroir.BranchTargetDrop{
+					getBranchTargetDropFromFrameID(1),
+					getBranchTargetDropFromFrameID(2),
+				},
+				Default: getBranchTargetDropFromFrameID(6),
+			},
+			index:         100000,
+			expectedValue: 6,
+		},
+		{
+			name: "select first with two targets",
+			o: &wazeroir.OperationBrTable{
+				Targets: []*wazeroir.BranchTargetDrop{
+					getBranchTargetDropFromFrameID(1),
+					getBranchTargetDropFromFrameID(2),
+				},
+				Default: getBranchTargetDropFromFrameID(5),
+			},
+			index:         0,
+			expectedValue: 1,
+		},
+		{
+			name: "select last with two targets",
+			o: &wazeroir.OperationBrTable{
+				Targets: []*wazeroir.BranchTargetDrop{
+					getBranchTargetDropFromFrameID(1),
+					getBranchTargetDropFromFrameID(2),
+				},
+				Default: getBranchTargetDropFromFrameID(6),
+			},
+			index:         1,
+			expectedValue: 2,
+		},
+		{
+			name: "select first with five targets",
+			o: &wazeroir.OperationBrTable{
+				Targets: []*wazeroir.BranchTargetDrop{
+					getBranchTargetDropFromFrameID(1),
+					getBranchTargetDropFromFrameID(2),
+					getBranchTargetDropFromFrameID(3),
+					getBranchTargetDropFromFrameID(4),
+					getBranchTargetDropFromFrameID(5),
+				},
+				Default: getBranchTargetDropFromFrameID(5),
+			},
+			index:         0,
+			expectedValue: 1,
+		},
+		{
+			name: "select middle with five targets",
+			o: &wazeroir.OperationBrTable{
+				Targets: []*wazeroir.BranchTargetDrop{
+					getBranchTargetDropFromFrameID(1),
+					getBranchTargetDropFromFrameID(2),
+					getBranchTargetDropFromFrameID(3),
+					getBranchTargetDropFromFrameID(4),
+					getBranchTargetDropFromFrameID(5),
+				},
+				Default: getBranchTargetDropFromFrameID(5),
+			},
+			index:         2,
+			expectedValue: 3,
+		},
+		{
+			name: "select last with five targets",
+			o: &wazeroir.OperationBrTable{
+				Targets: []*wazeroir.BranchTargetDrop{
+					getBranchTargetDropFromFrameID(1),
+					getBranchTargetDropFromFrameID(2),
+					getBranchTargetDropFromFrameID(3),
+					getBranchTargetDropFromFrameID(4),
+					getBranchTargetDropFromFrameID(5),
+				},
+				Default: getBranchTargetDropFromFrameID(5),
+			},
+			index:         4,
+			expectedValue: 5,
+		},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			env := newJITEnvironment()
+			compiler := env.requireNewCompiler(t)
+			compiler.ir = &wazeroir.CompilationResult{LabelCallers: map[string]uint32{}}
+
+			err := compiler.compilePreamble()
+			require.NoError(t, err)
+
+			err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(tc.index)})
+			require.NoError(t, err)
+
+			err = compiler.compileBrTable(tc.o)
+			require.NoError(t, err)
+
+			require.Len(t, compiler.locationStack.usedRegisters, 0)
+
+			requireRunAndExpectedValueReturned(t, compiler, tc.expectedValue)
+		})
+	}
+}