feat(asm): assemble the extended instruction layer on arm64

Assisted-by: GLM 5.3 Flash

asm/arm64_assemble.go

@@ -9,6 +9,7 @@ import (
 	"strconv"
 	"strings"
 
+	"sourcedock.dev/petrbalvin/gasm-sdk/arch"
 	"sourcedock.dev/petrbalvin/gasm-sdk/ast"
 )
 
@@ -269,6 +270,13 @@ func arm64InstrSize(instr *ast.Instr, fi arm64FrameInfo, pos int) int {
 	case "NO_LOCAL_POINTERS", "GO_ARGS", "GO_RESULTS_INITIALIZED", "END", "FUNCDATA", "PCDATA":
 		return 0
 	}
+	// The extended-instruction layer is one instruction word in every form
+	// the registry takes: pass 1 must size a pinned statement at the 4 bytes
+	// the encoder will lay down, ahead of the scalar immediate expansions
+	// below, whose sizes would misread a Z destination (asm/arm64_ext.go).
+	if _, pinned, _ := arm64ExtStatement(mnem, ops); pinned {
+		return 4
+	}
 	switch mnem {
 	case "MOV", "MOVD", "MOVW", "MOVWU", "MOVH", "MOVHU", "MOVB", "MOVBU",
 		"FMOVS", "FMOVD":
@@ -307,6 +315,23 @@ func encodeARM64Instr(instr *ast.Instr, pc int, offsets map[string]int, fi arm64
 	mnem := strings.ToUpper(instr.Mnemonic.Text)
 	ops := instr.Operands
 
+	// The extended-instruction layer: a statement whose mnemonic is
+	// registered in the extension registry and whose operands carry a
+	// scalable vector or predicate register encodes through the registry,
+	// before any scalar route can misread those operands.  Scalar, NEON and
+	// FP operand lists never pin, so everything below runs exactly as it
+	// did (asm/arm64_ext.go).
+	if extops, pinned, convErr := arm64ExtStatement(mnem, ops); pinned {
+		if convErr != nil {
+			return nil, convErr
+		}
+		code, encErr := EncodeExtension(arch.ARM64, mnem, extops...)
+		if encErr != nil {
+			return nil, encErr
+		}
+		return code, nil
+	}
+
 	// Pseudo-instructions and special cases first.
 	switch mnem {
 	case "RET":

asm/arm64_ext.go

@@ -0,0 +1,214 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: BSD-3-Clause
+
+// The assembler's side of the extended-instruction layer: this file turns a
+// parsed arm64 statement into the operand form arch.ExtInstr.Encode consumes
+// and routes statements only the layer can encode through the registry.  It
+// sits beside the main arm64 encoders, never inside them: the generated
+// tables and the scalar, NEON and FP paths are untouched, and a statement
+// reaches this file only when the mnemonic is registered in the extension
+// layer and at least one operand is a scalable vector or predicate register.
+//
+// The spellings are the layer's own Plan 9 forms, the ones its metadata
+// documents: Zn, Zm, Zd for the unpredicated three-vector class, Zm, Pg/M,
+// Zdn for the predicated class, imm{, LSL #8}, Zdn for the immediate classes.
+
+package asm
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+
+	"sourcedock.dev/petrbalvin/gasm-sdk/arch"
+	"sourcedock.dev/petrbalvin/gasm-sdk/ast"
+)
+
+// arm64ExtStatement converts one instruction's operands into the extended
+// layer's operand form.  pinned reports that the statement belongs to the
+// layer: the mnemonic is registered in the registry and the operand list
+// carries at least one scalable vector or predicate register.  A pinned
+// statement can only encode through the layer, so every operand is read
+// here and its diagnostic replaces whatever the scalar paths would have
+// said about operands they cannot read; err is non-nil for a pinned
+// statement whose operands the layer refuses, and extops is complete only
+// when err is nil.  Unpinned means the statement is nobody's: the caller
+// falls through to the ordinary arm64 encoders, which keep their exact
+// behaviour for every scalar, NEON and FP operand list.
+func arm64ExtStatement(mnem string, ops []*ast.Operand) (extops []arch.ExtOperand, pinned bool, err error) {
+	if _, ok := LookupExtension(arch.ARM64, mnem); !ok {
+		return nil, false, nil
+	}
+	if !arm64ExtPinned(ops) {
+		return nil, false, nil
+	}
+	out := make([]arch.ExtOperand, 0, len(ops))
+	for i, op := range ops {
+		text := strings.Join(strings.Fields(op.Raw), "")
+
+		// The spelled shift of an immediate class: the shift is an attribute
+		// of the preceding immediate operand (imm{, LSL #8}, Zdn), never an
+		// operand of its own.
+		if amount, ok := strings.CutPrefix(text, "LSL#"); ok {
+			if len(out) == 0 || out[len(out)-1].Kind != arch.ExtImm || out[len(out)-1].HasShift {
+				return nil, true, fmt.Errorf("%s: operand %d (%s): LSL belongs straight after an immediate", mnem, i+1, op.Raw)
+			}
+			n, convErr := strconv.Atoi(amount)
+			if convErr != nil {
+				return nil, true, fmt.Errorf("%s: operand %d (%s): %q is not an LSL amount", mnem, i+1, op.Raw, amount)
+			}
+			out[len(out)-1].Shift, out[len(out)-1].HasShift = n, true
+			continue
+		}
+
+		if op.Kind == ast.OpImmediate {
+			ext, ok := arm64ExtImmediate(op)
+			if !ok {
+				return nil, true, fmt.Errorf("%s: operand %d (%s) is not an immediate the layer can read", mnem, i+1, op.Raw)
+			}
+			out = append(out, ext)
+			continue
+		}
+		if ext, ok := arm64ExtVector(text); ok {
+			out = append(out, ext)
+			continue
+		}
+		if ext, ok := arm64ExtPredicate(text); ok {
+			out = append(out, ext)
+			continue
+		}
+		return nil, true, fmt.Errorf("%s: operand %d (%s) is not an extended-layer operand: want a scalable vector register, a predicate register or an immediate", mnem, i+1, op.Raw)
+	}
+	return out, true, nil
+}
+
+// arm64ExtPinned reports whether any operand is a scalable vector or
+// predicate register, the shapes only the extension layer reads.  The test
+// is deliberately loose about the suffixes: P0/B is not a spelling the
+// layer takes, but the P of it makes the statement the layer's, and the
+// conversion then diagnoses the operand precisely instead of leaving it to
+// a scalar path that would report an unrelated register error.
+func arm64ExtPinned(ops []*ast.Operand) bool {
+	for _, op := range ops {
+		if op.Kind == ast.OpImmediate {
+			continue
+		}
+		text := strings.Join(strings.Fields(op.Raw), "")
+		if _, ok := arm64ExtVector(text); ok {
+			return true
+		}
+		if arm64ExtPredicateShape(text) {
+			return true
+		}
+	}
+	return false
+}
+
+// arm64ExtPredicateShape reports whether text spells a predicate register at
+// all: P, digits, an optional arrangement suffix and an optional qualifier
+// after a slash, whatever the qualifier says.  The strict parse in
+// arm64ExtPredicate judges the suffix; this shape only decides who the
+// operand belongs to.
+func arm64ExtPredicateShape(text string) bool {
+	if text == "" || text[0] != 'P' {
+		return false
+	}
+	text = text[1:]
+	if i := strings.IndexByte(text, '/'); i >= 0 {
+		text = text[:i]
+	}
+	if i := strings.IndexByte(text, '.'); i >= 0 {
+		text = text[:i]
+	}
+	_, err := strconv.Atoi(text)
+	return err == nil && text != ""
+}
+
+// arm64ExtImmediate converts a $ immediate into the layer's form.  The
+// parser folds a parenthesised constant expression in full ($(255<<8)) and
+// reads a bare literal greedily, dropping any trailing operator tokens:
+// $255<<8 parses as 255 with the shift silently gone.  Encoding that silent
+// prefix would assemble what the text did not say, so an unparenthesised
+// immediate is accepted only when its whole text reads back as one integer
+// carrying the parser's value.
+func arm64ExtImmediate(op *ast.Operand) (arch.ExtOperand, bool) {
+	if op.Kind != ast.OpImmediate || !op.Imm.HasVal {
+		return arch.ExtOperand{}, false
+	}
+	text := strings.Join(strings.Fields(strings.TrimPrefix(op.Raw, "$")), "")
+	if !strings.HasPrefix(text, "(") {
+		if _, parseErr := strconv.ParseInt(text, 0, 64); parseErr != nil {
+			return arch.ExtOperand{}, false
+		}
+	}
+	v := op.Imm.Val
+	if op.Imm.Neg {
+		v = -v
+	}
+	return arch.ExtOperand{Kind: arch.ExtImm, Imm: v}, true
+}
+
+// arm64ExtVector parses a scalable vector register operand: Z0..Z31 with an
+// optional element-size suffix, Z0.S.  The arrangement is carried as written
+// and the encoding validates it against the form.
+func arm64ExtVector(text string) (arch.ExtOperand, bool) {
+	reg, arr, ok := arm64ExtReg(text, 'Z')
+	if !ok {
+		return arch.ExtOperand{}, false
+	}
+	return arch.ExtOperand{Kind: arch.ExtZReg, Reg: reg, Arr: arr}, true
+}
+
+// arm64ExtPredicate parses a predicate register operand: P0..P15 with an
+// optional element-size suffix and an optional qualifier, P0/M, P0.Z, P0.B/M.
+func arm64ExtPredicate(text string) (arch.ExtOperand, bool) {
+	qual := arch.ExtQualNone
+	if base, suffix, found := strings.Cut(text, "/"); found {
+		switch suffix {
+		case "M":
+			qual = arch.ExtQualMerging
+		case "Z":
+			qual = arch.ExtQualZeroing
+		default:
+			return arch.ExtOperand{}, false
+		}
+		text = base
+	}
+	reg, arr, ok := arm64ExtReg(text, 'P')
+	if !ok {
+		return arch.ExtOperand{}, false
+	}
+	return arch.ExtOperand{Kind: arch.ExtPReg, Reg: reg, Arr: arr, Qual: qual}, true
+}
+
+// arm64ExtReg parses Pn or Zn with an optional arrangement suffix off a
+// normalised operand text.  The register range is left to the encoding: the
+// layer's own diagnostics name the range a form carries.
+func arm64ExtReg(text string, letter byte) (reg int, arr arch.ExtArrangement, ok bool) {
+	if len(text) < 2 || text[0] != letter {
+		return 0, 0, false
+	}
+	digits := text[1:]
+	if base, suffix, found := strings.Cut(digits, "."); found {
+		switch suffix {
+		case "B":
+			arr = arch.ExtArrB
+		case "H":
+			arr = arch.ExtArrH
+		case "S":
+			arr = arch.ExtArrS
+		case "D":
+			arr = arch.ExtArrD
+		case "Q":
+			arr = arch.ExtArrQ
+		default:
+			return 0, 0, false
+		}
+		digits = base
+	}
+	n, err := strconv.Atoi(digits)
+	if err != nil || n < 0 {
+		return 0, 0, false
+	}
+	return n, arr, true
+}

asm/arm64_ext_asm_test.go

@@ -0,0 +1,181 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: BSD-3-Clause
+
+package asm
+
+import (
+	"encoding/binary"
+	"strings"
+	"testing"
+
+	"sourcedock.dev/petrbalvin/gasm-sdk/parser"
+)
+
+// assembleArm64Words parses src, assembles it for arm64 and returns the
+// first function's body as little-endian instruction words.  Every statement
+// must encode: a failure is the test's.
+func assembleArm64Words(t *testing.T, src string) []uint32 {
+	t.Helper()
+	f, errs := parser.Parse("ext_arm64.s", src)
+	if len(errs) > 0 {
+		t.Fatalf("parse: %v", errs)
+	}
+	img, err := AssembleFileARM64(f)
+	if err != nil {
+		t.Fatalf("assemble: %v", err)
+	}
+	if len(img.Funcs) != 1 {
+		t.Fatalf("got %d functions, want 1", len(img.Funcs))
+	}
+	body := img.Code[img.Funcs[0].Offset:][:img.Funcs[0].Size]
+	if len(body)%4 != 0 {
+		t.Fatalf("body is %d bytes, not a whole number of instructions", len(body))
+	}
+	words := make([]uint32, 0, len(body)/4)
+	for i := 0; i < len(body); i += 4 {
+		words = append(words, binary.LittleEndian.Uint32(body[i:]))
+	}
+	return words
+}
+
+// assembleArm64SrcError parses and assembles src and returns the assembler's
+// error text.
+func assembleArm64SrcError(t *testing.T, src string) string {
+	t.Helper()
+	f, errs := parser.Parse("ext_arm64.s", src)
+	if len(errs) > 0 {
+		t.Fatalf("parse: %v", errs)
+	}
+	_, err := AssembleFileARM64(f)
+	if err == nil {
+		t.Fatal("assembled, want an error")
+	}
+	return err.Error()
+}
+
+const arm64ExtProbeHead = "#include \"textflag.h\"\nTEXT ·t(SB), NOSPLIT, $0\n"
+
+// TestArm64AssembleExtensionGolden drives the wired layer through the full
+// assembler: text in, instruction word out.  Each want is the encoding the
+// ARM Architecture Reference Manual's field layout gives for the statement:
+// the fixed class word, the size field from the arrangement, and the
+// register and immediate fields in the class's own places.  The arch-level
+// golden vectors in arch/arm64_ext_test.go pin the same bytes at the
+// metadata layer; these pin the text-to-bytes path.
+func TestArm64AssembleExtensionGolden(t *testing.T) {
+	tests := []struct {
+		stmt string
+		want uint32
+	}{
+		// Unpredicated three-vector: Zn, Zm, Zd, one shared arrangement.
+		{"ADD Z1.S, Z2.S, Z0.S", 0x04a20020},
+		{"ADD Z0.B, Z1.B, Z2.B", 0x04210002},
+		{"SUB Z31.D, Z30.D, Z29.D", 0x04fe07fd},
+		{"SQADD Z5.H, Z6.H, Z7.H", 0x046610a7},
+		{"UQADD Z8.S, Z9.S, Z10.S", 0x04a9150a},
+		{"SQSUB Z5.H, Z6.H, Z7.H", 0x046618a7},
+		{"UQSUB Z8.S, Z9.S, Z10.S", 0x04a91d0a},
+		{"MUL Z0.B, Z1.B, Z2.B", 0x04216002},
+		{"SMULH Z11.D, Z12.D, Z13.D", 0x04ec696d},
+		{"UMULH Z0.B, Z1.B, Z2.B", 0x04216c02},
+		// Governed destructive, merging: Zm, Pg/M, Zdn; the governing
+		// predicate is a 3-bit field, so P0-P7 alone.
+		{"ADD Z1.S, P0/M, Z0.S", 0x04800020},
+		{"SUBR Z1.S, P7/M, Z0.S", 0x04831c20},
+		{"MUL Z3.D, P2/M, Z5.D", 0x04d00865},
+		{"SUBR Z0.B, P5/M, Z31.B", 0x0403141f},
+		// Immediate classes: imm{, LSL #8}, Zdn.  A bare multiple of 256
+		// derives the shift bit, the spelling the layer canonicalises.
+		{"ADD $255, Z0.S", 0x25a0dfe0},
+		{"ADD $65280, Z0.H", 0x2560ffe0},
+		{"ADD $255, LSL #8, Z0.S", 0x25a0ffe0},
+		{"ADD $(255<<8), Z0.S", 0x25a0ffe0},
+		{"MUL $-128, Z0.B", 0x2530d000},
+	}
+	for _, tt := range tests {
+		words := assembleArm64Words(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
+		if len(words) != 2 {
+			t.Fatalf("%s: got %d words, want the statement and RET", tt.stmt, len(words))
+		}
+		if words[0] != tt.want {
+			t.Errorf("%s:\n  got  %08x\n  want %08x", tt.stmt, words[0], tt.want)
+		}
+		if words[1] != 0xd65f03c0 {
+			t.Errorf("%s: RET encoded %08x", tt.stmt, words[1])
+		}
+	}
+}
+
+// TestArm64AssembleExtensionRefusals pins the diagnostics a pinned statement
+// gets from the layer instead of a scalar path's register complaint.
+func TestArm64AssembleExtensionRefusals(t *testing.T) {
+	tests := []struct {
+		stmt string
+		want string
+	}{
+		{"ADD Z0, Z1.S, Z2.S", "carries no arrangement suffix"},
+		{"ADD Z33.S, Z1.S, Z2.S", "outside Z0-Z31"},
+		{"ADD Z1.S, P0/Z, Z0.S", "/M"},
+		{"ADD Z1.S, P9/M, Z0.S", "outside P0-P7"},
+		{"ADD $300, Z0.S", "immediate 300"},
+		{"ADD $255<<8, Z0.S", "not an immediate the layer can read"},
+		{"ADD Z1.S, P0/M, R0", "not an extended-layer operand"},
+		{"ADD Z1.S, P0/B, Z0.S", "not an extended-layer operand"},
+		{"MUL $200, Z0.B", "outside the signed 8-bit range"},
+		{"ADD Z1.S, LSL #8, Z0.S", "LSL belongs straight after an immediate"},
+	}
+	for _, tt := range tests {
+		got := assembleArm64SrcError(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
+		if !strings.Contains(got, tt.want) {
+			t.Errorf("%s: error %q does not name %q", tt.stmt, got, tt.want)
+		}
+	}
+}
+
+// TestArm64AssembleExtensionLabelOffsets proves pass 1 and pass 2 agree on a
+// function that mixes the layer with ordinary statements: the label after an
+// SVE instruction lands on the 4 bytes the encoder laid down, and the branch
+// back to it encodes the distance in words.
+func TestArm64AssembleExtensionLabelOffsets(t *testing.T) {
+	src := arm64ExtProbeHead + `
+	ADD Z0.S, Z1.S, Z2.S
+loop:
+	ADD $255, Z0.S
+	B loop
+	MUL $-128, Z0.B
+	RET
+`
+	words := assembleArm64Words(t, src)
+	want := []uint32{0x04a10002, 0x25a0dfe0, 0x17ffffff, 0x2530d000, 0xd65f03c0}
+	if len(words) != len(want) {
+		t.Fatalf("got %d words, want %d", len(words), len(want))
+	}
+	for i := range want {
+		if words[i] != want[i] {
+			t.Errorf("word %d: got %08x, want %08x", i, words[i], want[i])
+		}
+	}
+}
+
+// TestArm64AssembleExtensionLeavesScalarsAlone pins the non-invasion
+// promise: statements whose operands the scalar paths already read keep
+// their exact encodings, scalar and NEON alike, with the layer wired in.
+func TestArm64AssembleExtensionLeavesScalarsAlone(t *testing.T) {
+	tests := []struct {
+		stmt string
+		want uint32
+	}{
+		{"ADD R0, R1, R2", 0x8b000022},
+		{"ADD $255, R0", 0x9103fc00},
+		{"SUB R0, R1, R2", 0xcb000022},
+	}
+	for _, tt := range tests {
+		words := assembleArm64Words(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
+		if len(words) != 2 {
+			t.Fatalf("%s: got %d words, want the statement and RET", tt.stmt, len(words))
+		}
+		if words[0] != tt.want {
+			t.Errorf("%s:\n  got  %08x\n  want %08x", tt.stmt, words[0], tt.want)
+		}
+	}
+}