fix(parser): substitute macro parameters behind element selectors

Assisted-by: GLM 5.3 Flash

parser/preproc.go

@@ -349,7 +349,7 @@ func (pp *preproc) expandTokens(in []token.Token) []token.Token {
 			consecutive = 0
 			continue
 		}
-		def := pp.macros[t.Text]
+		def, suffix := pp.macroFor(t.Text)
 		if def == nil {
 			i++
 			consecutive = 0
@@ -363,7 +363,14 @@ func (pp *preproc) expandTokens(in []token.Token) []token.Token {
 			return nil
 		}
 		if def.args == nil {
-			s = append(s[:i], append(restamp(def.body, t.Pos), s[i+1:]...)...)
+			body := restamp(def.body, t.Pos)
+			if suffix != "" {
+				// The macro was reached only through a compound spelling
+				// (ACC0.B16 over "#define ACC0 V8"), so the selector has
+				// to travel with the expansion.
+				body = appendSelector(body, suffix, t.Pos)
+			}
+			s = append(s[:i], append(body, s[i+1:]...)...)
 			continue
 		}
 		// A parameterised macro invoked without its parentheses stands
@@ -394,6 +401,17 @@ func (pp *preproc) expandTokens(in []token.Token) []token.Token {
 					sub = append(sub, restamp(args[k], t.Pos)...)
 					continue
 				}
+				// A parameter used with an element or lane selector: the
+				// lexer folds A.S4 into one identifier, so the whole-token
+				// match above cannot see the parameter.  The toolchain
+				// lexes the period separately and substitutes the name
+				// alone; splitting at the FIRST period and pasting the
+				// argument back in front of the selector is the equivalent
+				// for this lexer.
+				if k, sel := parameterSelector(bt.Text, def.args); k >= 0 {
+					sub = append(sub, restamp(pasteSelector(args[k], sel), t.Pos)...)
+					continue
+				}
 			}
 			sub = append(sub, bt)
 		}
@@ -402,6 +420,75 @@ func (pp *preproc) expandTokens(in []token.Token) []token.Token {
 	return s
 }
 
+// macroFor finds the macro a use names.  The lexer folds NAME.selector into
+// one identifier token, so a macro written behind a selector suffix
+// (ACC0.B16 over "#define ACC0 V8") never matches a whole-token table
+// lookup; the toolchain splits on the period and reads the two halves, so
+// the prefix before the FIRST period is tried here as well and the caller
+// re-attaches the suffix to whatever the macro expands to.  Only a whole
+// name counts: AB.S4 does not reach a macro named A, and a parameterised
+// macro is not hidden behind a selector, because its invocation would need
+// the parentheses to follow the bare name.
+func (pp *preproc) macroFor(text string) (*macroDef, string) {
+	if def := pp.macros[text]; def != nil {
+		return def, ""
+	}
+	if j := strings.IndexByte(text, '.'); j > 0 {
+		if def := pp.macros[text[:j]]; def != nil && def.args == nil {
+			return def, text[j:]
+		}
+	}
+	return nil, ""
+}
+
+// appendSelector glues a selector suffix onto an object macro's expansion:
+// the selector binds to the identifier the expansion ends with, the way the
+// toolchain's operand parser reads V0 and .B16 back as one register
+// spelling.  An expansion that does not end in an identifier carries the
+// selector as its own token, which the parser then reports where it cannot
+// parse it.
+func appendSelector(body []token.Token, suffix string, pos token.Position) []token.Token {
+	if n := len(body); n > 0 && body[n-1].Kind == token.Ident {
+		body[n-1].Text += suffix
+		return body
+	}
+	return append(body, token.Token{Kind: token.Ident, Text: suffix, Pos: pos, End: pos})
+}
+
+// parameterSelector reports the argument a compound body token names: the
+// parameter whose whole name occupies the text before the token's FIRST
+// period, with the selector that follows.  k is negative when no parameter
+// matches, which leaves tokens like AB.S4 untouched even though a parameter
+// A is bound.
+func parameterSelector(text string, args []string) (int, string) {
+	j := strings.IndexByte(text, '.')
+	if j <= 0 {
+		return -1, ""
+	}
+	if k := slices.Index(args, text[:j]); k >= 0 {
+		return k, text[j:]
+	}
+	return -1, ""
+}
+
+// pasteSelector joins an argument with the selector a compound body token
+// carries, textually: the selector binds to the identifier the argument
+// ends with, so A.S4 over the argument V0.B16 spells V0.B16.S4, exactly the
+// operand the toolchain's split-then-substitute leaves behind.  An argument
+// with no trailing identifier carries the selector as a separate token,
+// which the parser then reports where it cannot parse it.
+func pasteSelector(val []token.Token, suffix string) []token.Token {
+	if len(val) == 0 {
+		return []token.Token{{Kind: token.Ident, Text: suffix}}
+	}
+	out := slices.Clone(val)
+	if n := len(out); out[n-1].Kind == token.Ident {
+		out[n-1].Text += suffix
+		return out
+	}
+	return append(out, token.Token{Kind: token.Ident, Text: suffix})
+}
+
 // collectArgs reads the actual argument tokens of an invocation; the opening
 // parenthesis is at start.  Commas separate arguments except inside nested
 // parentheses.  A nil result means the list was unterminated, which is a

