Repository: tetratelabs/wazero
Pull request #1637 (2023-08-18 02:24:37)

Adds LookupFunction in experimental/table package

mathetake: PTAL @jerbob92 
jerbob92: @mathetake Looks good to me! Shall I use this branch to change my implementation to so that we can validate it? I'm also not completely sure I always know the params/return type when I'm looking up the function, so have to try that out.
mathetake: Cool that sounds great! Please let me know if you have any issue

Patch:
diff --git a/experimental/table/lookup.go b/experimental/table/lookup.go
new file mode 100644
index 00000000..7f144401
--- /dev/null
+++ b/experimental/table/lookup.go
@@ -0,0 +1,33 @@
+package table
+
+import (
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/wasm"
+)
+
+// LookupFunction tries to get an api.Function from the table instance specified by `tableIndex` and `tableOffset` in the
+// given api.Module. The user of this function must be well aware of the structure of the given api.Module,
+// and the offset and table index must be valid. If this fails to find it, e.g. table is not found,
+// table offset is out of range, violates the expected types, this panics according to the same semantics as
+// call_indirect instruction: https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/exec/instructions.html#xref-syntax-instructions-syntax-instr-control-mathsf-call-indirect-x-y
+//
+//   - `module` is a module instance to look up the function.
+//   - `tableIndex` is the index of the table instance in the module.
+//   - `tableOffset` is the offset of the lookup target in the table.
+//   - `expectedParamTypes` and `expectedResultTypes` are used to check the type of the function found in the table.
+//
+// Note: the returned api.Function is always valid, i.e. not nil, if this returns without panic.
+func LookupFunction(
+	module api.Module, tableIndex uint32, tableOffset uint32,
+	expectedParamTypes, expectedResultTypes []api.ValueType,
+) api.Function {
+	m := module.(*wasm.ModuleInstance)
+	typ := &wasm.FunctionType{Params: expectedParamTypes, Results: expectedResultTypes}
+	typ.CacheNumInUint64()
+	typeID := m.GetFunctionTypeID(typ)
+	if int(tableIndex) >= len(m.Tables) {
+		panic("table index out of range")
+	}
+	table := m.Tables[tableIndex]
+	return m.LookupFunction(table, typeID, tableOffset)
+}
diff --git a/internal/emscripten/emscripten.go b/internal/emscripten/emscripten.go
index e6aa79f0..4492df58 100644
--- a/internal/emscripten/emscripten.go
+++ b/internal/emscripten/emscripten.go
@@ -93,10 +93,7 @@ func (v *InvokeFunc) Call(ctx context.Context, mod api.Module, stack []uint64) {
 	m := mod.(*wasm.ModuleInstance)
 
 	// Lookup the type of the function we are calling indirectly.
-	typeID, err := m.GetFunctionTypeID(v.FunctionType)
-	if err != nil {
-		panic(err)
-	}
+	typeID := m.GetFunctionTypeID(v.FunctionType)
 
 	// This needs copy (not reslice) because the stack is reused for results.
 	// Consider invoke_i (zero arguments, one result): index zero (tableOffset)
@@ -129,7 +126,7 @@ func (v *InvokeFunc) Call(ctx context.Context, mod api.Module, stack []uint64) {
 	var savedStack [2]uint64
 	callOrPanic(ctx, mod, "stackSave", savedStack[:])
 
-	err = f.CallWithStack(ctx, stack)
+	err := f.CallWithStack(ctx, stack)
 	if err != nil {
 		// Module closed: any calls will just fail with the same error.
 		if _, ok := err.(*sys.ExitError); ok {
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 0e752a81..19493457 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -259,9 +259,9 @@ func (c *compiler) AllocateVRegWithSSAType(regType regalloc.RegType, typ ssa.Typ
 
 // Init implements Compiler.Init.
 func (c *compiler) Init(needGoEntryPreamble bool) {
-	for i := regalloc.VRegID(0); i < c.nextVRegID; i++ {
+	for i, v := range c.ssaValueToVRegs {
 		c.ssaValueToVRegs[i] = regalloc.VRegInvalid
-		delete(c.ssaTypeOfVRegID, i)
+		delete(c.ssaTypeOfVRegID, v.ID())
 	}
 	c.currentGID = 0
 	c.nextVRegID = 0
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index cf601b47..5e2126d0 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -50,6 +50,7 @@ var defKinds = [numInstructionKinds]defKind{
 	aluRRRShift:     defKindRD,
 	aluRRImmShift:   defKindRD,
 	aluRRRExtend:    defKindRD,
+	bitRR:           defKindRD,
 	movZ:            defKindRD,
 	movN:            defKindRD,
 	mov32:           defKindRD,
@@ -88,8 +89,7 @@ var defKinds = [numInstructionKinds]defKind{
 	fpuStore128:     defKindNone,
 	udf:             defKindNone,
 	cSel:            defKindRD,
-	fpuCSel32:       defKindRD,
-	fpuCSel64:       defKindRD,
+	fpuCSel:         defKindRD,
 }
 
 // defs returns the list of regalloc.VReg that are defined by the instruction.
@@ -143,6 +143,7 @@ var useKinds = [numInstructionKinds]useKind{
 	aluRRRShift:     useKindRNRM,
 	aluRRImmShift:   useKindRN,
 	aluRRRExtend:    useKindRNRM,
+	bitRR:           useKindRN,
 	movZ:            useKindNone,
 	movN:            useKindNone,
 	mov32:           useKindRN,
@@ -180,8 +181,7 @@ var useKinds = [numInstructionKinds]useKind{
 	loadFpuConst32:  useKindNone,
 	loadFpuConst64:  useKindNone,
 	cSel:            useKindRNRM,
-	fpuCSel32:       useKindRNRM,
-	fpuCSel64:       useKindRNRM,
+	fpuCSel:         useKindRNRM,
 }
 
 // uses returns the list of regalloc.VReg that are used by the instruction.
@@ -464,15 +464,14 @@ func (i *instruction) asCSel(rd, rn, rm operand, c condFlag, _64bit bool) {
 }
 
 func (i *instruction) asFpuCSel(rd, rn, rm operand, c condFlag, _64bit bool) {
-	if _64bit {
-		i.kind = fpuCSel64
-	} else {
-		i.kind = fpuCSel32
-	}
+	i.kind = fpuCSel
 	i.rd = rd
 	i.rn = rn
 	i.rm = rm
 	i.u1 = uint64(c)
+	if _64bit {
+		i.u3 = 1
+	}
 }
 
 func (i *instruction) asBr(target label) {
@@ -612,6 +611,15 @@ func (i *instruction) asALUBitmaskImm(aluOp aluOp, rn, rd regalloc.VReg, imm uin
 	}
 }
 
+func (i *instruction) asBitRR(bitOp bitOp, rd, rn regalloc.VReg, is64bit bool) {
+	i.kind = bitRR
+	i.rn, i.rd = operandNR(rn), operandNR(rd)
+	i.u1 = uint64(bitOp)
+	if is64bit {
+		i.u2 = 1
+	}
+}
+
 func (i *instruction) asFpuRRR(op fpuBinOp, rd, rn, rm operand, dst64bit bool) {
 	i.kind = fpuRRR
 	i.u1 = uint64(op)
@@ -722,7 +730,12 @@ func (i *instruction) String() (str string) {
 			e,
 		)
 	case bitRR:
-		panic("TODO")
+		size := is64SizeBitToSize(i.u2)
+		str = fmt.Sprintf("%s %s, %s",
+			bitOp(i.u1),
+			formatVRegSized(i.rd.nr(), size),
+			formatVRegSized(i.rn.nr(), size),
+		)
 	case uLoad8:
 		str = fmt.Sprintf("ldrb %s, %s", formatVRegSized(i.rd.nr(), 32), i.amode.format(32))
 	case sLoad8:
@@ -839,18 +852,12 @@ func (i *instruction) String() (str string) {
 		panic("TODO")
 	case intToFpu:
 		panic("TODO")
-	case fpuCSel32:
-		str = fmt.Sprintf("fcsel %s, %s, %s, %s",
-			formatVRegSized(i.rd.nr(), 32),
-			formatVRegSized(i.rn.nr(), 32),
-			formatVRegSized(i.rm.nr(), 32),
-			condFlag(i.u1),
-		)
-	case fpuCSel64:
+	case fpuCSel:
+		size := is64SizeBitToSize(i.u3)
 		str = fmt.Sprintf("fcsel %s, %s, %s, %s",
-			formatVRegSized(i.rd.nr(), 64),
-			formatVRegSized(i.rn.nr(), 64),
-			formatVRegSized(i.rm.nr(), 64),
+			formatVRegSized(i.rd.nr(), size),
+			formatVRegSized(i.rn.nr(), size),
+			formatVRegSized(i.rm.nr(), size),
 			condFlag(i.u1),
 		)
 	case fpuRound:
@@ -1066,10 +1073,8 @@ const (
 	fpuToInt
 	// intToFpu represents a conversion from integer to FP.
 	intToFpu
-	// fpuCSel32 represents a 32-bit FP conditional select.
-	fpuCSel32
-	// fpuCSel64 represents a 64-bit FP conditional select.
-	fpuCSel64
+	// fpuCSel represents a 32/64-bit FP conditional select.
+	fpuCSel
 	// fpuRound represents a rounding to integer operation.
 	fpuRound
 	// movToFpu represents a move from a GPR to a scalar FP register.
@@ -1235,6 +1240,28 @@ const (
 	aluOpMSub
 )
 
+// bitOp determines the type of bitwise operation. Instructions whose kind is one of
+// bitOpRbit and bitOpClz would use this type.
+type bitOp int
+
+// String implements fmt.Stringer.
+func (b bitOp) String() string {
+	switch b {
+	case bitOpRbit:
+		return "rbit"
+	case bitOpClz:
+		return "clz"
+	}
+	panic(int(b))
+}
+
+const (
+	// 32/64-bit Rbit.
+	bitOpRbit bitOp = iota
+	// 32/64-bit Clz.
+	bitOpClz
+)
+
 // fpuBinOp represents a binary floating-point unit (FPU) operation.
 type fpuBinOp byte
 
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
index 52baefed..abb2aee7 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
@@ -157,6 +157,13 @@ func (i *instruction) encode(c backend.Compiler) {
 			i.u2,
 			i.u3 == 1,
 		))
+	case bitRR:
+		c.Emit4Bytes(encodeBitRR(
+			bitOp(i.u1),
+			regNumberInEncoding[i.rd.realReg()],
+			regNumberInEncoding[i.rn.realReg()],
+			uint32(i.u2)),
+		)
 	case aluRRImm12:
 		imm12, shift := i.rm.imm12()
 		c.Emit4Bytes(encodeAluRRImm12(
@@ -221,27 +228,21 @@ func (i *instruction) encode(c backend.Compiler) {
 			condFlag(i.u1),
 			i.u3 == 1,
 		))
-	case fpuCSel32:
-		c.Emit4Bytes(encodeFpuCSel(
-			regNumberInEncoding[i.rd.realReg()],
-			regNumberInEncoding[i.rn.realReg()],
-			regNumberInEncoding[i.rm.realReg()],
-			condFlag(i.u1),
-			false,
-		))
-	case fpuCSel64:
+	case fpuCSel:
 		c.Emit4Bytes(encodeFpuCSel(
 			regNumberInEncoding[i.rd.realReg()],
 			regNumberInEncoding[i.rn.realReg()],
 			regNumberInEncoding[i.rm.realReg()],
 			condFlag(i.u1),
-			true,
+			i.u3 == 1,
 		))
 	default:
 		panic(i.String())
 	}
 }
 
+// encodeFpuCSel encodes as "Floating-point conditional select" in
+// https://developer.arm.com/documentation/ddi0596/2020-12/Index-by-Encoding/Data-Processing----Scalar-Floating-Point-and-Advanced-SIMD?lang=en
 func encodeFpuCSel(rd, rn, rm uint32, c condFlag, _64bit bool) uint32 {
 	var ftype uint32
 	if _64bit {
@@ -309,6 +310,21 @@ func encodeAluRRRR(op aluOp, rd, rn, rm, ra, _64bit uint32) uint32 {
 	return _64bit<<31 | 0b11011<<24 | op31<<21 | rm<<16 | oO<<15 | ra<<10 | rn<<5 | rd
 }
 
+// encodeBitRR encodes as Data-processing (1 source) in
+// https://developer.arm.com/documentation/ddi0596/2020-12/Index-by-Encoding/Data-Processing----Register?lang=en
+func encodeBitRR(op bitOp, rd, rn, _64bit uint32) uint32 {
+	var opcode2, opcode uint32
+	switch op {
+	case bitOpRbit:
+		opcode2, opcode = 0b00000, 0b000000
+	case bitOpClz:
+		opcode2, opcode = 0b00000, 0b000100
+	default:
+		panic("TODO/BUG")
+	}
+	return _64bit<<31 | 0b1_0_11010110<<21 | opcode2<<15 | opcode<<10 | rn<<5 | rd
+}
+
 func encodeAsMov32(rn, rd uint32) uint32 {
 	// This is an alias of ORR (shifted register):
 	// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/MOV--register---Move--register---an-alias-of-ORR--shifted-register--
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index f9ccd7ca..55d61c75 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -147,6 +147,14 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeSelect:
 		c, x, y := instr.SelectData()
 		m.lowerSelect(c, x, y, instr.Return())
+	case ssa.OpcodeClz:
+		x := instr.UnaryData()
+		result := instr.Return()
+		m.lowerClz(x, result)
+	case ssa.OpcodeCtz:
+		x := instr.UnaryData()
+		result := instr.Return()
+		m.lowerCtz(x, result)
 	default:
 		panic("TODO: lowering " + instr.Opcode().String())
 	}
@@ -295,6 +303,26 @@ func (m *machine) lowerImul(x, y, result ssa.Value) {
 	m.insert(mul)
 }
 
+func (m *machine) lowerClz(x, result ssa.Value) {
+	rd := m.compiler.VRegOf(result)
+	rn := m.getOperand_NR(m.compiler.ValueDefinition(x), extModeNone)
+	clz := m.allocateInstr()
+	clz.asBitRR(bitOpClz, rd, rn.nr(), x.Type().Bits() == 64)
+	m.insert(clz)
+}
+
+func (m *machine) lowerCtz(x, result ssa.Value) {
+	rd := m.compiler.VRegOf(result)
+	rn := m.getOperand_NR(m.compiler.ValueDefinition(x), extModeNone)
+	rbit := m.allocateInstr()
+	rbit.asBitRR(bitOpRbit, tmpRegVReg, rn.nr(), x.Type().Bits() == 64)
+	m.insert(rbit)
+
+	clz := m.allocateInstr()
+	clz.asBitRR(bitOpClz, rd, tmpRegVReg, x.Type().Bits() == 64)
+	m.insert(clz)
+}
+
 const exitWithCodeEncodingSize = exitSequenceSize + 8
 
 // lowerExitWithCode lowers the lowerExitWithCode takes a context pointer as argument.
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
index cbd82fb0..d16e6324 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
@@ -301,20 +301,25 @@ func (m *machine) getOperand_NR(def *backend.SSAValueDefinition, mode extMode) (
 		}
 	}
 
+	r := v
 	switch inBits := def.SSAValue().Type().Bits(); {
 	case mode == extModeNone:
 	case inBits == 32 && (mode == extModeZeroExtend32 || mode == extModeSignExtend32):
 	case inBits == 32 && mode == extModeZeroExtend64:
+		extended := m.compiler.AllocateVReg(regalloc.RegTypeInt)
 		ext := m.allocateInstr()
-		ext.asExtend(v, v, 32, 64, false)
+		ext.asExtend(extended, v, 32, 64, false)
 		m.insert(ext)
+		r = extended
 	case inBits == 32 && mode == extModeSignExtend64:
+		extended := m.compiler.AllocateVReg(regalloc.RegTypeInt)
 		ext := m.allocateInstr()
-		ext.asExtend(v, v, 32, 64, true)
+		ext.asExtend(extended, v, 32, 64, true)
 		m.insert(ext)
+		r = extended
 	case inBits == 64 && (mode == extModeZeroExtend64 || mode == extModeSignExtend64):
 	}
-	return operandNR(v)
+	return operandNR(r)
 }
 
 func asImm12Operand(val uint64) (op operand, ok bool) {
diff --git a/internal/engine/wazevo/backend/regalloc/reg.go b/internal/engine/wazevo/backend/regalloc/reg.go
index ab7fa404..31e39732 100644
--- a/internal/engine/wazevo/backend/regalloc/reg.go
+++ b/internal/engine/wazevo/backend/regalloc/reg.go
@@ -87,7 +87,7 @@ func (v VReg) String() string {
 	if v.IsRealReg() {
 		return fmt.Sprintf("r%d", v.ID()-vRegIDReservedForRealBegin)
 	}
-	return fmt.Sprintf("r%d?", v.ID())
+	return fmt.Sprintf("v%d?", v.ID())
 }
 
 // RegType represents the type of a register.
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index f00961ee..c0a24448 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -188,6 +188,19 @@ func (a *Allocator) livenessAnalysis(f Function) {
 				if def.IsRealReg() {
 					info.realRegDefs[def] = append(info.realRegDefs[def], pos)
 				} else {
+					if _, ok := info.defs[def]; ok {
+						// This must be a bug in lowering logic in the ISA-specific code.
+						// In short, a virtual register must be defined only once because that's the invariant of
+						// liveness analysis result used in the current implementation. If you reach here, you can either
+						// use "temporary" physical register (like x27 in AArch64), or simply allocate a new virtual
+						// register for the additional definition. In any ways, such a new virtual register doesn't cost
+						// at all in the final code because they are instantly killed and won't interfere with other
+						// virtual registers later on.
+						//
+						// TODO: this should be enforced globally, not just per block. But that can be costly to check.
+						panic(fmt.Sprintf("BUG: multiple definitions found for a virtual register %s", def.String()))
+					}
+
 					info.defs[def] = pos
 					a.vs = append(a.vs, def)
 				}
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index 9445b65b..08d60abe 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -427,6 +427,26 @@ func (c *Compiler) lowerOpcode(op wasm.Opcode) {
 		builder.InsertInstruction(ishl)
 		value := ishl.Return()
 		state.push(value)
+	case wasm.OpcodeI32Clz, wasm.OpcodeI64Clz:
+		if state.unreachable {
+			return
+		}
+		x := state.pop()
+		clz := builder.AllocateInstruction()
+		clz.AsClz(x)
+		builder.InsertInstruction(clz)
+		value := clz.Return()
+		state.push(value)
+	case wasm.OpcodeI32Ctz, wasm.OpcodeI64Ctz:
+		if state.unreachable {
+			return
+		}
+		x := state.pop()
+		ctz := builder.AllocateInstruction()
+		ctz.AsCtz(x)
+		builder.InsertInstruction(ctz)
+		value := ctz.Return()
+		state.push(value)
 	case wasm.OpcodeGlobalGet:
 		index := c.readI32u()
 		if state.unreachable {
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index b7dfb217..130ad435 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -836,6 +836,8 @@ var instructionSideEffects = [opcodeEnd]sideEffect{
 	OpcodeIcmp:               sideEffectFalse,
 	OpcodeFcmp:               sideEffectFalse,
 	OpcodeFadd:               sideEffectFalse,
+	OpcodeClz:                sideEffectFalse,
+	OpcodeCtz:                sideEffectFalse,
 	OpcodeLoad:               sideEffectFalse,
 	OpcodeUload8:             sideEffectFalse,
 	OpcodeUload16:            sideEffectFalse,
@@ -933,6 +935,8 @@ var instructionReturnTypes = [opcodeEnd]returnTypesFn{
 	OpcodeFmin:               returnTypesFnSingle,
 	OpcodeF32const:           returnTypesFnF32,
 	OpcodeF64const:           returnTypesFnF64,
+	OpcodeClz:                returnTypesFnSingle,
+	OpcodeCtz:                returnTypesFnSingle,
 	OpcodeStore:              returnTypesFnNoReturns,
 	OpcodeIstore8:            returnTypesFnNoReturns,
 	OpcodeIstore16:           returnTypesFnNoReturns,
@@ -1302,6 +1306,25 @@ func (i *Instruction) CallIndirectData() (funcPtr Value, sigID SignatureID, args
 	return
 }
 
+// AsClz initializes this instruction as a Count Leading Zeroes instruction with OpcodeClz.
+func (i *Instruction) AsClz(x Value) {
+	i.opcode = OpcodeClz
+	i.v = x
+	i.typ = x.Type()
+}
+
+// AsCtz initializes this instruction as a Count Trailing Zeroes instruction with OpcodeCtz.
+func (i *Instruction) AsCtz(x Value) {
+	i.opcode = OpcodeCtz
+	i.v = x
+	i.typ = x.Type()
+}
+
+// UnaryData return the operand for a unary instruction.
+func (i *Instruction) UnaryData() Value {
+	return i.v
+}
+
 // AsSExtend initializes this instruction as a sign extension instruction with OpcodeSExtend.
 func (i *Instruction) AsSExtend(v Value, from, to byte) {
 	i.opcode = OpcodeSExtend
@@ -1438,6 +1461,8 @@ func (i *Instruction) Format(b Builder) string {
 	case OpcodeIshl, OpcodeSshr, OpcodeUshr:
 		instSuffix = fmt.Sprintf(" %s, %s", i.v.Format(b), i.v2.Format(b))
 	case OpcodeUndefined:
+	case OpcodeClz, OpcodeCtz:
+		instSuffix = " " + i.v.Format(b)
 	default:
 		panic(fmt.Sprintf("TODO: format for %s", i.opcode))
 	}
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index da531307..d45d1e0c 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -153,8 +153,14 @@ type (
 const maximumFunctionTypes = 1 << 27
 
 // GetFunctionTypeID is used by emscripten.
-func (m *ModuleInstance) GetFunctionTypeID(t *FunctionType) (FunctionTypeID, error) {
-	return m.s.GetFunctionTypeID(t)
+func (m *ModuleInstance) GetFunctionTypeID(t *FunctionType) FunctionTypeID {
+	id, err := m.s.GetFunctionTypeID(t)
+	if err != nil {
+		// This is not recoverable in practice since the only error GetFunctionTypeID returns is
+		// when there's too many function types in the store.
+		panic(err)
+	}
+	return id
 }
 
 func (m *ModuleInstance) buildElementInstances(elements []ElementSegment) {


Test patch:
diff --git a/experimental/table/lookup_test.go b/experimental/table/lookup_test.go
new file mode 100644
index 00000000..26825ef5
--- /dev/null
+++ b/experimental/table/lookup_test.go
@@ -0,0 +1,92 @@
+package table_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental/table"
+	"github.com/tetratelabs/wazero/internal/testing/binaryencoding"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+	"github.com/tetratelabs/wazero/internal/wasm"
+)
+
+func TestLookupFunction(t *testing.T) {
+	const i32 = wasm.ValueTypeI32
+	bytes := binaryencoding.EncodeModule(&wasm.Module{
+		TypeSection: []wasm.FunctionType{
+			{Results: []wasm.ValueType{i32}},
+			{Params: []wasm.ValueType{i32, i32}, Results: []wasm.ValueType{i32, i32}},
+		},
+		FunctionSection: []wasm.Index{0, 1},
+		CodeSection: []wasm.Code{
+			{Body: []byte{
+				wasm.OpcodeI32Const, 1,
+				wasm.OpcodeEnd,
+			}},
+			{Body: []byte{
+				// Swap the two i32s params.
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeEnd,
+			}},
+		},
+		ElementSection: []wasm.ElementSegment{
+			{
+				OffsetExpr: wasm.ConstantExpression{Opcode: wasm.OpcodeI32Const, Data: []byte{0}},
+				Init:       []wasm.Index{0, 1},
+				Type:       wasm.RefTypeFuncref,
+			},
+		},
+		TableSection: []wasm.Table{{Type: wasm.RefTypeFuncref, Min: 100}},
+	})
+
+	r := wazero.NewRuntime(context.Background())
+	m, err := r.Instantiate(context.Background(), bytes)
+	require.NoError(t, err)
+	require.NotNil(t, m)
+
+	t.Run("v_i32", func(t *testing.T) {
+		f := table.LookupFunction(m, 0, 0, nil, []api.ValueType{i32})
+		var result [1]uint64
+		err = f.CallWithStack(context.Background(), result[:])
+		require.NoError(t, err)
+		require.Equal(t, uint64(1), result[0])
+	})
+	t.Run("i32i32_i32i32", func(t *testing.T) {
+		f := table.LookupFunction(m, 0, 1, []api.ValueType{i32, i32}, []api.ValueType{i32, i32})
+		stack := [2]uint64{100, 200}
+		err = f.CallWithStack(context.Background(), stack[:])
+		require.NoError(t, err)
+		require.Equal(t, uint64(200), stack[0])
+		require.Equal(t, uint64(100), stack[1])
+	})
+
+	t.Run("panics", func(t *testing.T) {
+		err := require.CapturePanic(func() {
+			table.LookupFunction(m, 0, 2000, nil, []api.ValueType{i32})
+		})
+		require.Equal(t, "invalid table access", err.Error())
+		err = require.CapturePanic(func() {
+			table.LookupFunction(m, 1000, 0, nil, []api.ValueType{i32})
+		})
+		require.Equal(t, "table index out of range", err.Error())
+		err = require.CapturePanic(func() {
+			table.LookupFunction(m, 0, 0, nil, []api.ValueType{api.ValueTypeF32})
+		})
+		require.Equal(t, "indirect call type mismatch", err.Error())
+		err = require.CapturePanic(func() {
+			table.LookupFunction(m, 0, 0, []api.ValueType{i32}, nil)
+		})
+		require.Equal(t, "indirect call type mismatch", err.Error())
+		err = require.CapturePanic(func() {
+			table.LookupFunction(m, 0, 1, []api.ValueType{i32, i32}, nil)
+		})
+		require.Equal(t, "indirect call type mismatch", err.Error())
+		err = require.CapturePanic(func() {
+			table.LookupFunction(m, 0, 1, []api.ValueType{i32, i32}, []api.ValueType{i32, api.ValueTypeF32})
+		})
+		require.Equal(t, "indirect call type mismatch", err.Error())
+	})
+}
diff --git a/internal/engine/wazevo/backend/backend_test.go b/internal/engine/wazevo/backend/backend_test.go
index e36ee2a1..409fff70 100644
--- a/internal/engine/wazevo/backend/backend_test.go
+++ b/internal/engine/wazevo/backend/backend_test.go
@@ -1415,6 +1415,37 @@ L1 (SSA Block: blk0):
 	sxth w6, w8
 	ldr x30, [sp], #0x10
 	ret
+`,
+		},
+		{
+			name: "integer bit counts", m: testcases.IntegerBitCounts.Module,
+			afterLoweringARM64: `
+L1 (SSA Block: blk0):
+	mov x2?, x2
+	mov x3?, x3
+	clz w4?, w2?
+	rbit w27, w2?
+	clz w5?, w27
+	clz x6?, x3?
+	rbit x27, x3?
+	clz x7?, x27
+	mov x3, x7?
+	mov x2, x6?
+	mov x1, x5?
+	mov x0, x4?
+	ret
+`,
+			afterFinalizeARM64: `
+L1 (SSA Block: blk0):
+	str x30, [sp, #-0x10]!
+	clz w0, w2
+	rbit w27, w2
+	clz w1, w27
+	clz x2, x3
+	rbit x27, x3
+	clz x3, x27
+	ldr x30, [sp], #0x10
+	ret
 `,
 		},
 		{
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
index 98774814..e57b7568 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
@@ -434,6 +434,10 @@ func TestInstruction_encode(t *testing.T) {
 		{want: "400041d3", setup: func(i *instruction) {
 			i.asALUShift(aluOpLsl, operandNR(x0VReg), operandNR(x2VReg), operandShiftImm(63), true)
 		}},
+		{want: "4000c05a", setup: func(i *instruction) { i.asBitRR(bitOpRbit, x0VReg, x2VReg, false) }},
+		{want: "4000c0da", setup: func(i *instruction) { i.asBitRR(bitOpRbit, x0VReg, x2VReg, true) }},
+		{want: "4010c05a", setup: func(i *instruction) { i.asBitRR(bitOpClz, x0VReg, x2VReg, false) }},
+		{want: "4010c0da", setup: func(i *instruction) { i.asBitRR(bitOpClz, x0VReg, x2VReg, true) }},
 	} {
 		tc := tc
 		t.Run(tc.want, func(t *testing.T) {
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
index b6d0f2dc..ce7664c5 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
@@ -53,8 +53,8 @@ func TestMachine_getOperand_NR(t *testing.T) {
 				def = &backend.SSAValueDefinition{BlkParamVReg: regToVReg(x4), BlockParamValue: v}
 				return def, extModeZeroExtend64
 			},
-			exp:          operandNR(regToVReg(x4)),
-			instructions: []string{"uxtw x4, w4"},
+			exp:          operandNR(regalloc.VReg(1).SetRegType(regalloc.RegTypeInt)),
+			instructions: []string{"uxtw x1?, w4"},
 		},
 		{
 			name: "block param - sign extend",
@@ -64,8 +64,8 @@ func TestMachine_getOperand_NR(t *testing.T) {
 				def = &backend.SSAValueDefinition{BlkParamVReg: regToVReg(x4), BlockParamValue: v}
 				return def, extModeSignExtend64
 			},
-			exp:          operandNR(regToVReg(x4)),
-			instructions: []string{"sxtw x4, w4"},
+			exp:          operandNR(regalloc.VReg(1).SetRegType(regalloc.RegTypeInt)),
+			instructions: []string{"sxtw x1?, w4"},
 		},
 		{
 			name: "const instr",
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
index 0c3b6aa1..31ca7353 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
@@ -394,10 +394,10 @@ func TestMachine_collectAddends(t *testing.T) {
 						require.True(t, lowerd)
 					}
 					// Param must be zero-extended.
-					require.Equal(t, "uxtw x1000?, w1000?", formatEmittedInstructionsInCurrentBlock(m))
+					require.Equal(t, "uxtw x1?, w1000?", formatEmittedInstructionsInCurrentBlock(m))
 				}
 			},
-			exp64s: []regalloc.VReg{v1000},
+			exp64s: []regalloc.VReg{regalloc.VReg(1).SetRegType(regalloc.RegTypeInt)},
 			offset: 1 + 2 + 3 - 1,
 		},
 		{
diff --git a/internal/engine/wazevo/backend/regalloc/coloring_test.go b/internal/engine/wazevo/backend/regalloc/coloring_test.go
index 965cb0a4..88b02af8 100644
--- a/internal/engine/wazevo/backend/regalloc/coloring_test.go
+++ b/internal/engine/wazevo/backend/regalloc/coloring_test.go
@@ -1,7 +1,6 @@
 package regalloc
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/tetratelabs/wazero/internal/testing/require"
@@ -219,7 +218,6 @@ func TestAllocator_coloringFor(t *testing.T) {
 			var actual []string
 			for _, n := range testNodes {
 				actual = append(actual, n.r.String())
-				fmt.Println(n)
 			}
 			var exp []string
 			for _, r := range tc.expRegs {
diff --git a/internal/engine/wazevo/backend/regalloc/reg_test.go b/internal/engine/wazevo/backend/regalloc/reg_test.go
index 16c67fe0..ba311e61 100644
--- a/internal/engine/wazevo/backend/regalloc/reg_test.go
+++ b/internal/engine/wazevo/backend/regalloc/reg_test.go
@@ -15,8 +15,8 @@ func TestRegTypeOf(t *testing.T) {
 }
 
 func TestVReg_String(t *testing.T) {
-	require.Equal(t, "r0?", VReg(0).String())
-	require.Equal(t, "r100?", VReg(100).String())
+	require.Equal(t, "v0?", VReg(0).String())
+	require.Equal(t, "v100?", VReg(100).String())
 	require.Equal(t, "r5", FromRealReg(5, RegTypeInt).String())
 }
 
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index 6b6014a0..9bf4aed1 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -609,6 +609,17 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64)
 	v9:i32 = SExtend v2, 8->32
 	v10:i32 = SExtend v2, 16->32
 	Jump blk_ret, v4, v5, v6, v7, v8, v9, v10
+`,
+		},
+		{
+			name: "integer bit counts", m: testcases.IntegerBitCounts.Module,
+			exp: `
+blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64)
+	v4:i32 = Clz v2
+	v5:i32 = Ctz v2
+	v6:i64 = Clz v3
+	v7:i64 = Ctz v3
+	Jump blk_ret, v4, v5, v6, v7
 `,
 		},
 		{
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index 28666ca0..6487aeb0 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -731,6 +731,27 @@ var (
 			wasm.OpcodeEnd,
 		}, []wasm.ValueType{}),
 	}
+	IntegerBitCounts = TestCase{
+		Name: "integer_bit_counts",
+		Module: SingleFunctionModule(wasm.FunctionType{
+			Params:  []wasm.ValueType{i32, i64},
+			Results: []wasm.ValueType{i32, i32, i64, i64},
+		}, []byte{
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeI32Clz,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeI32Ctz,
+
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeI64Clz,
+
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeI64Ctz,
+
+			wasm.OpcodeEnd,
+		}, []wasm.ValueType{}),
+	}
 	FloatComparisons = TestCase{
 		Name: "float_comparisons",
 		Module: SingleFunctionModule(wasm.FunctionType{