compiler: avoid alloc with stack pointer ceil and reuse bytes reader (#1344)

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

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)

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)

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.
+	// assignStackPointerCeilNeeded holds an asm.Node whose AssignDestinationConstant must be called with the determined stack pointer ceiling.
-	onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64)
+	assignStackPointerCeilNeeded asm.Node
-	withListener                         bool
+	withListener                 bool
-	typ                                  *wasm.FunctionType
+	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.assignStackPointerCeil(stackPointerCeil)
-		c.onStackPointerCeilDeterminedCallBack(stackPointerCeil)
-		c.onStackPointerCeilDeterminedCallBack = nil
-	}
 
 	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) {
+	c.assignStackPointerCeilNeeded = cmpWithStackPointerCeil
-		cmpWithStackPointerCeil.AssignDestinationConstant(int64(stackPointerCeil) << 3)
-	}
 
 	// Jump if we have no need to grow.
 	jmpIfNoNeedToGrowStack := c.assembler.CompileJump(amd64.JCC)

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

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.
+	// assignStackPointerCeilNeeded holds an asm.Node whose AssignDestinationConstant must be called with the determined stack pointer ceiling.
-	onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64)
+	assignStackPointerCeilNeeded asm.Node
-	withListener                         bool
+	withListener                 bool
-	typ                                  *wasm.FunctionType
+	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.assignStackPointerCeil(stackPointerCeil)
-		c.onStackPointerCeilDeterminedCallBack(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) {
+	c.assignStackPointerCeilNeeded = loadStackPointerCeil
-		loadStackPointerCeil.AssignSourceConstant(int64(stackPointerCeil) << 3)
-	}
 
 	// Compare tmpX (len(ce.stack) - ce.stackBasePointer) and tmpY (ce.stackPointerCeil)
 	c.assembler.CompileTwoRegistersToNone(arm64.CMP, tmpX, tmpY)

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