table grow and size on arm64

Signed-off-by: Takeshi Yoneda <[email protected]>

internal/integration_test/post1_0/reference-types/spec_test.go

@@ -24,6 +24,8 @@ func TestReferenceTypes_JIT(t *testing.T) {
 	testRefFunc(t, wazero.NewRuntimeConfigJIT)
 	testTableSet(t, wazero.NewRuntimeConfigJIT)
 	testTableGet(t, wazero.NewRuntimeConfigJIT)
+	testTableGrow(t, wazero.NewRuntimeConfigJIT)
+	// testTableFill(t, wazero.NewRuntimeConfigJIT)
 }
 
 func TestReferenceTypes_Interpreter(t *testing.T) {

internal/wasm/jit/compiler.go

@@ -389,4 +389,19 @@ type compiler interface {
 	//
 	// https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/valid/instructions.html#xref-syntax-instructions-syntax-instr-table-mathsf-table-set-x
 	compileTableSet(*wazeroir.OperationTableSet) error
+	// compileTableGrow adds instructions to perform operations corresponding to wasm.OpcodeMiscTableGrow instruction introduced in
+	// wasm.FeatureReferenceTypes.
+	//
+	// https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/valid/instructions.html#xref-syntax-instructions-syntax-instr-table-mathsf-table-grow-x
+	compileTableGrow(*wazeroir.OperationTableGrow) error
+	// compileTableSize adds instructions to perform operations corresponding to wasm.OpcodeMiscTableSize instruction introduced in
+	// wasm.FeatureReferenceTypes.
+	//
+	// https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/valid/instructions.html#xref-syntax-instructions-syntax-instr-table-mathsf-table-size-x
+	compileTableSize(*wazeroir.OperationTableSize) error
+	// compileTableFill adds instructions to perform operations corresponding to wasm.OpcodeMiscTableFill instruction introduced in
+	// wasm.FeatureReferenceTypes.
+	//
+	// https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/valid/instructions.html#xref-syntax-instructions-syntax-instr-table-mathsf-table-fill-x
+	compileTableFill(*wazeroir.OperationTableFill) error
 }

internal/wasm/jit/engine.go

@@ -691,6 +691,7 @@ const (
 	builtinFunctionIndexMemoryGrow wasm.Index = iota
 	builtinFunctionIndexGrowValueStack
 	builtinFunctionIndexGrowCallFrameStack
+	builtinFunctionIndexTableGrow
 	// builtinFunctionIndexBreakPoint is internal (only for wazero developers). Disabled by default.
 	builtinFunctionIndexBreakPoint
 )
@@ -735,13 +736,16 @@ jitentry:
 		case jitCallStatusCodeCallBuiltInFunction:
 			switch ce.exitContext.builtinFunctionCallIndex {
 			case builtinFunctionIndexMemoryGrow:
-				callercode := ce.callFrameTop().function
-				ce.builtinFunctionMemoryGrow(ctx, callercode.source.Module.Memory)
+				callerFunction := ce.callFrameTop().function
+				ce.builtinFunctionMemoryGrow(ctx, callerFunction.source.Module.Memory)
 			case builtinFunctionIndexGrowValueStack:
 				callercode := ce.callFrameTop().function
 				ce.builtinFunctionGrowValueStack(callercode.stackPointerCeil)
 			case builtinFunctionIndexGrowCallFrameStack:
 				ce.builtinFunctionGrowCallFrameStack()
+			case builtinFunctionIndexTableGrow:
+				caller := ce.callFrameTop().function
+				ce.builtinFunctionTableGrow(ctx, caller.source.Module.Tables)
 			}
 			if buildoptions.IsDebugMode {
 				if ce.exitContext.builtinFunctionCallIndex == builtinFunctionIndexBreakPoint {
@@ -801,6 +805,15 @@ func (ce *callEngine) builtinFunctionMemoryGrow(ctx context.Context, mem *wasm.M
 	ce.moduleContext.memoryElement0Address = bufSliceHeader.Data
 }
 
+func (ce *callEngine) builtinFunctionTableGrow(ctx context.Context, tables []*wasm.TableInstance) {
+	tableIndex := ce.popValue()
+	table := tables[tableIndex] // verifed not to be out of range by the func validation at compilation phase.
+	num := ce.popValue()
+	ref := ce.popValue()
+	res := table.Grow(ctx, uint32(num), uintptr(ref))
+	ce.pushValue(uint64(res))
+}
+
 func compileHostFunction(sig *wasm.FunctionType) (*code, error) {
 	compiler, err := newCompiler(&wazeroir.CompilationResult{Signature: sig})
 	if err != nil {
@@ -1012,6 +1025,12 @@ func compileWasmFunction(enabledFeatures wasm.Features, ir *wazeroir.Compilation
 			err = compiler.compileTableGet(o)
 		case *wazeroir.OperationTableSet:
 			err = compiler.compileTableSet(o)
+		case *wazeroir.OperationTableGrow:
+			err = compiler.compileTableGrow(o)
+		case *wazeroir.OperationTableSize:
+			err = compiler.compileTableSize(o)
+		case *wazeroir.OperationTableFill:
+			err = compiler.compileTableFill(o)
 		}
 		if err != nil {
 			return nil, fmt.Errorf("operation %s: %w", op.Kind().String(), err)

internal/wasm/jit/jit_impl_arm64.go

@@ -3561,6 +3561,7 @@ func (c *arm64Compiler) compileTableGet(o *wazeroir.OperationTableGet) error {
 	return nil
 }
 
+// compileTableSet implements compiler.compileTableSet for the arm64 architecture.
 func (c *arm64Compiler) compileTableSet(o *wazeroir.OperationTableSet) error {
 	ref := c.locationStack.pop()
 	if err := c.compileEnsureOnGeneralPurposeRegister(ref); err != nil {
@@ -3620,6 +3621,67 @@ func (c *arm64Compiler) compileTableSet(o *wazeroir.OperationTableSet) error {
 	return nil
 }
 
+// compileTableGrow implements compiler.compileTableGrow for the arm64 architecture.
+func (c *arm64Compiler) compileTableGrow(o *wazeroir.OperationTableGrow) error {
+	c.maybeCompileMoveTopConditionalToFreeGeneralPurposeRegister()
+
+	if err := c.compileConstI32(&wazeroir.OperationConstI32{Value: o.TableIndex}); err != nil {
+		return err
+	}
+
+	// Table grow cannot be done in assembly just like memory grow as it involves with allocation in Go.
+	// Therefore, call out to the built function for this purpose.
+	if err := c.compileCallGoFunction(jitCallStatusCodeCallBuiltInFunction, builtinFunctionIndexTableGrow); err != nil {
+		return err
+	}
+
+	// TableGrow consumes three values (table index, number of items, initial value).
+	for i := 0; i < 3; i++ {
+		c.locationStack.pop()
+	}
+
+	// Then, the previous length was pushed as the result.
+	c.locationStack.pushValueLocationOnStack()
+
+	// After return, we re-initialize reserved registers just like preamble of functions.
+	c.compileReservedStackBasePointerRegisterInitialization()
+	return nil
+}
+
+// compileTableSize implements compiler.compileTableSize for the arm64 architecture.
+func (c *arm64Compiler) compileTableSize(o *wazeroir.OperationTableSize) error {
+	result, err := c.allocateRegister(generalPurposeRegisterTypeInt)
+	if err != nil {
+		return err
+	}
+	c.markRegisterUsed(result)
+
+	// arm64ReservedRegisterForTemporary = &tables[0]
+	c.assembler.CompileMemoryToRegister(arm64.MOVD,
+		arm64ReservedRegisterForCallEngine, callEngineModuleContextTablesElement0AddressOffset,
+		arm64ReservedRegisterForTemporary)
+	// arm64ReservedRegisterForTemporary = [arm64ReservedRegisterForTemporary + TableIndex*8]
+	//                                   = [&tables[0] + TableIndex*sizeOf(*tableInstance)]
+	//                                   = [&tables[TableIndex]] = tables[TableIndex].
+	c.assembler.CompileMemoryToRegister(arm64.MOVD,
+		arm64ReservedRegisterForTemporary, int64(o.TableIndex)*8,
+		arm64ReservedRegisterForTemporary)
+
+	// result = [&tables[TableIndex] + tableInstanceTableLenOffset] = len(tables[TableIndex])
+	c.assembler.CompileMemoryToRegister(arm64.MOVD,
+		arm64ReservedRegisterForTemporary, tableInstanceTableLenOffset,
+		result,
+	)
+
+	c.pushValueLocationOnRegister(result)
+	return nil
+}
+
+// compileTableFill implements compiler.compileTableFill for the arm64 architecture.
+func (c *arm64Compiler) compileTableFill(*wazeroir.OperationTableFill) error {
+	return nil
+}
+
 func (c *arm64Compiler) pushZeroValue() {
 	c.pushValueLocationOnRegister(arm64.REGZERO)
 }