From 4fc96decc40602bf79c25c7e1b99c7b05f508efa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Petr=20Balv=C3=ADn?= <opensource@petrbalvin.org>
Date: Wed, 7 Oct 2026 13:14:33 +0200
Subject: [PATCH] fix(asm): size the arm64 logical-immediate materialisation
 exactly

Assisted-by: GLM 5.3
---
 asm/arm64_assemble.go    | 15 +++++++++++----
 asm/arm64_encode_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/asm/arm64_assemble.go b/asm/arm64_assemble.go
index 187ab42..cd9c219 100644
--- a/asm/arm64_assemble.go
+++ b/asm/arm64_assemble.go
@@ -313,22 +313,26 @@ func arm64InstrSize(instr *ast.Instr, fi arm64FrameInfo, pos int) int {
 		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.
+	// real destination, or the flags-only TST spellings), otherwise the
+	// constant materialisation into REGTMP plus the register-form tail.
+	// Mirrors encodeARM64DPSR's decision exactly, the materialisation word
+	// count included: a MOVZ plus up to three MOVKs makes five words.
 	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 {
+				written := v
 				switch mnem {
 				case "BIC", "BICW", "BICS", "BICSW", "ORN", "ORNW", "EON", "EONW":
 					v = ^v
 				}
 				width := 64
+				mwMnem := "MOVD"
 				if strings.HasSuffix(mnem, "W") {
 					width = 32
+					mwMnem = "MOVW"
 				}
 				zrDest := false
 				if mnem != "TST" && mnem != "TSTW" {
@@ -337,7 +341,10 @@ func arm64InstrSize(instr *ast.Instr, fi arm64FrameInfo, pos int) int {
 				if _, _, _, bc := a64LogicalImm(v, width); bc && (mnem == "TST" || mnem == "TSTW" || !zrDest) {
 					return 4
 				}
-				return 8
+				if mw, err := encodeARM64LoadImm(27, written, mwMnem); err == nil {
+					return len(mw) + 4
+				}
+				return 4
 			}
 		}
 	}
diff --git a/asm/arm64_encode_test.go b/asm/arm64_encode_test.go
index 0df75ee..a9d8919 100644
--- a/asm/arm64_encode_test.go
+++ b/asm/arm64_encode_test.go
@@ -2070,6 +2070,40 @@ func TestArm64ConRnRejections(t *testing.T) {
 	}
 }
 
+// TestArm64LogicalMaterialisationBranch pins a forward branch over the
+// three-word logical materialisation against `go tool asm -S` (Go 1.27,
+// arm64): the size pass must count the MOVZ/MOVK pair the encoder lays down,
+// or the label offsets desynchronise from the bytes and the branch lands a
+// word early.
+func TestArm64LogicalMaterialisationBranch(t *testing.T) {
+	// The body closes with its own RET under the end label, so the file is
+	// parsed as written rather than through arm64Words' appended RET.
+	f, errs := parser.Parse("test_arm64.s", "#include \"textflag.h\"\n\nTEXT ·f(SB), NOSPLIT, $0-0\n\tCBZ\tR2, end\n\tTST\t$0x4900000049, R0\nend:\tRET\n")
+	if len(errs) > 0 {
+		t.Fatalf("parse: %v", errs)
+	}
+	img, err := AssembleFileARM64(f)
+	if err != nil {
+		t.Fatalf("AssembleFileARM64: %v", err)
+	}
+	got := leWords(img.Code)
+	want := []uint32{
+		0xb4000082, // CBZ R2, +16 (four words ahead)
+		0xd280093b, // MOVZ $0x49, R27
+		0xf2c0093b, // MOVK $(0x49<<32), R27
+		0xea1b001f, // TST R27, R0
+		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])
+		}
+	}
+}
+
 // TestArm64FPImmediate pins the FP immediate moves against `go tool asm`
 // words: the FMOV (immediate) instruction for the 8-bit encodable values and
 // the FMOV-from-ZR move for zero, plus the rejections the toolchain raises