diff --git a/internal/engine/compiler/compiler_initialization_test.go b/internal/engine/compiler/compiler_initialization_test.go index 690f8765..c2c153f7 100644 --- a/internal/engine/compiler/compiler_initialization_test.go +++ b/internal/engine/compiler/compiler_initialization_test.go @@ -184,11 +184,9 @@ func TestCompiler_compileMaybeGrowStack(t *testing.T) { err := compiler.compilePreamble() require.NoError(t, err) - require.NotNil(t, compiler.getOnStackPointerCeilDeterminedCallBack()) - stackLen := uint64(len(env.stack())) stackBasePointer := stackLen - baseOffset // Ceil <= stackLen - stackBasePointer = no need to grow! - compiler.getOnStackPointerCeilDeterminedCallBack()(stackPointerCeil) + compiler.assignStackPointerCeil(stackPointerCeil) env.setStackBasePointer(stackBasePointer) compiler.compileExitFromNativeCode(nativeCallStatusCodeReturned) diff --git a/internal/engine/compiler/compiler_test.go b/internal/engine/compiler/compiler_test.go index ce4dcd9b..2cc096e9 100644 --- a/internal/engine/compiler/compiler_test.go +++ b/internal/engine/compiler/compiler_test.go @@ -237,14 +237,14 @@ func (j *compilerEnv) requireNewCompiler(t *testing.T, functionType *wasm.Functi return ret } -// CompilerImpl is the interface used for architecture-independent unit tests in this pkg. +// compilerImpl is the interface used for architecture-independent unit tests in this pkg. // This is currently implemented by amd64 and arm64. type compilerImpl interface { compiler compileExitFromNativeCode(nativeCallStatusCode) compileMaybeGrowStack() error compileReturnFunction() error - getOnStackPointerCeilDeterminedCallBack() func(uint64) + assignStackPointerCeil(uint64) setStackPointerCeil(uint64) compileReleaseRegisterToStack(loc *runtimeValueLocation) setRuntimeValueLocationStack(runtimeValueLocationStack) diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go index 0a903dd1..77c36b2f 100644 --- a/internal/engine/compiler/impl_amd64.go +++ b/internal/engine/compiler/impl_amd64.go @@ -92,10 +92,11 @@ type amd64Compiler struct { labels [wazeroir.LabelKindNum][]amd64LabelInfo // stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation. stackPointerCeil uint64 - // onStackPointerCeilDeterminedCallBack hold a callback which are called when the max stack pointer is determined BEFORE generating native code. - onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64) - withListener bool - typ *wasm.FunctionType + // assignStackPointerCeilNeeded holds an asm.Node whose AssignDestinationConstant must be called with the determined stack pointer ceiling. + assignStackPointerCeilNeeded asm.Node + withListener bool + typ *wasm.FunctionType + br *bytes.Reader } func newAmd64Compiler() compiler { @@ -103,6 +104,7 @@ func newAmd64Compiler() compiler { assembler: amd64.NewAssembler(), locationStack: newRuntimeValueLocationStack(), cpuFeatures: platform.CpuFeatures, + br: bytes.NewReader(nil), } return c } @@ -123,6 +125,7 @@ func (c *amd64Compiler) Init(typ *wasm.FunctionType, ir *wazeroir.CompilationRes withListener: withListener, labels: c.labels, typ: typ, + br: c.br, } } @@ -248,17 +251,15 @@ func (c *amd64Compiler) compile() (code []byte, stackPointerCeil uint64, err err // Now that the max stack pointer is determined, we are invoking the callback. // Note this MUST be called before Assemble() below. - if c.onStackPointerCeilDeterminedCallBack != nil { - c.onStackPointerCeilDeterminedCallBack(stackPointerCeil) - c.onStackPointerCeilDeterminedCallBack = nil - } + c.assignStackPointerCeil(stackPointerCeil) code, err = c.assembler.Assemble() if err != nil { return } - code, err = platform.MmapCodeSegment(bytes.NewReader(code), len(code)) + c.br.Reset(code) + code, err = platform.MmapCodeSegment(c.br, len(code)) return } @@ -268,6 +269,13 @@ func (c *amd64Compiler) compileUnreachable() error { return nil } +// assignStackPointerCeil implements compilerImpl.assignStackPointerCeil for the amd64 architecture. +func (c *amd64Compiler) assignStackPointerCeil(ceil uint64) { + if c.assignStackPointerCeilNeeded != nil { + c.assignStackPointerCeilNeeded.AssignDestinationConstant(int64(ceil) << 3) + } +} + // compileSet implements compiler.compileSet for the amd64 architecture. func (c *amd64Compiler) compileSet(o *wazeroir.UnionOperation) error { depth := int(o.U1) @@ -4949,9 +4957,7 @@ func (c *amd64Compiler) compileMaybeGrowStack() error { // If stack base pointer + max stack pointer > stackLen, we need to grow the stack. cmpWithStackPointerCeil := c.assembler.CompileRegisterToConst(amd64.CMPQ, tmpRegister, 0) - c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) { - cmpWithStackPointerCeil.AssignDestinationConstant(int64(stackPointerCeil) << 3) - } + c.assignStackPointerCeilNeeded = cmpWithStackPointerCeil // Jump if we have no need to grow. jmpIfNoNeedToGrowStack := c.assembler.CompileJump(amd64.JCC) diff --git a/internal/engine/compiler/impl_amd64_test.go b/internal/engine/compiler/impl_amd64_test.go index 4acbf3ee..26a3b52e 100644 --- a/internal/engine/compiler/impl_amd64_test.go +++ b/internal/engine/compiler/impl_amd64_test.go @@ -577,11 +577,6 @@ func collectRegistersFromRuntimeValues(locs []*runtimeValueLocation) []asm.Regis return out } -// compile implements compilerImpl.getOnStackPointerCeilDeterminedCallBack for the amd64 architecture. -func (c *amd64Compiler) getOnStackPointerCeilDeterminedCallBack() func(uint64) { - return c.onStackPointerCeilDeterminedCallBack -} - // compile implements compilerImpl.setStackPointerCeil for the amd64 architecture. func (c *amd64Compiler) setStackPointerCeil(v uint64) { c.stackPointerCeil = v diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go index 745d6d76..24c2d34a 100644 --- a/internal/engine/compiler/impl_arm64.go +++ b/internal/engine/compiler/impl_arm64.go @@ -26,16 +26,18 @@ type arm64Compiler struct { labels [wazeroir.LabelKindNum][]arm64LabelInfo // stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation. stackPointerCeil uint64 - // onStackPointerCeilDeterminedCallBack hold a callback which are called when the ceil of stack pointer is determined before generating native code. - onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64) - withListener bool - typ *wasm.FunctionType + // assignStackPointerCeilNeeded holds an asm.Node whose AssignDestinationConstant must be called with the determined stack pointer ceiling. + assignStackPointerCeilNeeded asm.Node + withListener bool + typ *wasm.FunctionType + br *bytes.Reader } func newArm64Compiler() compiler { return &arm64Compiler{ assembler: arm64.NewAssembler(arm64ReservedRegisterForTemporary), locationStack: newRuntimeValueLocationStack(), + br: bytes.NewReader(nil), } } @@ -51,6 +53,7 @@ func (c *arm64Compiler) Init(typ *wasm.FunctionType, ir *wazeroir.CompilationRes assembler: assembler, locationStack: locationStack, ir: ir, withListener: withListener, labels: c.labels, typ: typ, + br: c.br, } } @@ -121,9 +124,7 @@ func (c *arm64Compiler) compile() (code []byte, stackPointerCeil uint64, err err // Now that the ceil of stack pointer is determined, we are invoking the callback. // Note: this must be called before Assemble() below. - if c.onStackPointerCeilDeterminedCallBack != nil { - c.onStackPointerCeilDeterminedCallBack(stackPointerCeil) - } + c.assignStackPointerCeil(stackPointerCeil) var original []byte original, err = c.assembler.Assemble() @@ -131,7 +132,8 @@ func (c *arm64Compiler) compile() (code []byte, stackPointerCeil uint64, err err return } - code, err = platform.MmapCodeSegment(bytes.NewReader(original), len(original)) + c.br.Reset(original) + code, err = platform.MmapCodeSegment(c.br, len(original)) return } @@ -143,6 +145,13 @@ type arm64LabelInfo struct { initialStack runtimeValueLocationStack } +// assignStackPointerCeil implements compilerImpl.assignStackPointerCeil for the arm64 architecture. +func (c *arm64Compiler) assignStackPointerCeil(ceil uint64) { + if c.assignStackPointerCeilNeeded != nil { + c.assignStackPointerCeilNeeded.AssignSourceConstant(int64(ceil) << 3) + } +} + func (c *arm64Compiler) label(label wazeroir.Label) *arm64LabelInfo { kind := label.Kind() frames := c.labels[kind] @@ -268,9 +277,7 @@ func (c *arm64Compiler) compileMaybeGrowStack() error { ) // At this point of compilation, we don't know the value of stack point ceil, // so we lazily resolve the value later. - c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) { - loadStackPointerCeil.AssignSourceConstant(int64(stackPointerCeil) << 3) - } + c.assignStackPointerCeilNeeded = loadStackPointerCeil // Compare tmpX (len(ce.stack) - ce.stackBasePointer) and tmpY (ce.stackPointerCeil) c.assembler.CompileTwoRegistersToNone(arm64.CMP, tmpX, tmpY) diff --git a/internal/engine/compiler/impl_arm64_test.go b/internal/engine/compiler/impl_arm64_test.go index cab0d2b8..f385fbd6 100644 --- a/internal/engine/compiler/impl_arm64_test.go +++ b/internal/engine/compiler/impl_arm64_test.go @@ -100,11 +100,6 @@ func TestArm64Compiler_readInstructionAddress(t *testing.T) { require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus()) } -// compile implements compilerImpl.getOnStackPointerCeilDeterminedCallBack for the amd64 architecture. -func (c *arm64Compiler) getOnStackPointerCeilDeterminedCallBack() func(uint64) { - return c.onStackPointerCeilDeterminedCallBack -} - // compile implements compilerImpl.setStackPointerCeil for the amd64 architecture. func (c *arm64Compiler) setStackPointerCeil(v uint64) { c.stackPointerCeil = v