compiler: remove embedding of pointers of jump tables

internal/asm/amd64/assembler.go

@@ -90,9 +90,9 @@ type Assembler interface {
 
 	// CompileStaticConstToRegister adds an instruction where the source operand is asm.StaticConst located in the
 	// memory and the destination is the dstReg.
-	CompileStaticConstToRegister(instruction asm.Instruction, c asm.StaticConst, dstReg asm.Register) error
+	CompileStaticConstToRegister(instruction asm.Instruction, c *asm.StaticConst, dstReg asm.Register) error
 
 	// CompileRegisterToStaticConst adds an instruction where the destination operand is asm.StaticConst located in the
 	// memory and the source is the srcReg.
-	CompileRegisterToStaticConst(instruction asm.Instruction, srcReg asm.Register, c asm.StaticConst) error
+	CompileRegisterToStaticConst(instruction asm.Instruction, srcReg asm.Register, c *asm.StaticConst) error
 }

internal/asm/amd64/impl.go

@@ -39,7 +39,7 @@ type NodeImpl struct {
 	// JumpOrigins hold all the nodes trying to jump into this node. In other words, all the nodes with .JumpTarget == this.
 	JumpOrigins map[*NodeImpl]struct{}
 
-	staticConst asm.StaticConst
+	staticConst *asm.StaticConst
 }
 
 type NodeFlag byte
