Repository: tetratelabs/wazero
Pull request #960 (2022-12-26 00:58:51)

asm(arm64): adds support for stack pointer manipulation 


Patch:
diff --git a/internal/asm/arm64/assembler.go b/internal/asm/arm64/assembler.go
index d49130d1..628350d4 100644
--- a/internal/asm/arm64/assembler.go
+++ b/internal/asm/arm64/assembler.go
@@ -32,6 +32,10 @@ type Assembler interface {
 	// constant `srcConst`, and destination operand is unspecified.
 	CompileRegisterAndConstToNone(instruction asm.Instruction, src asm.Register, srcConst asm.ConstantValue)
 
+	// CompileRegisterAndConstToRegister adds an instruction where source operands consist of one register `src` and
+	// constant `srcConst`, and destination operand is a register `dst`.
+	CompileRegisterAndConstToRegister(instruction asm.Instruction, src asm.Register, srcConst asm.ConstantValue, dst asm.Register)
+
 	// CompileLeftShiftedRegisterToRegister adds an instruction where source operand is the "left shifted register"
 	// represented as `srcReg << shiftNum` and the destination is the register `dstReg`.
 	CompileLeftShiftedRegisterToRegister(
diff --git a/internal/asm/arm64/consts.go b/internal/asm/arm64/consts.go
index 689f4519..343b013d 100644
--- a/internal/asm/arm64/consts.go
+++ b/internal/asm/arm64/consts.go
@@ -116,6 +116,8 @@ const (
 	RegR30
 	// RegRZR is the RZR register (read-only, always returning zero)
 	RegRZR
+	// RegSP is the SP register
+	RegSP
 
 	// Scalar floating point registers.
 
@@ -338,6 +340,8 @@ func RegisterName(r asm.Register) string {
 		return "R30"
 	case RegRZR:
 		return "RZR"
+	case RegSP:
+		return "SP"
 	case RegV0:
 		return "V0"
 	case RegV1:
diff --git a/internal/asm/arm64/impl.go b/internal/asm/arm64/impl.go
index 8b1e377b..bbf0063c 100644
--- a/internal/asm/arm64/impl.go
+++ b/internal/asm/arm64/impl.go
@@ -80,6 +80,8 @@ func (n *nodeImpl) String() (ret string) {
 		ret = fmt.Sprintf("%s (%s, %s)", instName, RegisterName(n.srcReg), RegisterName(n.srcReg2))
 	case operandTypesRegisterAndConstToNone:
 		ret = fmt.Sprintf("%s (%s, 0x%x)", instName, RegisterName(n.srcReg), n.srcConst)
+	case operandTypesRegisterAndConstToRegister:
+		ret = fmt.Sprintf("%s (%s, 0x%x), %s", instName, RegisterName(n.srcReg), n.srcConst, RegisterName(n.dstReg))
 	case operandTypesRegisterToMemory:
 		if n.dstReg2 != asm.NilRegister {
 			ret = fmt.Sprintf("%s %s, [%s + %s]", instName, RegisterName(n.srcReg), RegisterName(n.dstReg), RegisterName(n.dstReg2))
@@ -183,6 +185,7 @@ var (
 	operandTypesThreeRegistersToRegister           = operandTypes{operandTypeThreeRegisters, operandTypeRegister}
 	operandTypesTwoRegistersToNone                 = operandTypes{operandTypeTwoRegisters, operandTypeNone}
 	operandTypesRegisterAndConstToNone             = operandTypes{operandTypeRegisterAndConst, operandTypeNone}
+	operandTypesRegisterAndConstToRegister         = operandTypes{operandTypeRegisterAndConst, operandTypeRegister}
 	operandTypesRegisterToMemory                   = operandTypes{operandTypeRegister, operandTypeMemory}
 	operandTypesMemoryToRegister                   = operandTypes{operandTypeMemory, operandTypeRegister}
 	operandTypesConstToRegister                    = operandTypes{operandTypeConst, operandTypeRegister}
@@ -431,7 +434,7 @@ func (a *AssemblerImpl) encodeNode(n *nodeImpl) (err error) {
 		err = a.encodeRegisterToMemory(n)
 	case operandTypesMemoryToRegister:
 		err = a.encodeMemoryToRegister(n)
-	case operandTypesConstToRegister:
+	case operandTypesRegisterAndConstToRegister, operandTypesConstToRegister:
 		err = a.encodeConstToRegister(n)
 	case operandTypesRegisterToVectorRegister:
 		err = a.encodeRegisterToVectorRegister(n)
@@ -586,6 +589,19 @@ func (a *AssemblerImpl) CompileRegisterAndConstToNone(
 	n.srcConst = srcConst
 }
 
+// CompileRegisterAndConstToRegister implements Assembler.CompileRegisterAndConstToRegister
+func (a *AssemblerImpl) CompileRegisterAndConstToRegister(
+	instruction asm.Instruction,
+	src asm.Register,
+	srcConst asm.ConstantValue,
+	dst asm.Register,
+) {
+	n := a.newNode(instruction, operandTypesRegisterAndConstToRegister)
+	n.srcReg = src
+	n.srcConst = srcConst
+	n.dstReg = dst
+}
+
 // CompileLeftShiftedRegisterToRegister implements Assembler.CompileLeftShiftedRegisterToRegister
 func (a *AssemblerImpl) CompileLeftShiftedRegisterToRegister(
 	instruction asm.Instruction,
@@ -1163,8 +1179,20 @@ func (a *AssemblerImpl) encodeRegisterToRegister(n *nodeImpl) (err error) {
 		if err = checkRegisterToRegisterType(n.srcReg, n.dstReg, true, true); err != nil {
 			return
 		}
-
 		srcRegBits, dstRegBits := registerBits(n.srcReg), registerBits(n.dstReg)
+
+		if n.srcReg == RegSP || n.dstReg == RegSP {
+			// Moving between stack pointers.
+			// https://developer.arm.com/documentation/ddi0602/2021-12/Base-Instructions/MOV--to-from-SP---Move-between-register-and-stack-pointer--an-alias-of-ADD--immediate--
+			a.buf.Write([]byte{
+				(srcRegBits << 5) | dstRegBits,
+				srcRegBits >> 3,
+				0x0,
+				0b1001_0001,
+			})
+			return
+		}
+
 		if n.srcReg == RegRZR && inst == MOVD {
 			// If this is 64-bit mov from zero register, then we encode this as MOVK.
 			// See "Move wide (immediate)" in
@@ -2032,10 +2060,10 @@ func getLowestBit(x uint64) uint64 {
 	return x & (^x + 1)
 }
 
-func (a *AssemblerImpl) addOrSub64BitRegisters(sfops byte, src1RegBits byte, src2RegBits byte) {
+func (a *AssemblerImpl) addOrSub64BitRegisters(sfops byte, dstRegBits, src1RegBits, src2RegBits byte) {
 	// src1Reg = src1Reg +/- src2Reg
 	a.buf.Write([]byte{
-		(src1RegBits << 5) | src1RegBits,
+		(src1RegBits << 5) | dstRegBits,
 		src1RegBits >> 3,
 		src2RegBits,
 		sfops<<5 | 0b01011,
@@ -2139,6 +2167,14 @@ func (a *AssemblerImpl) encodeConstToRegister(n *nodeImpl) (err error) {
 	// TODO: refactor and generalize the following like ^ logicalImmediate, etc.
 	switch inst := n.instruction; inst {
 	case ADD, ADDS, SUB, SUBS:
+		srcRegBits := dstRegBits
+		if n.srcReg != asm.NilRegister {
+			srcRegBits, err = intRegisterBits(n.srcReg)
+			if err != nil {
+				return err
+			}
+		}
+
 		var sfops byte
 		if inst == ADD {
 			sfops = 0b100
@@ -2152,7 +2188,7 @@ func (a *AssemblerImpl) encodeConstToRegister(n *nodeImpl) (err error) {
 
 		if c == 0 {
 			// If the constant equals zero, we encode it as ADD (register) with zero register.
-			a.addOrSub64BitRegisters(sfops, dstRegBits, zeroRegisterBits)
+			a.addOrSub64BitRegisters(sfops, dstRegBits, srcRegBits, zeroRegisterBits)
 			return
 		}
 
@@ -2162,16 +2198,16 @@ func (a *AssemblerImpl) encodeConstToRegister(n *nodeImpl) (err error) {
 
 			if c <= 0xfff {
 				a.buf.Write([]byte{
-					(dstRegBits << 5) | dstRegBits,
-					(byte(c) << 2) | (dstRegBits >> 3),
+					(srcRegBits << 5) | dstRegBits,
+					(byte(c) << 2) | (srcRegBits >> 3),
 					byte(c >> 6),
 					sfops<<5 | 0b10001,
 				})
 			} else {
 				c >>= 12
 				a.buf.Write([]byte{
-					(dstRegBits << 5) | dstRegBits,
-					(byte(c) << 2) | (dstRegBits >> 3),
+					(srcRegBits << 5) | dstRegBits,
+					(byte(c) << 2) | (srcRegBits >> 3),
 					0b01<<6 /* shift by 12 */ | byte(c>>6),
 					sfops<<5 | 0b10001,
 				})
@@ -2189,7 +2225,7 @@ func (a *AssemblerImpl) encodeConstToRegister(n *nodeImpl) (err error) {
 			a.load16bitAlignedConst(c>>(16*t), byte(t), tmpRegBits, false, true)
 
 			// ADD/SUB tmpReg, dstReg
-			a.addOrSub64BitRegisters(sfops, dstRegBits, tmpRegBits)
+			a.addOrSub64BitRegisters(sfops, dstRegBits, srcRegBits, tmpRegBits)
 			return
 		} else if t := const16bitAligned(^c); t >= 0 {
 			// Also if the reverse of the const can fit within 16-bit range, do the same ^^.
@@ -2200,7 +2236,7 @@ func (a *AssemblerImpl) encodeConstToRegister(n *nodeImpl) (err error) {
 			a.load16bitAlignedConst(^c>>(16*t), byte(t), tmpRegBits, true, true)
 
 			// ADD/SUB tmpReg, dstReg
-			a.addOrSub64BitRegisters(sfops, dstRegBits, tmpRegBits)
+			a.addOrSub64BitRegisters(sfops, dstRegBits, srcRegBits, tmpRegBits)
 			return
 		}
 
@@ -2212,7 +2248,7 @@ func (a *AssemblerImpl) encodeConstToRegister(n *nodeImpl) (err error) {
 			a.loadConstViaBitMaskImmediate(uc, tmpRegBits, true)
 
 			// ADD/SUB tmpReg, dstReg
-			a.addOrSub64BitRegisters(sfops, dstRegBits, tmpRegBits)
+			a.addOrSub64BitRegisters(sfops, dstRegBits, srcRegBits, tmpRegBits)
 			return
 		}
 
@@ -2239,7 +2275,7 @@ func (a *AssemblerImpl) encodeConstToRegister(n *nodeImpl) (err error) {
 		// Otherwise we use MOVZ and MOVNs for loading const into tmpRegister.
 		tmpRegBits := registerBits(a.temporaryRegister)
 		a.load64bitConst(c, tmpRegBits)
-		a.addOrSub64BitRegisters(sfops, dstRegBits, tmpRegBits)
+		a.addOrSub64BitRegisters(sfops, dstRegBits, srcRegBits, tmpRegBits)
 	case MOVW:
 		if c == 0 {
 			a.buf.Write([]byte{
@@ -3792,7 +3828,7 @@ func (a *AssemblerImpl) encodeRegisterToVectorRegister(n *nodeImpl) (err error)
 var zeroRegisterBits byte = 0b11111
 
 func isIntRegister(r asm.Register) bool {
-	return RegR0 <= r && r <= RegRZR
+	return RegR0 <= r && r <= RegSP
 }
 
 func isVectorRegister(r asm.Register) bool {
@@ -3806,9 +3842,11 @@ func isConditionalRegister(r asm.Register) bool {
 func intRegisterBits(r asm.Register) (ret byte, err error) {
 	if !isIntRegister(r) {
 		err = fmt.Errorf("%s is not integer", RegisterName(r))
-	} else {
-		ret = byte(r - RegR0)
+	} else if r == RegSP {
+		// SP has the same bit representations as RegRZR.
+		r = RegRZR
 	}
+	ret = byte(r - RegR0)
 	return
 }
 
@@ -3823,6 +3861,10 @@ func vectorRegisterBits(r asm.Register) (ret byte, err error) {
 
 func registerBits(r asm.Register) (ret byte) {
 	if isIntRegister(r) {
+		if r == RegSP {
+			// SP has the same bit representations as RegRZR.
+			r = RegRZR
+		}
 		ret = byte(r - RegR0)
 	} else {
 		ret = byte(r - RegV0)


Test patch:
diff --git a/internal/asm/arm64/impl_1_test.go b/internal/asm/arm64/impl_1_test.go
index 422a7cb2..1cfe437e 100644
--- a/internal/asm/arm64/impl_1_test.go
+++ b/internal/asm/arm64/impl_1_test.go
@@ -471,6 +471,18 @@ func Test_CompileRegisterAndConstToNone(t *testing.T) {
 	require.Equal(t, operandTypeNone, actualNode.types.dst)
 }
 
+func Test_CompileRegisterAndConstToRegister(t *testing.T) {
+	a := NewAssembler(RegR10)
+	a.CompileRegisterAndConstToRegister(CMP, RegR27, 10, RegSP)
+	actualNode := a.current
+	require.Equal(t, CMP, actualNode.instruction)
+	require.Equal(t, RegR27, actualNode.srcReg)
+	require.Equal(t, int64(10), actualNode.srcConst)
+	require.Equal(t, RegSP, actualNode.dstReg)
+	require.Equal(t, operandTypeRegisterAndConst, actualNode.types.src)
+	require.Equal(t, operandTypeRegister, actualNode.types.dst)
+}
+
 func Test_CompileLeftShiftedRegisterToRegister(t *testing.T) {
 	a := NewAssembler(RegR10)
 	a.CompileLeftShiftedRegisterToRegister(ADD, RegR27, 10, RegR28, RegR5)
diff --git a/internal/asm/arm64/impl_2_test.go b/internal/asm/arm64/impl_2_test.go
index 0728b314..faa30b10 100644
--- a/internal/asm/arm64/impl_2_test.go
+++ b/internal/asm/arm64/impl_2_test.go
@@ -44,6 +44,64 @@ func TestAssemblerImpl_EncodeConstToRegister(t *testing.T) {
 		n    *nodeImpl
 		exp  []byte
 	}{
+		{
+			name: "add x29, sp, #0x200",
+			n: &nodeImpl{
+				instruction: ADD,
+				srcReg:      RegSP,
+				dstReg:      RegR29,
+				srcConst:    512,
+			},
+			exp: []byte{0xfd, 0x3, 0x8, 0x91, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		},
+		{
+			name: "add x29, sp, #0x100, lsl #12",
+			n: &nodeImpl{
+				instruction: ADD,
+				srcReg:      RegSP,
+				dstReg:      RegR29,
+				srcConst:    0x100000,
+			},
+			exp: []byte{0xfd, 0x3, 0x44, 0x91, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		},
+		{
+			name: "sub x29, sp, #0x200",
+			n: &nodeImpl{
+				instruction: SUB,
+				srcReg:      RegSP,
+				dstReg:      RegR29,
+				srcConst:    512,
+			},
+			exp: []byte{0xfd, 0x3, 0x8, 0xd1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		},
+		{
+			name: "sub x29, sp, #0x100, lsl #12",
+			n: &nodeImpl{
+				instruction: SUB,
+				srcReg:      RegSP,
+				dstReg:      RegR29,
+				srcConst:    0x100000,
+			},
+			exp: []byte{0xfd, 0x3, 0x44, 0xd1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		},
+		{
+			name: "add sp, sp, #0x200",
+			n: &nodeImpl{
+				instruction: ADD,
+				dstReg:      RegSP,
+				srcConst:    512,
+			},
+			exp: []byte{0xff, 0x3, 0x8, 0x91, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		},
+		{
+			name: "sub sp, sp, #0x200",
+			n: &nodeImpl{
+				instruction: SUB,
+				dstReg:      RegSP,
+				srcConst:    512,
+			},
+			exp: []byte{0xff, 0x3, 0x8, 0xd1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		},
 		{
 			name: "and w30, w30, #1",
 			n: &nodeImpl{
diff --git a/internal/asm/arm64/impl_3_test.go b/internal/asm/arm64/impl_3_test.go
index d313b496..f450bdc6 100644
--- a/internal/asm/arm64/impl_3_test.go
+++ b/internal/asm/arm64/impl_3_test.go
@@ -1,6 +1,7 @@
 package arm64
 
 import (
+	"encoding/hex"
 	"testing"
 
 	"github.com/tetratelabs/wazero/internal/asm"
@@ -708,6 +709,11 @@ func TestAssemblerImpl_EncodeRegisterToRegister(t *testing.T) {
 		src, dst asm.Register
 		exp      []byte
 	}{
+		{name: "MOV/src=RegSP,dst=R10", inst: MOVD, src: RegSP, dst: RegR10, exp: []byte{0xea, 0x3, 0x0, 0x91}},
+		{name: "MOV/src=RegSP,dst=R30", inst: MOVD, src: RegSP, dst: RegR30, exp: []byte{0xfe, 0x3, 0x0, 0x91}},
+		{name: "MOV/src=R10,dst=RegSP", inst: MOVD, src: RegR10, dst: RegSP, exp: []byte{0x5f, 0x1, 0x0, 0x91}},
+		{name: "MOV/src=R10,dst=R10", inst: MOVD, src: RegR30, dst: RegSP, exp: []byte{0xdf, 0x3, 0x0, 0x91}},
+		{name: "MOV/src=R10,dst=RegSP", inst: MOVD, src: RegSP, dst: RegSP, exp: []byte{0xff, 0x3, 0x0, 0x91}},
 		{name: "ADD/src=RZR,dst=RZR", inst: ADD, src: RegRZR, dst: RegRZR, exp: []byte{0xff, 0x3, 0x1f, 0x8b}},
 		{name: "ADD/src=RZR,dst=R10", inst: ADD, src: RegRZR, dst: RegR10, exp: []byte{0x4a, 0x1, 0x1f, 0x8b}},
 		{name: "ADD/src=RZR,dst=R30", inst: ADD, src: RegRZR, dst: RegR30, exp: []byte{0xde, 0x3, 0x1f, 0x8b}},
@@ -1189,7 +1195,7 @@ func TestAssemblerImpl_EncodeRegisterToRegister(t *testing.T) {
 			require.NoError(t, err)
 
 			actual := a.bytes()
-			require.Equal(t, tc.exp, actual[:4])
+			require.Equal(t, tc.exp, actual[:4], hex.EncodeToString(actual[:4]))
 		})
 	}
 }