gasm-sdk/asm/arm64_ext.go

// 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, and for the predicate family Pm.B, Pn.B, Pg/Z (or Pg.Z), Pd.B
// for the logical operations, Pn.B, Pg.Z, Pd.B for the breaks, Pm.T, Pn.T,
// Pd.T for the permutations, Rm, Rn, Pd.T for the while compares, PN8-PN15
// for the counter destinations, and the bare SETFFR.  Stage three adds the
// crypto family (Zn.T, Zd.T, Zd.T read-back and the in-place Zd.T, Zd.T),
// the predicate counters (Pn.T, Pg, Rd; Pn.T, ZR; Rd, Pn.T, Rd; ZR and R
// terminators) and the reductions (Zn.T, Pg, Vd over the SIMD register
// V0-V31, with ZR and RSP accepted where the classes take them).

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, predicate or predicate-as-counter
// register, or no operands at all (the zero-operand forms such as SETFFR,
// which no scalar path could mean instead).  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(mnem, 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
		}
		// The gather/scatter destination list, [Z13.B]: one scalable vector
		// in brackets, its arrangement part of the instruction's identity.
		if strings.HasPrefix(text, "[") && strings.HasSuffix(text, "]") {
			if ext, ok := arm64ExtVector(strings.Trim(text, "[]")); ok {
				out = append(out, ext)
				continue
			}
		}
		// The gather/scatter memory operand: a parenthesised register pair,
		// an immediate-offset base or a lone vector base.
		if ext, ok := arm64ExtSveMem(text); ok {
			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
		}
		if ext, ok := arm64ExtCounter(text); ok {
			out = append(out, ext)
			continue
		}
		if text == "ZR" {
			out = append(out, arch.ExtZeroRegister())
			continue
		}
		if text == "RSP" {
			out = append(out, arch.ExtStackPointer())
			continue
		}
		if ext, ok := arm64ExtSIMD(text); ok {
			out = append(out, ext)
			continue
		}
		if ext, ok := arm64ExtGeneral(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, predicate, general or counter register, or an immediate", mnem, i+1, op.Raw)
	}
	return out, true, nil
}

// arm64ExtPinned reports whether the statement belongs to the layer.  A
// mnemonic the extension layer registers on its own, one the generated
// arm64 table does not know, owns every one of its statements: no scalar
// path could mean it instead, and the layer's diagnostics replace the
// unsupported-instruction complaint.  A mnemonic both tables carry (the
// SVE aliases of ADD, SUB and MUL) keeps the operand-shape test: any
// operand is a scalable vector, predicate or predicate-as-counter
// register, the shapes only the extension layer reads, or the statement
// carries no operands at all and the mnemonic's zero-operand forms claim
// it.  The shape 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(mnem string, ops []*ast.Operand) bool {
	if len(ops) == 0 {
		return true
	}
	if _, shared := a64InstrTable[mnem]; !shared {
		return true
	}
	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 or
// predicate-as-counter register at all: PN or P, digits, an optional
// arrangement suffix and an optional qualifier after a slash, whatever the
// qualifier says.  The strict parses in arm64ExtPredicate and
// arm64ExtCounter judge 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 rest, found := strings.CutPrefix(text, "N"); found {
		text = rest
	}
	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 (P0.B) and an optional qualifier in either
// spelling the corpus and the wired forms use, P0/M and P0.Z.
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
	} else if base, suffix, found := strings.Cut(text, "."); found &&
		(suffix == "Z" || suffix == "M") {
		// The dot qualifier stands in place of an arrangement, the spelling
		// the toolchain's corpus writes (P1.Z, P14.M).
		qual = arch.ExtQualMerging
		if suffix == "Z" {
			qual = arch.ExtQualZeroing
		}
		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
}

// arm64ExtCounter parses a predicate-as-counter register operand: PN8..PN15
// with an optional element-size suffix, PN14.S.  The register range is the
// counter range the layer's convention carries; the encoding validates it.
func arm64ExtCounter(text string) (arch.ExtOperand, bool) {
	rest, ok := strings.CutPrefix(text, "PN")
	if !ok {
		return arch.ExtOperand{}, false
	}
	reg, arr, ok := arm64ExtRegDigits(rest)
	if !ok {
		return arch.ExtOperand{}, false
	}
	return arch.ExtOperand{Kind: arch.ExtPNReg, Reg: reg, Arr: arr}, true
}

// arm64ExtGeneral parses a general register operand: R0..R30, the plain
// spelling the while-compare forms take, beside the ZR and RSP spellings of
// the thirty-first slot the conversion above reads.  The register range is
// left to the encoding, whose diagnostics name it.
func arm64ExtGeneral(text string) (arch.ExtOperand, bool) {
	rest, ok := strings.CutPrefix(text, "R")
	if !ok {
		return arch.ExtOperand{}, false
	}
	reg, arr, ok := arm64ExtRegDigits(rest)
	if !ok || arr != arch.ExtArrNone {
		return arch.ExtOperand{}, false
	}
	return arch.ExtOperand{Kind: arch.ExtGReg, Reg: reg}, true
}

