Repository: tetratelabs/wazero
Pull request #1750 (2023-10-03 07:17:32)

wazevo: passes table.size spec test
#1496 


Patch:
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index 401433e4..f0e992b0 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -74,6 +74,8 @@ type (
 		savedRegisters [64][2]uint64
 		// goFunctionCallCalleeModuleContextOpaque is the pointer to the target Go function's moduleContextOpaque.
 		goFunctionCallCalleeModuleContextOpaque uintptr
+		// tableGrowTrampolineAddress holds the address of table grow trampoline function.
+		tableGrowTrampolineAddress *byte
 	}
 )
 
@@ -260,6 +262,14 @@ func (c *callEngine) callWithStack(ctx context.Context, paramResultStack []uint6
 			}
 			c.execCtx.exitCode = wazevoapi.ExitCodeOK
 			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall)))
+		case wazevoapi.ExitCodeTableGrow:
+			mod := c.callerModuleInstance()
+			s := goCallStackView(c.execCtx.stackPointerBeforeGoCall)
+			tableIndex, num, ref := s[0], uint32(s[1]), uintptr(s[2])
+			table := mod.Tables[tableIndex]
+			s[0] = uint64(uint32(int32(table.Grow(num, ref))))
+			c.execCtx.exitCode = wazevoapi.ExitCodeOK
+			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall)))
 		case wazevoapi.ExitCodeCallGoFunction:
 			index := wazevoapi.GoFunctionIndexFromExitCode(ec)
 			f := hostModuleGoFuncFromOpaque[api.GoFunction](index, c.execCtx.goFunctionCallCalleeModuleContextOpaque)
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index 5d42c09e..412a2ce7 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -46,13 +46,15 @@ type (
 	}
 
 	sharedFunctions struct {
-		// memoryGrowExecutable is a compiled executable for memory.grow builtin function.
+		// memoryGrowExecutable is a compiled trampoline executable for memory.grow builtin function.
 		memoryGrowExecutable []byte
-		// checkModuleExitCode is a compiled executable for checking module instance exit code. This
+		// checkModuleExitCode is a compiled trampoline executable for checking module instance exit code. This
 		// is used when ensureTermination is true.
 		checkModuleExitCode []byte
 		// stackGrowExecutable is a compiled executable for growing stack builtin function.
-		stackGrowExecutable       []byte
+		stackGrowExecutable []byte
+		// tableGrowExecutable is a compiled trampoline executable for table.grow builtin function.
+		tableGrowExecutable       []byte
 		entryPreambles            map[*wasm.FunctionType][]byte
 		listenerBeforeTrampolines map[*wasm.FunctionType][]byte
 		listenerAfterTrampolines  map[*wasm.FunctionType][]byte
@@ -535,12 +537,21 @@ func (e *engine) compileSharedFunctions() {
 	e.be.Init()
 	{
 		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeGrowMemory, &ssa.Signature{
-			Params:  []ssa.Type{ssa.TypeI32 /* exec context */, ssa.TypeI32},
+			Params:  []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32},
 			Results: []ssa.Type{ssa.TypeI32},
 		}, false)
 		e.sharedFunctions.memoryGrowExecutable = mmapExecutable(src)
 	}
 
