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.

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(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
		}
		if ext, ok := arm64ExtCounter(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: 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 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 {
	if len(ops) == 0 {
		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.  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
}

// 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
}