// arm64ExtSIMD parses a 128-bit SIMD register operand: V0..V31, written
// bare, the scalar destination the reductions and the crypto read-back
// forms take.  The register range is left to the encoding.
func arm64ExtSIMD(text string) (arch.ExtOperand, bool) {
	rest, ok := strings.CutPrefix(text, "V")
	if !ok {
		return arch.ExtOperand{}, false
	}
	reg, arr, ok := arm64ExtRegDigits(rest)
	if !ok || arr != arch.ExtArrNone {
		return arch.ExtOperand{}, false
	}
	return arch.ExtOperand{Kind: arch.ExtVReg, Reg: reg}, true
}

// arm64ExtSveMem parses the gather/scatter memory operand off a normalised
// operand text: the parenthesised pair (R6)(R14), (Z23.D<<1)(R24) and
// (Z4.S.UXTW)(R3), the immediate-offset base 6(Z7.S), and the lone vector
// base (Z5.D) of the stores.  The second parenthesis accepts the
// stack-pointer spelling RSP; the ranges and the mode's own rules are left
// to the encoding, whose diagnostics name them.
func arm64ExtSveMem(text string) (arch.ExtOperand, bool) {
	// The immediate-offset spelling: digits straight before the parenthesis.
	if i := strings.IndexByte(text, '('); i > 0 && i == strings.LastIndexByte(text, '(') {
		disp, err := strconv.ParseUint(text[:i], 10, 32)
		if err == nil && strings.HasSuffix(text, ")") {
			op, ok := arm64ExtSveMemGroup(text[i+1 : len(text)-1])
			if !ok {
				return arch.ExtOperand{}, false
			}
			if !op.BaseVec || op.Extend != 0 || op.Shift != 0 {
				return arch.ExtOperand{}, false
			}
			op.Imm = int64(disp)
			return op, true
		}
	}
	// The parenthesised forms: one group or two.
	rest, ok := strings.CutPrefix(text, "(")
	if !ok || !strings.HasSuffix(text, ")") {
		return arch.ExtOperand{}, false
	}
	rest = rest[:len(rest)-1]
	first := rest
	op := arch.ExtOperand{Off: -1}
	if base, second, found := strings.Cut(rest, ")("); found {
		first = base
		off, ok := arm64ExtSveMemOffset(second)
		if !ok {
			return arch.ExtOperand{}, false
		}
		op = off
	}
	group, ok := arm64ExtSveMemGroup(first)
	if !ok {
		return arch.ExtOperand{}, false
	}
	group.Off = op.Off
	group.Reg31 = op.Reg31
	return group, true
}

// arm64ExtSveMemGroup parses one parenthesised memory register: R6, R6<<3,
// Z23.D, Z23.D<<1, Z4.S.UXTW or Z7.D.SXTW.  The general registers run
// R0-R30 and the scalable vectors Z0-Z31 with an .S or .D element size and
// an optional UXTW or SXTW extension; the ranges are left to the encoding.
func arm64ExtSveMemGroup(text string) (arch.ExtOperand, bool) {
	op := arch.ExtOperand{Kind: arch.ExtSveMem, Off: -1}
	if base, shift, found := strings.Cut(text, "<<"); found {
		n, err := strconv.Atoi(shift)
		if err != nil || n < 0 {
			return arch.ExtOperand{}, false
		}
		op.Shift = n
		text = base
	}
	parts := strings.Split(text, ".")
	switch parts[0][0] {
	case 'R':
		reg, err := strconv.Atoi(parts[0][1:])
		if err != nil || len(parts) != 1 {
			return arch.ExtOperand{}, false
		}
		op.Reg = reg
	case 'Z':
		reg, err := strconv.Atoi(parts[0][1:])
		if err != nil || len(parts) < 2 || len(parts) > 3 {
			return arch.ExtOperand{}, false
		}
		switch parts[1] {
		case "S":
			op.Arr = arch.ExtArrS
		case "D":
			op.Arr = arch.ExtArrD
		default:
			return arch.ExtOperand{}, false
		}
		op.Reg = reg
		op.BaseVec = true
		if len(parts) == 3 {
			switch parts[2] {
			case "UXTW":
				op.Extend = 1
			case "SXTW":
				op.Extend = 2
			default:
				return arch.ExtOperand{}, false
			}
		}
	default:
		return arch.ExtOperand{}, false
	}
	return op, true
}

// arm64ExtSveMemOffset parses the second parenthesis of a gather/scatter
// memory operand: a plain R0-R30 or the stack-pointer spelling RSP.
func arm64ExtSveMemOffset(text string) (arch.ExtOperand, bool) {
	if text == "RSP" {
		return arch.ExtOperand{Off: 31, Reg31: 2}, true
	}
	if len(text) < 2 || text[0] != 'R' {
		return arch.ExtOperand{}, false
	}
	reg, err := strconv.Atoi(text[1:])
	if err != nil {
		return arch.ExtOperand{}, false
	}
	return arch.ExtOperand{Off: reg}, true
}

// arm64ExtRegDigits parses the digits and optional arrangement suffix of a
// register spelling once the letter prefix is gone.
func arm64ExtRegDigits(text string) (reg int, arr arch.ExtArrangement, ok bool) {
	if base, suffix, found := strings.Cut(text, "."); 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
		}
		text = base
	}
	n, err := strconv.Atoi(text)
	if err != nil || n < 0 {
		return 0, 0, false
	}
	return n, arr, 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
}