+	e.be.Init()
+	{
+		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeTableGrow, &ssa.Signature{
+			Params:  []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32 /* table index */, ssa.TypeI32 /* num */, ssa.TypeI64 /* ref */},
+			Results: []ssa.Type{ssa.TypeI32},
+		}, false)
+		e.sharedFunctions.tableGrowExecutable = mmapExecutable(src)
+	}
+
 	e.be.Init()
 	{
 		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeCheckModuleExitCode, &ssa.Signature{
@@ -571,6 +582,9 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	if err := platform.MunmapCodeSegment(sf.stackGrowExecutable); err != nil {
 		panic(err)
 	}
+	if err := platform.MunmapCodeSegment(sf.tableGrowExecutable); err != nil {
+		panic(err)
+	}
 	for _, f := range sf.entryPreambles {
 		if err := platform.MunmapCodeSegment(f); err != nil {
 			panic(err)
@@ -590,6 +604,7 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	sf.memoryGrowExecutable = nil
 	sf.checkModuleExitCode = nil
 	sf.stackGrowExecutable = nil
+	sf.tableGrowExecutable = nil
 	sf.entryPreambles = nil
 	sf.listenerBeforeTrampolines = nil
 	sf.listenerAfterTrampolines = nil
diff --git a/internal/engine/wazevo/frontend/frontend.go b/internal/engine/wazevo/frontend/frontend.go
index e03aa2a3..fe1ba59c 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -22,6 +22,7 @@ type Compiler struct {
 	listenerSignatures     map[*wasm.FunctionType][2]*ssa.Signature
 	memoryGrowSig          ssa.Signature
 	checkModuleExitCodeSig ssa.Signature
+	tableGrowSig           ssa.Signature
 	checkModuleExitCodeArg [1]ssa.Value
 	ensureTermination      bool
 
@@ -96,7 +97,7 @@ func (c *Compiler) declareSignatures(listenerOn bool) {
 		ID: begin,
 		// Takes execution context and the page size to grow.
 		Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32},
-		// Returns the new size.
+		// Returns the previous page size.
 		Results: []ssa.Type{ssa.TypeI32},
 	}
 	c.ssaBuilder.DeclareSignature(&c.memoryGrowSig)
@@ -107,6 +108,15 @@ func (c *Compiler) declareSignatures(listenerOn bool) {
 		Params: []ssa.Type{ssa.TypeI64},
 	}
 	c.ssaBuilder.DeclareSignature(&c.checkModuleExitCodeSig)
+
+	c.tableGrowSig = ssa.Signature{
+		ID: c.checkModuleExitCodeSig.ID + 1,
+		// Takes execution context and the page size to grow.
+		Params: []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32 /* table index */, ssa.TypeI32 /* num */, ssa.TypeI64 /* ref */},
+		// Returns the previous size.
+		Results: []ssa.Type{ssa.TypeI32},
+	}
+	c.ssaBuilder.DeclareSignature(&c.tableGrowSig)
 }
 
 // SignatureForWasmFunctionType returns the ssa.Signature for the given wasm.FunctionType.
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index c4a00909..3c92e1fb 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -540,6 +540,52 @@ func (c *Compiler) lowerCurrentOpcode() {
 				true,
 			).Insert(builder).Return()
 			state.push(ret)
+
+		case wasm.OpcodeMiscTableSize:
+			tableIndex := c.readI32u()
+			if state.unreachable {
+				break
+			}
+
+			// Load the table.
+			loadTableInstancePtr := builder.AllocateInstruction()
+			loadTableInstancePtr.AsLoad(c.moduleCtxPtrValue, c.offset.TableOffset(int(tableIndex)).U32(), ssa.TypeI64)
+			builder.InsertInstruction(loadTableInstancePtr)
+			tableInstancePtr := loadTableInstancePtr.Return()
+
+			// Load the table's length.
+			loadTableLen := builder.AllocateInstruction().
+				AsLoad(tableInstancePtr, tableInstanceLenOffset, ssa.TypeI32).
+				Insert(builder)
+			state.push(loadTableLen.Return())
+
+		case wasm.OpcodeMiscTableGrow:
+			tableIndex := c.readI32u()
+			if state.unreachable {
+				break
+			}
+
+			c.storeCallerModuleContext()
+
+			tableIndexVal := builder.AllocateInstruction().AsIconst32(tableIndex).Insert(builder).Return()
+
+			num := state.pop()
+			r := state.pop()
+
+			tableGrowPtr := builder.AllocateInstruction().
+				AsLoad(c.execCtxPtrValue,
+					wazevoapi.ExecutionContextOffsetTableGrowTrampolineAddress.U32(),
+					ssa.TypeI64,
+				).Insert(builder).Return()
+
+			// TODO: reuse the slice.
+			args := []ssa.Value{c.execCtxPtrValue, tableIndexVal, num, r}
+			callGrowRet := builder.
+				AllocateInstruction().
+				AsCallIndirect(tableGrowPtr, &c.tableGrowSig, args).
+				Insert(builder).Return()
+			state.push(callGrowRet)
+
 		default:
 			panic("Unknown MiscOp " + strconv.Itoa(int(miscOpUint)))
 		}
@@ -831,7 +877,7 @@ func (c *Compiler) lowerCurrentOpcode() {
 		c.storeCallerModuleContext()
 
 		pages := state.pop()
-		loadPtr := builder.AllocateInstruction().
+		memoryGrowPtr := builder.AllocateInstruction().
 			AsLoad(c.execCtxPtrValue,
 				wazevoapi.ExecutionContextOffsetMemoryGrowTrampolineAddress.U32(),
 				ssa.TypeI64,
@@ -842,7 +888,7 @@ func (c *Compiler) lowerCurrentOpcode() {
 
 		callGrowRet := builder.
 			AllocateInstruction().
-			AsCallIndirect(loadPtr, &c.memoryGrowSig, args).
+			AsCallIndirect(memoryGrowPtr, &c.memoryGrowSig, args).
 			Insert(builder).Return()
 		state.push(callGrowRet)
 
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index 226181d0..86426ba8 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -153,6 +153,7 @@ func (m *moduleEngine) NewFunction(index wasm.Index) api.Function {
 	ce.execCtx.memoryGrowTrampolineAddress = &m.parent.sharedFunctions.memoryGrowExecutable[0]
 	ce.execCtx.stackGrowCallTrampolineAddress = &m.parent.sharedFunctions.stackGrowExecutable[0]
 	ce.execCtx.checkModuleExitCodeTrampolineAddress = &m.parent.sharedFunctions.checkModuleExitCode[0]
+	ce.execCtx.tableGrowTrampolineAddress = &m.parent.sharedFunctions.tableGrowExecutable[0]
 	ce.init()
 	return ce
 }
diff --git a/internal/engine/wazevo/wazevoapi/exitcode.go b/internal/engine/wazevo/wazevoapi/exitcode.go
index 0ca364db..ff2b7694 100644
--- a/internal/engine/wazevo/wazevoapi/exitcode.go
+++ b/internal/engine/wazevo/wazevoapi/exitcode.go
@@ -24,6 +24,7 @@ const (
 	ExitCodeCallListenerAfter
 	ExitCodeCallGoModuleFunctionWithListener
 	ExitCodeCallGoFunctionWithListener
+	ExitCodeTableGrow
 	exitCodeMax
 )
 
@@ -68,6 +69,8 @@ func (e ExitCode) String() string {
 		return "call_go_function_with_listener"
 	case ExitCodeGrowMemory:
 		return "grow_memory"
+	case ExitCodeTableGrow:
+		return "table_grow"
 	}
 	panic("TODO")
 }
diff --git a/internal/engine/wazevo/wazevoapi/offsetdata.go b/internal/engine/wazevo/wazevoapi/offsetdata.go
index 372c5250..fc7c2aa1 100644
--- a/internal/engine/wazevo/wazevoapi/offsetdata.go
+++ b/internal/engine/wazevo/wazevoapi/offsetdata.go
@@ -44,6 +44,8 @@ const (
 	ExecutionContextOffsetSavedRegistersBegin Offset = 96
 	// ExecutionContextOffsetGoFunctionCallCalleeModuleContextOpaque is an offset of `goFunctionCallCalleeModuleContextOpaque` field in wazevo.executionContext
 	ExecutionContextOffsetGoFunctionCallCalleeModuleContextOpaque Offset = 1120
+	// ExecutionContextOffsetTableGrowTrampolineAddress is an offset of `tableGrowTrampolineAddress` field in wazevo.executionContext
+	ExecutionContextOffsetTableGrowTrampolineAddress Offset = 1128
 )
 
 // ModuleContextOffsetData allows the compilers to get the information about offsets to the fields of wazevo.moduleContextOpaque,


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 853a02fe..3026128a 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -148,6 +148,7 @@ func TestSpectestV2(t *testing.T) {
 		{"ref_is_null"},
 		{"table_get"},
 		{"table_set"},
+		{"table_size"},
 		{"simd_boolean"},
 		{"simd_bit_shift"},
 		{"simd_bitwise"},
diff --git a/internal/engine/wazevo/engine_test.go b/internal/engine/wazevo/engine_test.go
index 8d0781bf..d3c365f8 100644
--- a/internal/engine/wazevo/engine_test.go
+++ b/internal/engine/wazevo/engine_test.go
@@ -28,6 +28,8 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	require.NoError(t, err)
 	b5, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
+	b6, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
 
 	preabmles := map[*wasm.FunctionType][]byte{
 		{Params: []wasm.ValueType{}}:                  b4,
@@ -37,12 +39,14 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	sf.memoryGrowExecutable = b1
 	sf.stackGrowExecutable = b2
 	sf.checkModuleExitCode = b3
+	sf.tableGrowExecutable = b6
 	sf.entryPreambles = preabmles
 
 	sharedFunctionsFinalizer(sf)
 	require.Nil(t, sf.memoryGrowExecutable)
 	require.Nil(t, sf.stackGrowExecutable)
 	require.Nil(t, sf.checkModuleExitCode)
+	require.Nil(t, sf.tableGrowExecutable)
 	require.Nil(t, sf.entryPreambles)
 }
 
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index 73c2c731..be6a9875 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1877,7 +1877,7 @@ func TestCompiler_declareSignatures(t *testing.T) {
 		declaredSigs := builder.Signatures()
 		require.Equal(t,
 			4+
-				2, // memoryGrowSig and checkModuleExitCodeSig.
+				3, // memoryGrowSig and checkModuleExitCodeSig + tableGrowSig.
 			len(declaredSigs),
 		)
 
@@ -1888,9 +1888,10 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			{ID: 3, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI64, ssa.TypeI32}},
 			{ID: 4, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32}, Results: []ssa.Type{ssa.TypeI32}},
 			{ID: 5, Params: []ssa.Type{ssa.TypeI64}},
+			{ID: 6, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, ssa.TypeI32, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
 		}
 
-		for i := 0; i < 6; i++ {
+		for i := 0; i < 7; i++ {
 			require.Equal(t, expected[i].String(), declaredSigs[i].String(), i)
 		}
 	})
@@ -1903,7 +1904,7 @@ func TestCompiler_declareSignatures(t *testing.T) {
 		declaredSigs := builder.Signatures()
 		require.Equal(t,
 			4*3+
-				2, // memoryGrowSig and checkModuleExitCodeSig.
+				3, // memoryGrowSig and checkModuleExitCodeSig + tableGrowSig.
 			len(declaredSigs),
 		)
 
@@ -1925,6 +1926,7 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			// Misc.
 			{ID: 12, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32}, Results: []ssa.Type{ssa.TypeI32}},
 			{ID: 13, Params: []ssa.Type{ssa.TypeI64}},
+			{ID: 14, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, ssa.TypeI32, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
 		}
 
 		for i := 0; i < len(declaredSigs); i++ {
diff --git a/internal/engine/wazevo/wazevo_test.go b/internal/engine/wazevo/wazevo_test.go
index 22cfaa99..f992c439 100644
--- a/internal/engine/wazevo/wazevo_test.go
+++ b/internal/engine/wazevo/wazevo_test.go
@@ -63,4 +63,5 @@ func Test_ExecutionContextOffsets(t *testing.T) {
 		"SavedRegistersBegin must be aligned to 16 bytes")
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.savedRegisters)), wazevoapi.ExecutionContextOffsetSavedRegistersBegin)
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.goFunctionCallCalleeModuleContextOpaque)), wazevoapi.ExecutionContextOffsetGoFunctionCallCalleeModuleContextOpaque)
+	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.tableGrowTrampolineAddress)), wazevoapi.ExecutionContextOffsetTableGrowTrampolineAddress)
 }