fix(asm): route the arm64 logical immediates to ZR through REGTMP

Assisted-by: GLM 5.3 Flash

asm/arm64_assemble.go

@@ -311,6 +311,37 @@ func arm64InstrSize(instr *ast.Instr, fi arm64FrameInfo, pos int) int {
 	case "MOV", "MOVD", "MOVW", "MOVWU", "MOVH", "MOVHU", "MOVB", "MOVBU",
 		"FMOVS", "FMOVD", "FMOVQ":
 		return arm64MovSize(mnem, ops, fi)
+	}
+	// The logical-immediate family: one word on the bitmask fast path (a
+	// real destination, or the flags-only TST spellings), two words for the
+	// materialisation the ZR destinations and the non-bitmask immediates
+	// take.  Mirrors encodeARM64DPSR's decision exactly.
+	if len(ops) >= 2 && len(ops) <= 3 && isImmOperand(ops[0]) {
+		switch mnem {
+		case "AND", "ANDW", "ANDS", "ANDSW", "ORR", "ORRW", "EOR", "EORW",
+			"BIC", "BICW", "BICS", "BICSW", "ORN", "ORNW", "EON", "EONW",
+			"TST", "TSTW":
+			if v, ok := arm64ImmOperandValue(ops[0]); ok {
+				switch mnem {
+				case "BIC", "BICW", "BICS", "BICSW", "ORN", "ORNW", "EON", "EONW":
+					v = ^v
+				}
+				width := 64
+				if strings.HasSuffix(mnem, "W") {
+					width = 32
+				}
+				rd := 31
+				if mnem != "TST" && mnem != "TSTW" {
+					rd = arm64RegNum(operandRegName(ops[len(ops)-1]))
+				}
+				if _, _, _, bc := a64LogicalImm(v, width); bc && (rd != 31 || strings.HasPrefix(mnem, "TST")) {
+					return 4
+				}
+				return 8
+			}
+		}
+	}
+	switch mnem {
 	case "ADD", "ADDW", "SUB", "SUBW", "CMP", "CMPW", "CMN", "CMNW",
 		"ADDS", "ADDSW", "SUBS", "SUBSW":
 		if len(ops) >= 2 && isImmOperand(ops[0]) {
@@ -839,48 +870,6 @@ func encodeARM64DPSR(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, er
 			if strings.HasSuffix(mnem, "W") {
 				width = 32
 			}
-			n, immr, imms, ok := a64LogicalImm(v, width)
-			if !ok {
-				// Beyond the bitmask immediates the toolchain materialises
-				// the constant into REGTMP (R27) and uses the register form
-				// (asm7.go cases 62 and 13).  BIC/ORN/EON read the written
-				// value, so the materialisation uses v before any inversion.
-				written := v
-				if inverted {
-					written = ^v
-				}
-				width := mnem
-				if strings.HasSuffix(mnem, "W") {
-					width = "MOVW"
-				} else {
-					width = "MOVD"
-				}
-				mw, merr := encodeARM64LoadImm(27, written, width)
-				var rn, rd int
-				switch len(ops) {
-				case 3:
-					rn = arm64RegNum(operandRegName(ops[1]))
-					rd = arm64RegNum(operandRegName(ops[2]))
-				default:
-					rd = arm64RegNum(operandRegName(ops[1]))
-					rn = rd
-				}
-				if isCmp {
-					rd = 31
-				}
-				if merr != nil || rn < 0 || rd < 0 {
-					return nil, fmt.Errorf("%s: immediate %q is not a logical (bitmask) immediate", mnem, strings.Join(strings.Fields(ops[0].Raw), " "))
-				}
-				// The register tail against SP takes the extended form, like the
-				// plain register path below.
-				tail := baseOp | 27<<16 | uint32(rn)<<5 | uint32(rd)
-				if opt, spok := arm64SpExtendOpt(mnem, ops[1:]); spok {
-					tail = baseOp | 1<<21 | opt<<13 | 27<<16 | uint32(rn)<<5 | uint32(rd)
-				}
-				return append(mw, a64wordLE(tail)...), nil
-			}
-			opc := (baseOp >> 29) & 7
-			sf := (baseOp >> 31) & 1
 			var rn, rd int
 			switch len(ops) {
 			case 3:
@@ -890,11 +879,50 @@ func encodeARM64DPSR(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, er
 				rd = arm64RegNum(operandRegName(ops[1]))
 				rn = rd
 			}
+			if isCmp {
+				// CMP/CMN/TST write the flags alone: the destination is ZR
+				// whatever the spelling says.
+				rd = 31
+			}
 			if rn < 0 || rd < 0 {
 				return nil, fmt.Errorf("invalid register operand in %s", mnem)
 			}
-			return a64wordLE(sf<<31 | opc<<29 | 0x24<<23 | n<<22 | immr<<16 | imms<<10 |
+			n, immr, imms, ok := a64LogicalImm(v, width)
-				uint32(rn)<<5 | uint32(rd)), nil
+			// The toolchain's logical-immediate rows take a real destination
+			// only (omovconst guards the bitmask path with rt != REGZERO): a
+			// non-flag-setting logical to ZR materialises the constant into
+			// REGTMP (R27) and takes the register form.  The flags-only TST
+			// spellings keep the fast path: ANDS ZR, Rn, #imm is their form.
+			if ok && (rd != 31 || isCmp) {
+				opc := (baseOp >> 29) & 7
+				sf := (baseOp >> 31) & 1
+				return a64wordLE(sf<<31 | opc<<29 | 0x24<<23 | n<<22 | immr<<16 | imms<<10 |
+					uint32(rn)<<5 | uint32(rd)), nil
+			}
+			// Beyond the bitmask immediates, and for the ZR destinations, the
+			// toolchain materialises the constant into REGTMP and uses the
+			// register form (asm7.go cases 62 and 13).  BIC/ORN/EON read the
+			// written value, so the materialisation uses v before any
+			// inversion.
+			written := v
+			if inverted {
+				written = ^v
+			}
+			mwMnem := "MOVD"
+			if strings.HasSuffix(mnem, "W") {
+				mwMnem = "MOVW"
+			}
+			mw, merr := encodeARM64LoadImm(27, written, mwMnem)
+			if merr != nil {
+				return nil, fmt.Errorf("%s: immediate %q is not a logical (bitmask) immediate", mnem, strings.Join(strings.Fields(ops[0].Raw), " "))
+			}
+			// The register tail against SP takes the extended form, like the
+			// plain register path below.
+			tail := baseOp | 27<<16 | uint32(rn)<<5 | uint32(rd)
+			if opt, spok := arm64SpExtendOpt(mnem, ops[1:]); spok {
+				tail = baseOp | 1<<21 | opt<<13 | 27<<16 | uint32(rn)<<5 | uint32(rd)
+			}
+			return append(mw, a64wordLE(tail)...), nil
 		}
 	}
 

asm/arm64_encode_test.go

@@ -2119,3 +2119,41 @@ func TestArm64FPImmediate(t *testing.T) {
 		}
 	}
 }