parser/preproc_test.go

@@ -536,6 +536,127 @@ func TestConstantExpressionFoldsWithoutExpand(t *testing.T) {
 	}
 }
 
+func TestParameterWithSelectorSubstitutes(t *testing.T) {
+	// The lexer folds A.S4 into one identifier token, so a parameter used
+	// with an element or lane selector never matched the whole-token
+	// substitution; the toolchain's lexer splits on the period and its
+	// substitution sees the name alone.  Several parameters carry selectors
+	// in one body here, which is the chacha8_arm64.s QR shape in miniature.
+	_, got := expand(t, `
+#define QR(A, B, C, D) VADD A.S4, B.S4, C.S4; VEOR D.B16, A.B16, D.B16
+TEXT ·f(SB), NOSPLIT, $0
+	QR(V0, V1, V2, V3)
+	RET
+`)
+	wantLines(t, got,
+		"VADD V0.S4, V1.S4, V2.S4",
+		"VEOR V3.B16, V0.B16, V3.B16",
+		"RET",
+	)
+}
+
+func TestSelectorWithCompoundArgumentPastesTextually(t *testing.T) {
+	// An argument that is itself one compound identifier pastes verbatim:
+	// A.S4 over V0.B16 spells V0.B16.S4, the operand the toolchain's
+	// split-then-substitute leaves behind.
+	_, got := expand(t, `
+#define M(A) VADD A.S4, A.S4, A.S4
+TEXT ·f(SB), NOSPLIT, $0
+	M(V0.B16)
+	RET
+`)
+	wantLines(t, got, "VADD V0.B16.S4, V0.B16.S4, V0.B16.S4", "RET")
+}
+
+func TestSelectorAlongsideBareParameter(t *testing.T) {
+	// A body may use the parameter bare and suffixed, and the argument may
+	// itself end in a selector; neither disturbs the other.
+	_, got := expand(t, `
+#define M(A) VADD A, A.S4, A
+TEXT ·f(SB), NOSPLIT, $0
+	M(V0)
+	M(V1.B16)
+	RET
+`)
+	wantLines(t, got,
+		"VADD V0, V0.S4, V0",
+		"VADD V1.B16, V1.B16.S4, V1.B16",
+		"RET",
+	)
+}
+
+func TestSelectorKeepsNonParameterPrefixes(t *testing.T) {
+	// The prefix before the period must be the whole parameter name:
+	// AB.S4 never reaches a parameter A.
+	_, got := expand(t, `
+#define M(A) VADD AB.S4, A.S4, AB.S4
+TEXT ·f(SB), NOSPLIT, $0
+	M(V0)
+	RET
+`)
+	wantLines(t, got, "VADD AB.S4, V0.S4, AB.S4", "RET")
+}
+
+func TestSelectorExpandsMacroValuedArgument(t *testing.T) {
+	// gcm_arm64.s invokes mulRound(B1) where B1 is itself an object macro:
+	// the paste stays rescannable, so B1.D1 still expands to V1.D1 the way
+	// the toolchain's rescan of substituted tokens does.
+	_, got := expand(t, `
+#define B1 V1
+#define mulRound(X) VPMULL X.D1, T1.D1, T3.Q1
+TEXT ·f(SB), NOSPLIT, $0
+	mulRound(B1)
+	RET
+`)
+	wantLines(t, got, "VPMULL V1.D1, T1.D1, T3.Q1", "RET")
+}
+
+func TestObjectMacroBehindSelectorExpands(t *testing.T) {
+	// Ordinary code writes ACC0.B16 where ACC0 is an object macro; the
+	// toolchain expands the alias because its lexer reads the selector as
+	// its own token, and the lookup here must reach the macro through the
+	// compound spelling the same way.
+	_, got := expand(t, `
+#define ACC0 V8
+TEXT ·f(SB), NOSPLIT, $0
+	VEOR ACC0.B16, ACC0.B16, ACC0.B16
+	RET
+`)
+	wantLines(t, got, "VEOR V8.B16, V8.B16, V8.B16", "RET")
+}
+
+func TestChacha8QRMacroExpands(t *testing.T) {
+	// The real QR round of chacha8_arm64.s end to end: every parameter
+	// carries a selector somewhere, and the round is sixteen instructions.
+	_, got := expand(t, `
+#define QR(A, B, C, D) \
+	VADD A.S4, B.S4, A.S4; VEOR D.B16, A.B16, D.B16; VREV32 D.H8, D.H8; \
+	VADD C.S4, D.S4, C.S4; VEOR B.B16, C.B16, V30.B16; VSHL $12, V30.S4, B.S4; VSRI $20, V30.S4, B.S4; \
+	VADD A.S4, B.S4, A.S4; VEOR D.B16, A.B16, D.B16; VTBL V31.B16, [D.B16], D.B16; \
+	VADD C.S4, D.S4, C.S4; VEOR B.B16, C.B16, V30.B16; VSHL $7, V30.S4, B.S4; VSRI $25, V30.S4, B.S4
+TEXT ·f(SB), NOSPLIT, $0
+	QR(V0, V1, V2, V3)
+	RET
+`)
+	wantLines(t, got,
+		"VADD V0.S4, V1.S4, V0.S4",
+		"VEOR V3.B16, V0.B16, V3.B16",
+		"VREV32 V3.H8, V3.H8",
+		"VADD V2.S4, V3.S4, V2.S4",
+		"VEOR V1.B16, V2.B16, V30.B16",
+		"VSHL $12, V30.S4, V1.S4",
+		"VSRI $20, V30.S4, V1.S4",
+		"VADD V0.S4, V1.S4, V0.S4",
+		"VEOR V3.B16, V0.B16, V3.B16",
+		"VTBL V31.B16, [V3.B16], V3.B16",
+		"VADD V2.S4, V3.S4, V2.S4",
+		"VEOR V1.B16, V2.B16, V30.B16",
+		"VSHL $7, V30.S4, V1.S4",
+		"VSRI $25, V30.S4, V1.S4",
+		"RET",
+	)
+}
+
 func TestNotAnExpressionFallsBack(t *testing.T) {
 	// Symbol immediates and floats must keep their ordinary parse.
 	f, errs := ParseWithOptions("t_amd64.s", "TEXT ·f(SB), NOSPLIT, $0\n\tMOVQ $1.5, AX\n\tMOVQ $·sym(SB), AX\n\tRET\n", Options{})