Repository: tetratelabs/wazero
Pull request #1751 (2023-10-03 23:05:06)

wazevo: passes ref_func spectest
#1496 


Patch:
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index f0e992b0..383bc222 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -76,6 +76,8 @@ type (
 		goFunctionCallCalleeModuleContextOpaque uintptr
 		// tableGrowTrampolineAddress holds the address of table grow trampoline function.
 		tableGrowTrampolineAddress *byte
+		// refFuncTrampolineAddress holds the address of ref-func trampoline function.
+		refFuncTrampolineAddress *byte
 	}
 )
 
@@ -348,6 +350,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.ExitCodeRefFunc:
+			mod := c.callerModuleInstance()
+			s := goCallStackView(c.execCtx.stackPointerBeforeGoCall)
+			funcIndex := s[0]
+			ref := mod.Engine.FunctionInstanceReference(wasm.Index(funcIndex))
+			s[0] = uint64(ref)
+			c.execCtx.exitCode = wazevoapi.ExitCodeOK
+			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall)))
 		case wazevoapi.ExitCodeUnreachable:
 			panic(wasmruntime.ErrRuntimeUnreachable)
 		case wazevoapi.ExitCodeMemoryOutOfBounds:
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index 412a2ce7..880c0277 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -54,7 +54,9 @@ type (
 		// stackGrowExecutable is a compiled executable for growing stack builtin function.
 		stackGrowExecutable []byte
 		// tableGrowExecutable is a compiled trampoline executable for table.grow builtin function.
-		tableGrowExecutable       []byte
+		tableGrowExecutable []byte
+		// refFuncExecutable is a compiled trampoline executable for ref.func builtin function.
+		refFuncExecutable         []byte
 		entryPreambles            map[*wasm.FunctionType][]byte
 		listenerBeforeTrampolines map[*wasm.FunctionType][]byte
 		listenerAfterTrampolines  map[*wasm.FunctionType][]byte
@@ -561,7 +563,14 @@ func (e *engine) compileSharedFunctions() {
 		e.sharedFunctions.checkModuleExitCode = mmapExecutable(src)
 	}
 
-	// TODO: table grow, etc.
+	e.be.Init()
+	{
+		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeRefFunc, &ssa.Signature{
+			Params:  []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32 /* function index */},
+			Results: []ssa.Type{ssa.TypeI64}, // returns the function reference.
+		}, false)
+		e.sharedFunctions.refFuncExecutable = mmapExecutable(src)
+	}
 
 	e.be.Init()
 	{
@@ -585,6 +594,9 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	if err := platform.MunmapCodeSegment(sf.tableGrowExecutable); err != nil {
 		panic(err)
 	}
+	if err := platform.MunmapCodeSegment(sf.refFuncExecutable); err != nil {
+		panic(err)
+	}
 	for _, f := range sf.entryPreambles {
 		if err := platform.MunmapCodeSegment(f); err != nil {
 			panic(err)
@@ -605,6 +617,7 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	sf.checkModuleExitCode = nil
 	sf.stackGrowExecutable = nil
 	sf.tableGrowExecutable = nil
+	sf.refFuncExecutable = 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 fe1ba59c..8874cdd5 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -23,6 +23,7 @@ type Compiler struct {
 	memoryGrowSig          ssa.Signature
 	checkModuleExitCodeSig ssa.Signature
 	tableGrowSig           ssa.Signature
+	refFuncSig             ssa.Signature
 	checkModuleExitCodeArg [1]ssa.Value
 	ensureTermination      bool
 
@@ -110,13 +111,20 @@ func (c *Compiler) declareSignatures(listenerOn bool) {
 	c.ssaBuilder.DeclareSignature(&c.checkModuleExitCodeSig)
 
 	c.tableGrowSig = ssa.Signature{
-		ID: c.checkModuleExitCodeSig.ID + 1,
-		// Takes execution context and the page size to grow.
+		ID:     c.checkModuleExitCodeSig.ID + 1,
 		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)
+
+	c.refFuncSig = ssa.Signature{
+		ID:     c.tableGrowSig.ID + 1,
+		Params: []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32 /* func index */},
+		// Returns the function reference.
+		Results: []ssa.Type{ssa.TypeI64},
+	}
+	c.ssaBuilder.DeclareSignature(&c.refFuncSig)
 }
 
 // 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 3c92e1fb..3d81b052 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -2448,6 +2448,30 @@ func (c *Compiler) lowerCurrentOpcode() {
 		default:
 			panic("TODO: unsupported vector instruction: " + wasm.VectorInstructionName(vecOp))
 		}
+	case wasm.OpcodeRefFunc:
+		funcIndex := c.readI32u()
+		if state.unreachable {
+			break
+		}
+
+		c.storeCallerModuleContext()
+
+		funcIndexVal := builder.AllocateInstruction().AsIconst32(funcIndex).Insert(builder).Return()
+
+		refFuncPtr := builder.AllocateInstruction().
+			AsLoad(c.execCtxPtrValue,
+				wazevoapi.ExecutionContextOffsetRefFuncTrampolineAddress.U32(),
+				ssa.TypeI64,
+			).Insert(builder).Return()
+
+		// TODO: reuse the slice.
+		args := []ssa.Value{c.execCtxPtrValue, funcIndexVal}
+		refFuncRet := builder.
+			AllocateInstruction().
+			AsCallIndirect(refFuncPtr, &c.refFuncSig, args).
+			Insert(builder).Return()
+		state.push(refFuncRet)
+
 	case wasm.OpcodeRefNull:
 		c.loweringState.pc++ // skips the reference type as we treat both of them as i64(0).
 		if state.unreachable {
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index 86426ba8..92b9b20c 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -154,6 +154,7 @@ func (m *moduleEngine) NewFunction(index wasm.Index) api.Function {
 	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.execCtx.refFuncTrampolineAddress = &m.parent.sharedFunctions.refFuncExecutable[0]
 	ce.init()
 	return ce
 }
diff --git a/internal/engine/wazevo/wazevoapi/exitcode.go b/internal/engine/wazevo/wazevoapi/exitcode.go
index ff2b7694..fd7c3989 100644
--- a/internal/engine/wazevo/wazevoapi/exitcode.go
+++ b/internal/engine/wazevo/wazevoapi/exitcode.go
@@ -25,6 +25,7 @@ const (
 	ExitCodeCallGoModuleFunctionWithListener
 	ExitCodeCallGoFunctionWithListener
 	ExitCodeTableGrow
+	ExitCodeRefFunc
 	exitCodeMax
 )
 
@@ -71,6 +72,8 @@ func (e ExitCode) String() string {
 		return "grow_memory"
 	case ExitCodeTableGrow:
 		return "table_grow"
+	case ExitCodeRefFunc:
+		return "ref_func"
 	}
 	panic("TODO")
 }
diff --git a/internal/engine/wazevo/wazevoapi/offsetdata.go b/internal/engine/wazevo/wazevoapi/offsetdata.go
index fc7c2aa1..3dfc81bd 100644
--- a/internal/engine/wazevo/wazevoapi/offsetdata.go
+++ b/internal/engine/wazevo/wazevoapi/offsetdata.go
@@ -46,6 +46,8 @@ const (
 	ExecutionContextOffsetGoFunctionCallCalleeModuleContextOpaque Offset = 1120
 	// ExecutionContextOffsetTableGrowTrampolineAddress is an offset of `tableGrowTrampolineAddress` field in wazevo.executionContext
 	ExecutionContextOffsetTableGrowTrampolineAddress Offset = 1128
+	// ExecutionContextOffsetRefFuncTrampolineAddress is an offset of `refFuncTrampolineAddress` field in wazevo.executionContext
+	ExecutionContextOffsetRefFuncTrampolineAddress Offset = 1136
 )
 
 // 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 3026128a..75f0c98b 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -146,6 +146,7 @@ func TestSpectestV2(t *testing.T) {
 		{"loop"},
 		{"ref_null"},
 		{"ref_is_null"},
+		{"ref_func"},
 		{"table_get"},
 		{"table_set"},
 		{"table_size"},
diff --git a/internal/engine/wazevo/engine_test.go b/internal/engine/wazevo/engine_test.go
index d3c365f8..a8aa30b0 100644
--- a/internal/engine/wazevo/engine_test.go
+++ b/internal/engine/wazevo/engine_test.go
@@ -17,19 +17,18 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 
 	b1, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
-
 	b2, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
-
 	b3, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
-
 	b4, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
 	b5, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
 	b6, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
+	b7, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
 
 	preabmles := map[*wasm.FunctionType][]byte{
 		{Params: []wasm.ValueType{}}:                  b4,
@@ -40,6 +39,7 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	sf.stackGrowExecutable = b2
 	sf.checkModuleExitCode = b3
 	sf.tableGrowExecutable = b6
+	sf.refFuncExecutable = b7
 	sf.entryPreambles = preabmles
 
 	sharedFunctionsFinalizer(sf)
@@ -48,6 +48,7 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	require.Nil(t, sf.checkModuleExitCode)
 	require.Nil(t, sf.tableGrowExecutable)
 	require.Nil(t, sf.entryPreambles)
+	require.Nil(t, sf.refFuncExecutable)
 }
 
 func Test_compiledModuleFinalizer(t *testing.T) {
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index be6a9875..6e70d661 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1875,12 +1875,6 @@ func TestCompiler_declareSignatures(t *testing.T) {
 		c.declareSignatures(false)
 
 		declaredSigs := builder.Signatures()
-		require.Equal(t,
-			4+
-				3, // memoryGrowSig and checkModuleExitCodeSig + tableGrowSig.
-			len(declaredSigs),
-		)
-
 		expected := []*ssa.Signature{
 			{ID: 0, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64}},
 			{ID: 1, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI64, ssa.TypeI32}},
@@ -1889,9 +1883,11 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			{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}},
+			{ID: 7, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32}, Results: []ssa.Type{ssa.TypeI64}},
 		}
 