@@ -224,14 +224,14 @@ type AssemblerImpl struct {
 	// but have it as a field here for testability.
 	MaxDisplacementForConstantPool int
 
-	pool constPool
+	pool *asm.StaticConstPool
 }
 
 // compile-time check to ensure AssemblerImpl implements Assembler.
 var _ Assembler = &AssemblerImpl{}
 
 func NewAssemblerImpl() *AssemblerImpl {
-	return &AssemblerImpl{Buf: bytes.NewBuffer(nil), EnablePadding: true, pool: newConstPool(),
+	return &AssemblerImpl{Buf: bytes.NewBuffer(nil), EnablePadding: true, pool: asm.NewStaticConstPool(),
 		MaxDisplacementForConstantPool: defaultMaxDisplacementForConstantPool}
 }
 

internal/asm/amd64/impl_staticconst.go

@@ -8,69 +8,47 @@ import (
 	"github.com/tetratelabs/wazero/internal/asm"
 )
 
-type constPool struct {
-	firstUseOffsetInBinary *asm.NodeOffsetInBinary
-	consts                 []asm.StaticConst
-	poolSizeInBytes        int
-
-	// offsetFinalizedCallbacks are functions called when the offsets of the
-	// constants in the binary have been determined.
-	offsetFinalizedCallbacks map[string][]func(offsetOfConstInBinary int)
-}
-
-func newConstPool() constPool {
-	return constPool{offsetFinalizedCallbacks: map[string][]func(offsetOfConstInBinary int){}}
-}
-
-func (p *constPool) addConst(c asm.StaticConst) {
-	key := asm.StaticConstKey(c)
-	if _, ok := p.offsetFinalizedCallbacks[key]; !ok {
-		p.consts = append(p.consts, c)
-		p.poolSizeInBytes += len(c)
-		p.offsetFinalizedCallbacks[key] = []func(int){}
-	}
-}
-
 // defaultMaxDisplacementForConstantPool is the maximum displacement allowed for literal move instructions which access
 // the constant pool. This is set as 2 ^30 conservatively while the actual limit is 2^31 since we actually allow this
 // limit plus max(length(c) for c in the pool) so we must ensure that limit is less than 2^31.
 const defaultMaxDisplacementForConstantPool = 1 << 30
 
 func (a *AssemblerImpl) maybeFlushConstants(isEndOfFunction bool) {
-	if a.pool.firstUseOffsetInBinary == nil {
+	if a.pool.FirstUseOffsetInBinary == nil {
 		return
 	}
 
 	if isEndOfFunction ||
 		// If the distance between (the first use in binary) and (end of constant pool) can be larger
 		// than MaxDisplacementForConstantPool, we have to emit the constant pool now, otherwise
 		// a const might be unreachable by a literal move whose maximum offset is +- 2^31.
-		((a.pool.poolSizeInBytes+a.Buf.Len())-int(*a.pool.firstUseOffsetInBinary)) >= a.MaxDisplacementForConstantPool {
+		((a.pool.PoolSizeInBytes+a.Buf.Len())-int(*a.pool.FirstUseOffsetInBinary)) >= a.MaxDisplacementForConstantPool {
 		if !isEndOfFunction {
 			// Adds the jump instruction to skip the constants if this is not the end of function.
 			//
 			// TODO: consider NOP padding for this jump, though this rarely happens as most functions should be
 			// small enough to fit all consts after the end of function.
-			if a.pool.poolSizeInBytes >= math.MaxInt8-2 {
+			if a.pool.PoolSizeInBytes >= math.MaxInt8-2 {
 				// long (near-relative) jump: https://www.felixcloutier.com/x86/jmp
 				a.Buf.WriteByte(0xe9)
-				a.WriteConst(int64(a.pool.poolSizeInBytes), 32)
+				a.WriteConst(int64(a.pool.PoolSizeInBytes), 32)
 			} else {
 				// short jump: https://www.felixcloutier.com/x86/jmp
 				a.Buf.WriteByte(0xeb)
-				a.WriteConst(int64(a.pool.poolSizeInBytes), 8)
+				a.WriteConst(int64(a.pool.PoolSizeInBytes), 8)
 			}
 		}
 
-		for _, c := range a.pool.consts {
-			offset := a.Buf.Len()
-			a.Buf.Write(c)
-			for _, callback := range a.pool.offsetFinalizedCallbacks[asm.StaticConstKey(c)] {
+		for _, c := range a.pool.Consts {
+			offset := uint64(a.Buf.Len())
+			c.OffsetInBinary = offset
+			a.Buf.Write(c.Raw)
+			for _, callback := range c.OffsetFinalizedCallbacks {
 				callback(offset)
 			}
 		}
 
-		a.pool = newConstPool() // reset
+		a.pool = asm.NewStaticConstPool() // reset
 	}
 }
 
@@ -147,7 +125,7 @@ func (a *AssemblerImpl) encodeStaticConstToRegister(n *NodeImpl) (err error) {
 // encodeStaticConstImpl encodes an instruction where mod:r/m points to the memory location of the static constant n.staticConst,
 // and the other operand is the register given at n.SrcReg or n.DstReg.
 func (a *AssemblerImpl) encodeStaticConstImpl(n *NodeImpl, opcode []byte, rex RexPrefix, mandatoryPrefix byte) (err error) {
-	a.pool.addConst(n.staticConst)
+	a.pool.AddConst(n.staticConst, uint64(a.Buf.Len()))
 
 	var reg asm.Register
 	if n.DstReg != asm.NilRegister {
@@ -164,17 +142,12 @@ func (a *AssemblerImpl) encodeStaticConstImpl(n *NodeImpl, opcode []byte, rex Re
 	rexPrefix |= rex
 
 	var inst []byte
-	key := asm.StaticConstKey(n.staticConst)
-	a.pool.offsetFinalizedCallbacks[key] = append(a.pool.offsetFinalizedCallbacks[key],
-		func(offsetOfConstInBinary int) {
-			bin := a.Buf.Bytes()
-			displacement := offsetOfConstInBinary - int(n.OffsetInBinary()) - len(inst)
-			displacementOffsetInInstruction := n.OffsetInBinary() + uint64(len(inst)-4)
-			binary.LittleEndian.PutUint32(bin[displacementOffsetInInstruction:], uint32(int32(displacement)))
-		})
-
-	nodeOffset := uint64(a.Buf.Len())
-	a.pool.firstUseOffsetInBinary = &nodeOffset
+	n.staticConst.OffsetFinalizedCallbacks = append(n.staticConst.OffsetFinalizedCallbacks, func(offsetOfConstInBinary uint64) {
+		bin := a.Buf.Bytes()
+		displacement := int(offsetOfConstInBinary) - int(n.OffsetInBinary()) - len(inst)
+		displacementOffsetInInstruction := n.OffsetInBinary() + uint64(len(inst)-4)
+		binary.LittleEndian.PutUint32(bin[displacementOffsetInInstruction:], uint32(int32(displacement)))
+	})
 
 	// https://wiki.osdev.org/X86-64_Instruction_Encoding#32.2F64-bit_addressing
 	modRM := 0b00_000_101 | // Indicate "[RIP + 32bit displacement]" encoding.
@@ -198,9 +171,9 @@ func (a *AssemblerImpl) encodeStaticConstImpl(n *NodeImpl, opcode []byte, rex Re
 }
 
 // CompileStaticConstToRegister implements Assembler.CompileStaticConstToRegister.
-func (a *AssemblerImpl) CompileStaticConstToRegister(instruction asm.Instruction, c asm.StaticConst, dstReg asm.Register) (err error) {
-	if len(c)%2 != 0 {
-		err = fmt.Errorf("the length of a static constant must be even but was %d", len(c))
+func (a *AssemblerImpl) CompileStaticConstToRegister(instruction asm.Instruction, c *asm.StaticConst, dstReg asm.Register) (err error) {
+	if len(c.Raw)%2 != 0 {
+		err = fmt.Errorf("the length of a static constant must be even but was %d", len(c.Raw))
 		return
 	}
 
@@ -211,9 +184,9 @@ func (a *AssemblerImpl) CompileStaticConstToRegister(instruction asm.Instruction
 }
 
 // CompileRegisterToStaticConst implements Assembler.CompileRegisterToStaticConst.
-func (a *AssemblerImpl) CompileRegisterToStaticConst(instruction asm.Instruction, srcReg asm.Register, c asm.StaticConst) (err error) {
-	if len(c)%2 != 0 {
-		err = fmt.Errorf("the length of a static constant must be even but was %d", len(c))
+func (a *AssemblerImpl) CompileRegisterToStaticConst(instruction asm.Instruction, srcReg asm.Register, c *asm.StaticConst) (err error) {
+	if len(c.Raw)%2 != 0 {
+		err = fmt.Errorf("the length of a static constant must be even but was %d", len(c.Raw))
 		return
 	}
 

internal/asm/amd64/impl_staticconst_test.go

@@ -8,28 +8,14 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
-func TestConstPool_addConst(t *testing.T) {
-	p := newConstPool()
-	cons := []byte{1, 2, 3, 4}
-
-	// Loop twice to ensure that the same constant is cached and not added twice.
-	for i := 0; i < 2; i++ {
-		p.addConst(cons)
-		require.Equal(t, 1, len(p.consts))
-		require.Equal(t, len(cons), p.poolSizeInBytes)
-		_, ok := p.offsetFinalizedCallbacks[asm.StaticConstKey(cons)]
-		require.True(t, ok)
-	}
-}
-
 func TestAssemblerImpl_CompileStaticConstToRegister(t *testing.T) {
 	a := NewAssemblerImpl()
 	t.Run("odd count of bytes", func(t *testing.T) {
-		err := a.CompileStaticConstToRegister(MOVDQU, []byte{1}, RegAX)
+		err := a.CompileStaticConstToRegister(MOVDQU, asm.NewStaticConst([]byte{1}), RegAX)
 		require.Error(t, err)
 	})
 	t.Run("ok", func(t *testing.T) {
-		cons := []byte{1, 2, 3, 4}
+		cons := asm.NewStaticConst([]byte{1, 2, 3, 4})
 		err := a.CompileStaticConstToRegister(MOVDQU, cons, RegAX)
 		require.NoError(t, err)
 		actualNode := a.Current
@@ -43,11 +29,11 @@ func TestAssemblerImpl_CompileStaticConstToRegister(t *testing.T) {
 func TestAssemblerImpl_CompileRegisterToStaticConst(t *testing.T) {
 	a := NewAssemblerImpl()
 	t.Run("odd count of bytes", func(t *testing.T) {
-		err := a.CompileRegisterToStaticConst(MOVDQU, RegAX, []byte{1})
+		err := a.CompileRegisterToStaticConst(MOVDQU, RegAX, asm.NewStaticConst([]byte{1}))
 		require.Error(t, err)
 	})
 	t.Run("ok", func(t *testing.T) {
-		cons := []byte{1, 2, 3, 4}
+		cons := asm.NewStaticConst([]byte{1, 2, 3, 4})
 		err := a.CompileRegisterToStaticConst(MOVDQU, RegAX, cons)
 		require.NoError(t, err)
 		actualNode := a.Current
@@ -73,17 +59,17 @@ func TestAssemblerImpl_maybeFlushConstants(t *testing.T) {
 		endOfFunction           bool
 		dummyBodyBeforeFlush    []byte
 		firstUseOffsetInBinary  uint64
-		consts                  []asm.StaticConst
-		expectedOffsetForConsts []int
+		consts                  [][]byte
+		expectedOffsetForConsts []uint64
 		exp                     []byte
 		maxDisplacement         int
 	}{
 		{
 			name:                    "end of function",
 			endOfFunction:           true,
 			dummyBodyBeforeFlush:    []byte{'?', '?', '?', '?'},
-			consts:                  []asm.StaticConst{{1, 2, 3, 4, 5, 6, 7, 8}, {10, 11, 12, 13}},
-			expectedOffsetForConsts: []int{4, 4 + 8}, // 4 = len(dummyBodyBeforeFlush)
+			consts:                  [][]byte{{1, 2, 3, 4, 5, 6, 7, 8}, {10, 11, 12, 13}},
+			expectedOffsetForConsts: []uint64{4, 4 + 8}, // 4 = len(dummyBodyBeforeFlush)
 			firstUseOffsetInBinary:  0,
 			exp:                     []byte{'?', '?', '?', '?', 1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12, 13},
 			maxDisplacement:         1 << 31, // large displacement will emit the consts at the end of function.
@@ -92,7 +78,7 @@ func TestAssemblerImpl_maybeFlushConstants(t *testing.T) {
 			name:                   "not flush",
 			endOfFunction:          false,
 			dummyBodyBeforeFlush:   []byte{'?', '?', '?', '?'},
-			consts:                 []asm.StaticConst{{1, 2, 3, 4, 5, 6, 7, 8}, {10, 11, 12, 13}},
+			consts:                 [][]byte{{1, 2, 3, 4, 5, 6, 7, 8}, {10, 11, 12, 13}},
 			firstUseOffsetInBinary: 0,
 			exp:                    []byte{'?', '?', '?', '?'},
 			maxDisplacement:        1 << 31, // large displacement will emit the consts at the end of function.
@@ -101,8 +87,8 @@ func TestAssemblerImpl_maybeFlushConstants(t *testing.T) {
 			name:                    "not end of function but flush - short jump",
 			endOfFunction:           false,
 			dummyBodyBeforeFlush:    []byte{'?', '?', '?', '?'},
-			consts:                  []asm.StaticConst{{1, 2, 3, 4, 5, 6, 7, 8}, {10, 11, 12, 13}},
-			expectedOffsetForConsts: []int{4 + 2, 4 + 2 + 8}, // 4 = len(dummyBodyBeforeFlush), 2 = the size of jump
+			consts:                  [][]byte{{1, 2, 3, 4, 5, 6, 7, 8}, {10, 11, 12, 13}},
+			expectedOffsetForConsts: []uint64{4 + 2, 4 + 2 + 8}, // 4 = len(dummyBodyBeforeFlush), 2 = the size of jump
 			firstUseOffsetInBinary:  0,
 			exp: []byte{'?', '?', '?', '?',
 				0xeb, 0x0c, // short jump with offset = len(consts[0]) + len(consts[1]) = 12 = 0xc.
@@ -113,8 +99,8 @@ func TestAssemblerImpl_maybeFlushConstants(t *testing.T) {
 			name:                    "not end of function but flush - long jump",
 			endOfFunction:           false,
 			dummyBodyBeforeFlush:    []byte{'?', '?', '?', '?'},
-			consts:                  []asm.StaticConst{largeData},
-			expectedOffsetForConsts: []int{4 + 5}, // 4 = len(dummyBodyBeforeFlush), 5 = the size of jump
+			consts:                  [][]byte{largeData},
+			expectedOffsetForConsts: []uint64{4 + 5}, // 4 = len(dummyBodyBeforeFlush), 5 = the size of jump
 			firstUseOffsetInBinary:  0,
 			exp: append([]byte{'?', '?', '?', '?',
 				0xe9, 0x0, 0x1, 0x0, 0x0, // short jump with offset = 256 = 0x0, 0x1, 0x0, 0x0 (in Little Endian).
@@ -130,15 +116,15 @@ func TestAssemblerImpl_maybeFlushConstants(t *testing.T) {
 			a.Buf.Write(tc.dummyBodyBeforeFlush)
 
 			for i, c := range tc.consts {
-				a.pool.addConst(c)
-				key := asm.StaticConstKey(c)
+				sc := asm.NewStaticConst(c)
+				a.pool.AddConst(sc, 100)
 				i := i
-				a.pool.offsetFinalizedCallbacks[key] = append(a.pool.offsetFinalizedCallbacks[key], func(offsetOfConstInBinary int) {
+				sc.AddOffsetFinalizedCallback(func(offsetOfConstInBinary uint64) {
 					require.Equal(t, tc.expectedOffsetForConsts[i], offsetOfConstInBinary)
 				})
 			}
 
-			a.pool.firstUseOffsetInBinary = &tc.firstUseOffsetInBinary
+			a.pool.FirstUseOffsetInBinary = &tc.firstUseOffsetInBinary
 			a.maybeFlushConstants(tc.endOfFunction)
 
 			require.Equal(t, tc.exp, a.Buf.Bytes())
@@ -150,7 +136,7 @@ func TestAssemblerImpl_encodeRegisterToStaticConst(t *testing.T) {
 	tests := []struct {
 		name            string
 		ins             asm.Instruction
-		c               asm.StaticConst
+		c               []byte
 		reg             asm.Register
 		ud2sBeforeConst int
 		exp             []byte
@@ -220,7 +206,7 @@ func TestAssemblerImpl_encodeRegisterToStaticConst(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			a := NewAssemblerImpl()
 
-			err := a.CompileRegisterToStaticConst(tc.ins, tc.reg, tc.c)
+			err := a.CompileRegisterToStaticConst(tc.ins, tc.reg, asm.NewStaticConst(tc.c))
 			require.NoError(t, err)
 
 			for i := 0; i < tc.ud2sBeforeConst; i++ {
@@ -239,7 +225,7 @@ func TestAssemblerImpl_encodeStaticConstToRegister(t *testing.T) {
 	tests := []struct {
 		name            string
 		ins             asm.Instruction
-		c               asm.StaticConst
+		c               []byte
 		reg             asm.Register
 		ud2sBeforeConst int
 		exp             []byte
@@ -612,7 +598,7 @@ func TestAssemblerImpl_encodeStaticConstToRegister(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			a := NewAssemblerImpl()
 
-			err := a.CompileStaticConstToRegister(tc.ins, tc.c, tc.reg)
+			err := a.CompileStaticConstToRegister(tc.ins, asm.NewStaticConst(tc.c), tc.reg)
 			require.NoError(t, err)
 
 			for i := 0; i < tc.ud2sBeforeConst; i++ {

internal/asm/arm64/assembler.go

@@ -96,9 +96,13 @@ type Assembler interface {
 	CompileVectorRegisterToVectorRegisterWithConst(instruction asm.Instruction, srcReg, dstReg asm.Register,
 		arrangement VectorArrangement, c asm.ConstantValue)
 
-	// CompileLoadStaticConstToVectorRegister adds an instruction where the source operand is StaticConstant located in
+	// CompileStaticConstToRegister adds an instruction where the source operand is StaticConstant located in
 	// the memory and the destination is the dstReg.
-	CompileLoadStaticConstToVectorRegister(instruction asm.Instruction, c asm.StaticConst, dstReg asm.Register,
+	CompileStaticConstToRegister(instruction asm.Instruction, c *asm.StaticConst, dstReg asm.Register)
+
+	// CompileStaticConstToVectorRegister adds an instruction where the source operand is StaticConstant located in
+	// the memory and the destination is the dstReg.
+	CompileStaticConstToVectorRegister(instruction asm.Instruction, c *asm.StaticConst, dstReg asm.Register,
 		arrangement VectorArrangement)
 
 	// CompileTwoVectorRegistersToVectorRegister adds an instruction where source are two vectors and destination is one