+
+// TestArm64LogicalImmZR pins the logical-immediate spellings the destination
+// changes, against `go tool asm` words: TST keeps the fast ANDS-to-ZR form
+// while a non-flag-setting logical to ZR materialises the constant into
+// REGTMP and takes the register form (omovconst's rt != REGZERO guard).
+func TestArm64LogicalImmZR(t *testing.T) {
+	got := arm64Words(t,
+		"\tTSTW $0x600000006, R1\n"+
+			"\tTST $15, R2\n"+
+			"\tAND $1, ZR\n"+
+			"\tANDW $1, ZR\n"+
+			"\tEOR $1, ZR\n"+
+			"\tAND $15, R2, R3\n"+
+			"\tANDS $0xff, R2, R3\n"+
+			"\tORR $0x10, R2, R3\n")
+	want := []uint32{
+		0x721f043f, // ANDS W31, W1, #bitmask (TSTW)
+		0xf2400c5f, // ANDS X31, X2, #15 (TST)
+		0xb24003fb, // ORR X27, XZR, #1
+		0x8a1b03ff, // AND X31, X31, R27
+		0x320003fb, // ORR W27, WZR, #1 (ANDW to ZR)
+		0x0a1b03ff, // AND W31, W31, W27
+		0xb24003fb, // ORR X27, XZR, #1 (EOR to ZR)
+		0xca1b03ff, // EOR X31, X31, R27
+		0x92400c43, // AND X3, X2, #15
+		0xf2401c43, // ANDS X3, X2, #0xff
+		0xb27c0043, // ORR X3, X2, #0x10
+		0xd65f03c0, // RET
+	}
+	if len(got) != len(want) {
+		t.Fatalf("word count = %d, want %d", len(got), len(want))
+	}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Errorf("word %d = %08x, want %08x", i, got[i], want[i])
+		}
+	}
+}