-		for i := 0; i < 7; i++ {
+		require.Equal(t, len(expected), len(declaredSigs))
+		for i := 0; i < len(expected); i++ {
 			require.Equal(t, expected[i].String(), declaredSigs[i].String(), i)
 		}
 	})
@@ -1902,11 +1898,6 @@ func TestCompiler_declareSignatures(t *testing.T) {
 		c.declareSignatures(true)
 
 		declaredSigs := builder.Signatures()
-		require.Equal(t,
-			4*3+
-				3, // memoryGrowSig and checkModuleExitCodeSig + tableGrowSig.
-			len(declaredSigs),
-		)
 
 		expected := []*ssa.Signature{
 			{ID: 0, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64}},
@@ -1927,8 +1918,9 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			{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}},
+			{ID: 15, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32}, Results: []ssa.Type{ssa.TypeI64}},
 		}
-
+		require.Equal(t, len(expected), len(declaredSigs))
 		for i := 0; i < len(declaredSigs); i++ {
 			require.Equal(t, expected[i].String(), declaredSigs[i].String(), i)
 		}
diff --git a/internal/engine/wazevo/wazevo_test.go b/internal/engine/wazevo/wazevo_test.go
index f992c439..e1edc807 100644
--- a/internal/engine/wazevo/wazevo_test.go
+++ b/internal/engine/wazevo/wazevo_test.go
@@ -64,4 +64,5 @@ func Test_ExecutionContextOffsets(t *testing.T) {
 	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)
+	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.refFuncTrampolineAddress)), wazevoapi.ExecutionContextOffsetRefFuncTrampolineAddress)
 }