compiler: stop using map for label infos (#1327)

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

internal/engine/compiler/impl_amd64.go

@@ -89,7 +89,7 @@ type amd64Compiler struct {
 	// and each item is either placed in register or the actual memory stack.
 	locationStack runtimeValueLocationStack
 	// labels hold per wazeroir label specific information in this function.
-	labels map[wazeroir.LabelID]*amd64LabelInfo
+	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.
@@ -106,17 +106,22 @@ func newAmd64Compiler() compiler {
 	return c
 }
 
+// Init implements compiler.Init.
 func (c *amd64Compiler) Init(ir *wazeroir.CompilationResult, withListener bool) {
-	assembler, vstack := c.assembler, c.locationStack
+	assembler, locationStack := c.assembler, c.locationStack
 	assembler.Reset()
-	vstack.reset()
-	*c = amd64Compiler{
-		labels:       map[wazeroir.LabelID]*amd64LabelInfo{},
-		ir:           ir,
-		cpuFeatures:  c.cpuFeatures,
-		withListener: withListener,
+	locationStack.reset()
+	for i := range c.labels {
+		c.labels[i] = c.labels[i][:0]
+	}
+	*c = amd64Compiler{
+		ir:            ir,
+		assembler:     assembler,
+		locationStack: locationStack,
+		cpuFeatures:   c.cpuFeatures,
+		withListener:  withListener,
+		labels:        c.labels,
 	}
-	c.assembler, c.locationStack = assembler, vstack
 }
 
 // runtimeValueLocationStack implements compilerImpl.runtimeValueLocationStack for the amd64 architecture.
@@ -157,12 +162,16 @@ type amd64LabelInfo struct {
 }
 
 func (c *amd64Compiler) label(labelID wazeroir.LabelID) *amd64LabelInfo {
-	ret, ok := c.labels[labelID]
-	if ok {
-		return ret
+	kind := labelID.Kind()
+	frames := c.labels[kind]
+	frameID := labelID.FrameID()
+	// If the frameID is not allocated yet, expand the slice by twice of the diff,
+	// so that we could reduce the allocation in the subsequent compilation.
+	if diff := frameID - len(frames) + 1; diff > 0 {
+		frames = append(frames, make([]amd64LabelInfo, diff*2)...)
+		c.labels[kind] = frames
 	}
-	c.labels[labelID] = &amd64LabelInfo{}
-	return c.labels[labelID]
+	return &frames[frameID]
 }
 
 // compileBuiltinFunctionCheckExitCode implements compiler.compileBuiltinFunctionCheckExitCode for the amd64 architecture.

internal/engine/compiler/impl_arm64.go

@@ -23,7 +23,7 @@ type arm64Compiler struct {
 	// and each item is either placed in register or the actual memory stack.
 	locationStack runtimeValueLocationStack
 	// labels maps a label (e.g. ".L1_then") to *arm64LabelInfo.
-	labels map[wazeroir.LabelID]*arm64LabelInfo
+	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.
@@ -38,12 +38,18 @@ func newArm64Compiler() compiler {
 	}
 }
 
+// Init implements compiler.Init.
 func (c *arm64Compiler) Init(ir *wazeroir.CompilationResult, withListener bool) {
-	assembler, vstack := c.assembler, c.locationStack
+	assembler, locationStack := c.assembler, c.locationStack
 	assembler.Reset()
-	vstack.reset()
-	*c = arm64Compiler{labels: map[wazeroir.LabelID]*arm64LabelInfo{}, ir: ir, withListener: withListener}
-	c.assembler, c.locationStack = assembler, vstack
+	locationStack.reset()
+	for i := range c.labels {
+		c.labels[i] = c.labels[i][:0]
+	}
+	*c = arm64Compiler{
+		assembler: assembler, locationStack: locationStack,
+		ir: ir, withListener: withListener, labels: c.labels,
+	}
 }
 
 var (
@@ -135,13 +141,17 @@ type arm64LabelInfo struct {
 	initialStack runtimeValueLocationStack
 }
 
-func (c *arm64Compiler) label(labelKey wazeroir.LabelID) *arm64LabelInfo {
-	ret, ok := c.labels[labelKey]
-	if ok {
-		return ret
+func (c *arm64Compiler) label(labelID wazeroir.LabelID) *arm64LabelInfo {
+	kind := labelID.Kind()
+	frames := c.labels[kind]
+	frameID := labelID.FrameID()
+	// If the frameID is not allocated yet, expand the slice by twice of the diff,
+	// so that we could reduce the allocation in the subsequent compilation.
+	if diff := frameID - len(frames) + 1; diff > 0 {
+		frames = append(frames, make([]arm64LabelInfo, diff*2)...)
+		c.labels[kind] = frames
 	}
-	c.labels[labelKey] = &arm64LabelInfo{}
-	return c.labels[labelKey]
+	return &frames[frameID]
 }
 
 // runtimeValueLocationStack implements compilerImpl.runtimeValueLocationStack for the amd64 architecture.

internal/wazeroir/operations.go

@@ -819,6 +819,16 @@ type Label struct {
 // LabelID is the unique identifiers for blocks in a single function.
 type LabelID uint64
 
+// Kind returns the LabelKind encoded in this LabelID.
+func (l LabelID) Kind() LabelKind {
+	return LabelKind(uint32(l))
+}
+
+// FrameID returns the frame id encoded in this LabelID.
+func (l LabelID) FrameID() int {
+	return int(uint32(l >> 32))
+}
+
 // ID returns the LabelID for this Label.
 func (l Label) ID() (id LabelID) {
 	id = LabelID(l.Kind) | LabelID(l.FrameID)<<32
@@ -863,6 +873,7 @@ const (
 	// we have the continuation block (of if-block) corresponding to "return" opcode.
 	LabelKindContinuation
 	LabelKindReturn
+	LabelKindNum
 )
 
 func (l Label) asBranchTargetDrop() BranchTargetDrop {

internal/wazeroir/operations_test.go

@@ -65,3 +65,12 @@ func TestUnionOperation_String(t *testing.T) {
 		require.NotEqual(t, "", op.String())
 	}
 }
+
+func TestLabelID(t *testing.T) {
+	for k := LabelKind(0); k < LabelKindNum; k++ {
+		l := Label{Kind: k, FrameID: 12345}
+		id := l.ID()
+		require.Equal(t, k, id.Kind())
+		require.Equal(t, int(l.FrameID), id.FrameID())
+	}
+}