compiler: pass runtimeValueLocationStack by values to reduce allocati…

…ons (#1170)

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

internal/engine/compiler/compiler_stack_test.go

@@ -34,7 +34,7 @@ func TestCompiler_releaseRegisterToStack(t *testing.T) {
 			require.NoError(t, err)
 
 			// Set up the location stack so that we push the const on the specified height.
-			s := &runtimeValueLocationStack{
+			s := runtimeValueLocationStack{
 				sp:                                tc.stackPointer,
 				stack:                             make([]runtimeValueLocation, tc.stackPointer),
 				usedRegisters:                     map[asm.Register]struct{}{},
@@ -51,7 +51,7 @@ func TestCompiler_releaseRegisterToStack(t *testing.T) {
 			}
 			require.NoError(t, err)
 			// Release the register allocated value to the memory stack so that we can see the value after exiting.
-			compiler.compileReleaseRegisterToStack(s.peek())
+			compiler.compileReleaseRegisterToStack(compiler.runtimeValueLocationStack().peek())
 			compiler.compileExitFromNativeCode(nativeCallStatusCodeReturned)
 
 			// Generate the code under test.

internal/engine/compiler/compiler_test.go

@@ -254,7 +254,7 @@ type compilerImpl interface {
 	getOnStackPointerCeilDeterminedCallBack() func(uint64)
 	setStackPointerCeil(uint64)
 	compileReleaseRegisterToStack(loc *runtimeValueLocation)
-	setRuntimeValueLocationStack(*runtimeValueLocationStack)
+	setRuntimeValueLocationStack(runtimeValueLocationStack)
 	compileEnsureOnRegister(loc *runtimeValueLocation) error
 	compileModuleContextInitialization() error
 }

internal/engine/compiler/compiler_value_location.go

@@ -109,8 +109,8 @@ func (v *runtimeValueLocation) String() string {
 	return fmt.Sprintf("{type=%s,location=%s}", v.valueType, location)
 }
 
-func newRuntimeValueLocationStack() *runtimeValueLocationStack {
-	return &runtimeValueLocationStack{
+func newRuntimeValueLocationStack() runtimeValueLocationStack {
+	return runtimeValueLocationStack{
 		stack:                             make([]runtimeValueLocation, 10),
 		usedRegisters:                     map[asm.Register]struct{}{},
 		unreservedVectorRegisters:         unreservedVectorRegisters,
@@ -142,6 +142,10 @@ type runtimeValueLocationStack struct {
 	unreservedGeneralPurposeRegisters, unreservedVectorRegisters []asm.Register
 }
 
+func (v *runtimeValueLocationStack) initialized() bool {
+	return len(v.unreservedGeneralPurposeRegisters) > 0
+}
+
 func (v *runtimeValueLocationStack) reset() {
 	v.stackPointerCeil, v.sp = 0, 0
 	v.stack = v.stack[:0]
@@ -160,8 +164,8 @@ func (v *runtimeValueLocationStack) String() string {
 	return fmt.Sprintf("sp=%d, stack=[%s], used_registers=[%s]", v.sp, strings.Join(stackStr, ","), strings.Join(usedRegisters, ","))
 }
 
-func (v *runtimeValueLocationStack) clone() *runtimeValueLocationStack {
-	ret := &runtimeValueLocationStack{}
+func (v *runtimeValueLocationStack) clone() runtimeValueLocationStack {
+	ret := runtimeValueLocationStack{}
 	ret.sp = v.sp
 	ret.usedRegisters = make(map[asm.Register]struct{}, len(ret.usedRegisters))
 	for r := range v.usedRegisters {

internal/engine/compiler/impl_amd64.go

@@ -87,7 +87,7 @@ type amd64Compiler struct {
 	cpuFeatures platform.CpuFeatureFlags
 	// locationStack holds the state of wazeroir virtual stack.
 	// and each item is either placed in register or the actual memory stack.
-	locationStack *runtimeValueLocationStack
+	locationStack runtimeValueLocationStack
 	// labels hold per wazeroir label specific information in this function.
 	labels map[wazeroir.LabelID]*amd64LabelInfo
 	// stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation.
@@ -121,13 +121,13 @@ func (c *amd64Compiler) Init(ir *wazeroir.CompilationResult, withListener bool)
 
 // runtimeValueLocationStack implements compilerImpl.runtimeValueLocationStack for the amd64 architecture.
 func (c *amd64Compiler) runtimeValueLocationStack() *runtimeValueLocationStack {
-	return c.locationStack
+	return &c.locationStack
 }
 
 // setLocationStack sets the given runtimeValueLocationStack to .locationStack field,
 // while allowing us to track runtimeValueLocationStack.stackPointerCeil across multiple stacks.
 // This is called when we branch into different block.
-func (c *amd64Compiler) setLocationStack(newStack *runtimeValueLocationStack) {
+func (c *amd64Compiler) setLocationStack(newStack runtimeValueLocationStack) {
 	if c.stackPointerCeil < c.locationStack.stackPointerCeil {
 		c.stackPointerCeil = c.locationStack.stackPointerCeil
 	}
@@ -153,7 +153,7 @@ type amd64LabelInfo struct {
 	// initialInstruction is the initial instruction for this label so other block can jump into it.
 	initialInstruction asm.Node
 	// initialStack is the initial value location stack from which we start compiling this label.
-	initialStack *runtimeValueLocationStack
+	initialStack runtimeValueLocationStack
 	// labelBeginningCallbacks holds callbacks should to be called with initialInstruction
 	labelBeginningCallbacks []func(asm.Node)
 }
@@ -400,7 +400,7 @@ func (c *amd64Compiler) branchInto(target wazeroir.Label) error {
 		// with the appropriate value locations. Note we clone the stack here as we maybe
 		// manipulate the stack before compiler reaches the label.
 		targetLabel := c.label(labelID)
-		if targetLabel.initialStack == nil {
+		if !targetLabel.initialStack.initialized() {
 			// It seems unnecessary to clone as branchInto is always the tail of the current block.
 			// TODO: verify ^^.
 			targetLabel.initialStack = c.locationStack.clone()
@@ -497,7 +497,7 @@ func (c *amd64Compiler) compileBrIf(o wazeroir.OperationBrIf) error {
 		// with the appropriate value locations. Note we clone the stack here as we maybe
 		// manipulate the stack before compiler reaches the label.
 		labelInfo := c.label(elseLabelID)
-		if labelInfo.initialStack == nil {
+		if !labelInfo.initialStack.initialized() {
 			labelInfo.initialStack = c.locationStack
 		}
 
@@ -527,7 +527,7 @@ func (c *amd64Compiler) compileBrIf(o wazeroir.OperationBrIf) error {
 		// with the appropriate value locations. Note we clone the stack here as we maybe
 		// manipulate the stack before compiler reaches the label.
 		labelInfo := c.label(thenLabelID)
-		if labelInfo.initialStack == nil {
+		if !labelInfo.initialStack.initialized() {
 			labelInfo.initialStack = c.locationStack
 		}
 		thenJmp := c.assembler.CompileJump(amd64.JMP)
@@ -631,7 +631,7 @@ func (c *amd64Compiler) compileBrTable(o wazeroir.OperationBrTable) error {
 		// Assembler would optimize out this NOP during code generation, so this is harmless.
 		labelInitialInstructions[i] = c.assembler.CompileStandAlone(amd64.NOP)
 
-		var locationStack *runtimeValueLocationStack
+		var locationStack runtimeValueLocationStack
 		var target *wazeroir.BranchTargetDrop
 		if i < len(o.Targets) {
 			target = o.Targets[i]
@@ -674,7 +674,7 @@ func (c *amd64Compiler) compileLabel(o wazeroir.OperationLabel) (skipLabel bool)
 	labelInfo := c.label(labelID)
 
 	// If initialStack is not set, that means this label has never been reached.
-	if labelInfo.initialStack == nil {
+	if !labelInfo.initialStack.initialized() {
 		skipLabel = true
 		return
 	}

internal/engine/compiler/impl_amd64_test.go

@@ -589,6 +589,6 @@ func (c *amd64Compiler) setStackPointerCeil(v uint64) {
 }
 
 // compile implements compilerImpl.setRuntimeValueLocationStack for the amd64 architecture.
-func (c *amd64Compiler) setRuntimeValueLocationStack(s *runtimeValueLocationStack) {
+func (c *amd64Compiler) setRuntimeValueLocationStack(s runtimeValueLocationStack) {
 	c.locationStack = s
 }

internal/engine/compiler/impl_arm64.go

@@ -21,7 +21,7 @@ type arm64Compiler struct {
 	ir        *wazeroir.CompilationResult
 	// locationStack holds the state of wazeroir virtual stack.
 	// and each item is either placed in register or the actual memory stack.
-	locationStack *runtimeValueLocationStack
+	locationStack runtimeValueLocationStack
 	// labels maps a label (e.g. ".L1_then") to *arm64LabelInfo.
 	labels map[wazeroir.LabelID]*arm64LabelInfo
 	// stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation.
@@ -132,7 +132,7 @@ type arm64LabelInfo struct {
 	// initialInstruction is the initial instruction for this label so other block can branch into it.
 	initialInstruction asm.Node
 	// initialStack is the initial value location stack from which we start compiling this label.
-	initialStack *runtimeValueLocationStack
+	initialStack runtimeValueLocationStack
 	// labelBeginningCallbacks holds callbacks should to be called with initialInstruction
 	labelBeginningCallbacks []func(asm.Node)
 }
@@ -148,7 +148,7 @@ func (c *arm64Compiler) label(labelKey wazeroir.LabelID) *arm64LabelInfo {
 
 // runtimeValueLocationStack implements compilerImpl.runtimeValueLocationStack for the amd64 architecture.
 func (c *arm64Compiler) runtimeValueLocationStack() *runtimeValueLocationStack {
-	return c.locationStack
+	return &c.locationStack
 }
 
 // pushRuntimeValueLocationOnRegister implements compiler.pushRuntimeValueLocationOnRegister for arm64.
@@ -412,7 +412,7 @@ func (c *arm64Compiler) compileGoDefinedHostFunction() error {
 // setLocationStack sets the given runtimeValueLocationStack to .locationStack field,
 // while allowing us to track runtimeValueLocationStack.stackPointerCeil across multiple stacks.
 // This is called when we branch into different block.
-func (c *arm64Compiler) setLocationStack(newStack *runtimeValueLocationStack) {
+func (c *arm64Compiler) setLocationStack(newStack runtimeValueLocationStack) {
 	if c.stackPointerCeil < c.locationStack.stackPointerCeil {
 		c.stackPointerCeil = c.locationStack.stackPointerCeil
 	}
@@ -437,7 +437,7 @@ func (c *arm64Compiler) compileLabel(o wazeroir.OperationLabel) (skipThisLabel b
 	labelInfo := c.label(labelKey)
 
 	// If initialStack is not set, that means this label has never been reached.
-	if labelInfo.initialStack == nil {
+	if !labelInfo.initialStack.initialized() {
 		skipThisLabel = true
 		return
 	}
@@ -748,7 +748,7 @@ func (c *arm64Compiler) compileBranchInto(target wazeroir.Label) error {
 		// with the appropriate value locations. Note we clone the stack here as we maybe
 		// manipulate the stack before compiler reaches the label.
 		targetLabel := c.label(labelID)
-		if targetLabel.initialStack == nil {
+		if !targetLabel.initialStack.initialized() {
 			targetLabel.initialStack = c.locationStack.clone()
 		}
 
@@ -876,7 +876,7 @@ func (c *arm64Compiler) compileBrTable(o wazeroir.OperationBrTable) error {
 		init := c.assembler.CompileStandAlone(arm64.NOP)
 		labelInitialInstructions[i] = init
 
-		var locationStack *runtimeValueLocationStack
+		var locationStack runtimeValueLocationStack
 		var target *wazeroir.BranchTargetDrop
 		if i < len(o.Targets) {
 			target = o.Targets[i]

internal/engine/compiler/impl_arm64_test.go

@@ -113,6 +113,6 @@ func (c *arm64Compiler) setStackPointerCeil(v uint64) {
 }
 
 // compile implements compilerImpl.setRuntimeValueLocationStack for the amd64 architecture.
-func (c *arm64Compiler) setRuntimeValueLocationStack(s *runtimeValueLocationStack) {
+func (c *arm64Compiler) setRuntimeValueLocationStack(s runtimeValueLocationStack) {
 	c.locationStack = s
 }