gasm-sdk/asm/extension.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

// The extended-instruction registry: the lookup over and above the generated
// architecture tables.  The generated tables (arch/*_gen.go) list the
// mnemonics the Go toolchain knows; the extension layer carries the
// instructions it does not, and this file indexes them per architecture so
// the assembler and the linter can consult the layer without touching the
// generated lists or the main encoders.  A later hook wires
// ExtensionEncodable into the Encodable mirror and EncodeExtension into the
// per-architecture assembly paths; nothing existing changes until then.

package asm

import (
	"fmt"
	"slices"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-sdk/arch"
)

// extensionIndex is the per-architecture index of the extension layer, keyed
// by upper-case mnemonic.  One mnemonic registers several forms (the SVE ADD
// carries unpredicated, predicated and immediate shapes), so the value is the
// full candidate list in table order.
type extensionIndex struct {
	byName map[string][]arch.ExtInstr
}

// extensionIndexes builds one index per known architecture.  Architectures
// whose extension layer is not built yet get an empty index, which keeps the
// queries answering false rather than failing on a missing entry.
var extensionIndexes = buildExtensionIndexes()

func buildExtensionIndexes() map[arch.Arch]*extensionIndex {
	m := make(map[arch.Arch]*extensionIndex)
	for _, a := range []arch.Arch{arch.AMD64, arch.ARM64, arch.RISCV, arch.LOONG64} {
		idx := &extensionIndex{byName: make(map[string][]arch.ExtInstr)}
		for _, in := range arch.Extensions(a) {
			key := strings.ToUpper(in.Name)
			idx.byName[key] = append(idx.byName[key], in)
		}
		m[a] = idx
	}
	return m
}

// LookupExtension returns the extended instructions registered for the
// mnemonic on a, outside the generated architecture table.  It reports false
// when a carries no extended layer or the mnemonic is not in it; a mnemonic
// the base table knows is not thereby covered, the layers stay independent.
func LookupExtension(a arch.Arch, mnemonic string) ([]arch.ExtInstr, bool) {
	idx, ok := extensionIndexes[a]
	if !ok || idx == nil {
		return nil, false
	}
	cands, ok := idx.byName[strings.ToUpper(mnemonic)]
	return cands, ok && len(cands) > 0
}

// ExtensionNames returns the mnemonics the extension layer of a registers,
// in table order, without duplicates.
func ExtensionNames(a arch.Arch) []string {
	var names []string
	seen := make(map[string]bool)
	for _, in := range arch.Extensions(a) {
		key := strings.ToUpper(in.Name)
		if !seen[key] {
			seen[key] = true
			names = append(names, in.Name)
		}
	}
	return names
}

// EncodeExtension encodes one extended instruction on a: it resolves the
// mnemonic through the extension registry, picks the registered form whose
// arity matches the operands and encodes against it.  The first form that
// encodes wins.  When every matching form rejects the operands, the error
// comes from the form whose operand kinds the list points at (the one with
// the most matching positions), so a mis-spelled predicate qualifier is
// diagnosed as one, not as the unpredicated form's register complaint.
func EncodeExtension(a arch.Arch, mnemonic string, ops ...arch.ExtOperand) ([]byte, error) {
	cands, ok := LookupExtension(a, mnemonic)
	if !ok {
		return nil, fmt.Errorf("%s registers no extended instruction %q", a, mnemonic)
	}
	var bestErr error
	var bestScore int
	var tried int
	for _, in := range cands {
		if in.Form.Arity() != len(ops) {
			continue
		}
		tried++
		b, err := in.Encode(ops)
		if err == nil {
			return b, nil
		}
		if score := kindScore(in.Form, ops); bestErr == nil || score > bestScore {
			bestErr, bestScore = err, score
		}
	}
	if tried == 0 {
		return nil, fmt.Errorf("%s: extended %q takes %s, got %d operands",
			a, mnemonic, extensionAritySummary(cands), len(ops))
	}
	return nil, bestErr
}

// kindScore counts the positions whose operand kind matches what the form
// wants, the tie-break that picks the most specific rejection.
func kindScore(form arch.ExtForm, ops []arch.ExtOperand) int {
	kinds := form.Kinds()
	score := 0
	for i, op := range ops {
		if i < len(kinds) && op.Kind == kinds[i] {
			score++
		}
	}
	return score
}

// ExtensionEncodable reports whether the extension layer of a encodes the
// mnemonic with these operands.  It mirrors asm.Encodable for the extension
// layer: the predicate the linter consults once the hook wires it in.
func ExtensionEncodable(a arch.Arch, mnemonic string, ops ...arch.ExtOperand) bool {
	_, err := EncodeExtension(a, mnemonic, ops...)
	return err == nil
}

// extensionAritySummary describes the operand counts the candidate forms
// take, "2 or 3" style, for the arity error.
func extensionAritySummary(cands []arch.ExtInstr) string {
	counts := make([]int, 0, len(cands))
	seen := make(map[int]bool)
	for _, in := range cands {
		n := in.Form.Arity()
		if !seen[n] {
			seen[n] = true
			counts = append(counts, n)
		}
	}
	slices.Sort(counts)
	var b strings.Builder
	for i, n := range counts {
		if i > 0 {
			if i == len(counts)-1 {
				b.WriteString(" or ")
			} else {
				b.WriteString(", ")
			}
		}
		fmt.Fprintf(&b, "%d", n)
	}
	b.WriteString(" operands")
	return b.String()
}