gasm-sdk/asm/encode.go

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

package asm

import (
	"fmt"
	"strings"
)

// Encode encodes one Plan 9 instruction (mnemonic plus operands, in source
// order) into x86-64 machine code.
func Encode(mnemonic string, ops ...Operand) ([]byte, error) {
	e := &enc{}
	if err := e.encode(mnemonic, ops); err != nil {
		return nil, err
	}
	return e.out, nil
}

type enc struct {
	out     []byte
	patches []encPatch // disp32 fields awaiting static-symbol resolution
}

// encPatch marks a 4-byte displacement field in enc.out that must receive the
// RIP-relative offset of a static symbol once the file layout is settled.
type encPatch struct {
	off    int
	name   string
	addend int64
}

func (e *enc) encode(mnem string, ops []Operand) error {
	upper := strings.ToUpper(mnem)

	// Fixed-name instructions (no size suffix).
	switch {
	case upper == "RET":
		return e.encodeRet()
	case upper == "NOP":
		return e.emit(&instr{opcode: []byte{0x90}, modrm: -1, sib: -1})
	case upper == "CALL" || upper == "JMP":
		// Through a register or memory: FF /2 (CALL) or FF /4 (JMP).
		// Anything else is a rel32 against a label resolved by the assembler.
		if len(ops) == 1 {
			switch ops[0].(type) {
			case Reg, Mem:
				return e.encodeIndirectBranch(upper, ops)
			}
		}
		opcode := []byte{0xE8}
		if upper == "JMP" {
			opcode = []byte{0xE9}
		}
		return e.encodeJmpRel(ops, opcode)
	}
	if cc, ok := condCode(upper); ok {
		return e.encodeJcc(cc, ops)
	}
	// No-operand system and string-control instructions (CPUID, RDTSC,
	// SYSCALL, the fences, UNDEF, …).
	if op, ok := noOperandTable[upper]; ok {
		if len(ops) != 0 {
			return fmt.Errorf("%s takes no operands, got %d", upper, len(ops))
		}
		return e.emit(&instr{opcode: op, modrm: -1, sib: -1})
	}
	// POPFQ/PUSHFQ are exact names: the bare POPF/PUSHF and the L spellings
	// are rejected by go tool asm in 64-bit mode, so they stay unsupported.
	switch upper {
	case "POPFQ":
		if len(ops) != 0 {
			return fmt.Errorf("POPFQ takes no operands, got %d", len(ops))
		}
		return e.emit(&instr{opcode: []byte{0x9D}, modrm: -1, sib: -1})
	case "PUSHFQ":
		if len(ops) != 0 {
			return fmt.Errorf("PUSHFQ takes no operands, got %d", len(ops))
		}
		return e.emit(&instr{opcode: []byte{0x9C}, modrm: -1, sib: -1})
	case "INT":
		return e.encodeInt(ops)
	case "LDMXCSR":
		return e.encodeMxcsr(2, ops)
	case "STMXCSR":
		return e.encodeMxcsr(3, ops)
	// CMPSD is the scalar double compare, whose predicate immediate comes
	// LAST in Plan 9 order (src, dst, $imm).
	case "CMPSD":
		return e.encodeCmpsd(ops)
	// SHA256RNDS2 carries the round constant in a literal X0 first operand.
	case "SHA256RNDS2":
		return e.encodeSha256rnds2(ops)
	// BYTE, WORD, LONG and QUAD write the immediate into the text stream
	// itself: 1, 2, 4 or 8 literal bytes, little-endian.  END is accepted
	// and ignored.  ADJSP adjusts SP by the immediate, sign-chosen between
	// the SUBQ and ADDQ forms.
	case "BYTE", "WORD", "LONG", "QUAD":
		return e.encodeData(upper, ops)
	case "END":
		return e.encodeEnd(ops)
	case "ADJSP":
		return e.encodeAdjsp(ops)
	}

	// VEX (AVX/AVX2) and EVEX (AVX-512) instructions: the trailing
	// B/W/L/Q/D is part of the mnemonic, not a size suffix, so dispatch
	// before splitSize.  EVEX suffixes (.Z, .SAE, rounding, .BCST) split
	// off the mnemonic too.
	base, sfx, err := parseEvexSuffix(upper)
	if err != nil {
		return err
	}
	if isVex(base) || isEvex(base) || isKOp(base) || isGather(base) || isScatter(base) ||
		isEvexPrefGather(base) ||
		base == "KMOVW" || base == "KMOVQ" || base == "KMOVB" || base == "KMOVD" {
		return e.encodeVec(base, ops, sfx)
	}
	if sfx.any() {
		return fmt.Errorf("%s: the suffix requires an EVEX instruction", mnem)
	}

	// CMOVcc and SETcc carry the condition in the mnemonic (CMOVLGT, SETNE).
	if strings.HasPrefix(upper, "CMOV") {
		return e.encodeCmov(upper, ops)
	}
	if strings.HasPrefix(upper, "SET") {
		return e.encodeSet(upper, ops)
	}

	base, size := splitSize(upper)
	if size == 0 {
		size = 8 // default operand size in 64-bit mode (e.g. PUSHQ)
	}
	// Legacy SSE imm8 shuffles whose names end in W/H (PSHUFLW,
	// PSHUFHW) must dispatch BEFORE the size-suffix split, and the
	// others ride along.
	if m, ok := sseShufTable[upper]; ok {
		return e.encodeSSEShuf(m, ops)
	}
	// Legacy SSE packed binaries dispatch on the full name: the packed
	// integer mnemonics carry real width suffixes (PADDB/PCMPGTW/...),
	// which the size split must not eat.
	if m, ok := sseBinTable[upper]; ok {
		return e.encodeSSEBin(m, ops)
	}
	if m, ok := sseBinTable[base]; ok {
		return e.encodeSSEBin(m, ops)
	}
	// The imm8-controlled legacy instructions, the lane extracts and inserts
	// and the packed integer shifts all dispatch on the full name: a trailing
	// width letter here belongs to the mnemonic, not to the size split.
	if m, ok := sseImm3Table[upper]; ok {
		return e.encodeSSEImm3(m, ops)
	}
	if m, ok := sseExtractTable[upper]; ok {
		return e.encodeSSEExtract(m, ops)
	}
	if m, ok := sseInsertTable[upper]; ok {
		return e.encodeSSEInsert(m, ops)
	}
	if _, ok := sseShiftImm[upper]; ok {
		return e.encodeSSEShift(upper, ops)
	}
	// PMOVMSKB ends in a width letter the size split would eat, so it
	// dispatches on the full name like the packed binaries above.
	if upper == "PMOVMSKB" {
		return e.encodePmovmskb(upper, ops)
	}
	switch base {
	case "MOV":
		return e.encodeMov(ops, size)
	// MOVD is the Go assembler's alias of MOVQ: the same byte forms, 64-bit
	// REX.W and all.
	case "MOVD":
		return e.encodeMov(ops, 8)
	case "ADD", "SUB", "AND", "OR", "XOR", "CMP", "ADC", "SBB":
		return e.encodeALU(aluOp[base], ops, size)
	case "TEST":
		return e.encodeTest(ops, size)
	case "LEA":
		return e.encodeLea(ops, size)
	case "INC", "DEC", "NEG", "NOT", "MUL", "DIV", "IDIV":
		return e.encodeUnary(unaryOp[base], ops, size)
	case "SHL", "SHR", "SAR", "SAL", "ROL", "ROR", "RCL", "RCR":
		return e.encodeShift(base, ops, size)
	case "BT", "BTS", "BTR", "BTC":
		return e.encodeBitTest(base, ops, size)
	case "XCHG":
		return e.encodeExchange(ops, size)
	case "CMPXCHG":
		return e.encodeRegRegOp(0xB0, 0xB1, base, ops, size)
	case "XADD":
		return e.encodeRegRegOp(0xC0, 0xC1, base, ops, size)
	case "CRC32":
		return e.encodeCrc32(ops, size)
	case "ADCX":
		return e.encodeCarryExt(0x66, ops, size)
	case "ADOX":
		return e.encodeCarryExt(0xF3, ops, size)
	case "MOVS", "STOS":
		return e.encodeStringOp(base, ops, size)
	case "IMUL", "IMUL3":
		return e.encodeImul(ops, size)
	case "PUSH":
		return e.encodePushPop(ops, size, true)
	case "POP":
		return e.encodePushPop(ops, size, false)
	case "BSF", "BSR", "LZCNT", "TZCNT", "POPCNT":
		return e.encodeCount(base, ops, size)
	case "BSWAP":
		return e.encodeBswap(ops, size)
	case "PREFETCHNTA", "PREFETCHT0", "PREFETCHT1", "PREFETCHT2":
		return e.encodePrefetch(base, ops)
	case "MOVBLZX", "MOVBQZX", "MOVWLZX", "MOVWQZX", "MOVWLSX", "MOVLQSX",
		"MOVBWZX", "MOVBWSX", "MOVBLSX", "MOVBQSX", "MOVWQSX", "MOVLQZX":
		return e.encodeMovExtend(base, ops)
	case "CVTSL2SD", "CVTSQ2SD":
		return e.encodeCvtsi2sd(base == "CVTSQ2SD", ops)
	case "CVTSD2S", "CVTTSD2S", "CVTSS2S", "CVTTSS2S":
		return e.encodeCvtInt(base, ops, size)
	case "FMOVD":
		return e.encodeFmov(ops)
	case "MOVOU", "MOVO", "MOVOA", "MOVUPS", "MOVAPS", "MOVUPD", "MOVAPD", "MOVSD", "MOVSS":
		return e.encodeSSEMove(sseMoveTable[base], ops)
	}
	return fmt.Errorf("unsupported instruction %q", mnem)
}

// encodePrefetch emits the 0F 18 /r prefetch hints: the reg field selects
// the locality (NTA=0, T0=1, T1=2, T2=3) and the single operand is memory.
func (e *enc) encodePrefetch(base string, ops []Operand) error {
	if len(ops) != 1 {
		return fmt.Errorf("%s expects one memory operand", base)
	}
	m, ok := ops[0].(Mem)
	if !ok {
		return fmt.Errorf("%s requires a memory operand", base)
	}
	i := newInstr(0, []byte{0x0F, 0x18})
	if err := setMem(i, prefetchVariant[base], m); err != nil {
		return err
	}
	return e.emit(i)
}

var prefetchVariant = map[string]int{
	"PREFETCHNTA": 0,
	"PREFETCHT0":  1,
	"PREFETCHT1":  2,
	"PREFETCHT2":  3,
}

// dataWidth is the literal byte count of each data-emission pseudo-op.
var dataWidth = map[string]int{
	"BYTE": 1,
	"WORD": 2,
	"LONG": 4,
	"QUAD": 8,
}

// encodeData emits the literal-data pseudo-ops: BYTE, WORD, LONG and QUAD
// write the immediate into the text stream as 1, 2, 4 or 8 bytes,
// little-endian, with no opcode lookup.  The value is truncated to the
// width rather than range-checked, exactly as go tool asm behaves (BYTE
// $0x1FF emits FF, WORD $0x12345 emits 45 23, both without an error), and
// exactly one immediate is accepted: the toolchain rejects a list such as
// BYTE $1, $2, $3.
func (e *enc) encodeData(mnem string, ops []Operand) error {
	if len(ops) != 1 {
		return fmt.Errorf("%s expects 1 immediate operand, got %d", mnem, len(ops))
	}
	imm, ok := ops[0].(Imm)
	if !ok {
		return fmt.Errorf("%s requires an integer immediate", mnem)
	}
	width := dataWidth[mnem]
	out := make([]byte, width)
	u := uint64(imm)
	for i := range width {
		out[i] = byte(u >> (8 * i))
	}
	e.out = append(e.out, out...)
	return nil
}

// encodeEnd accepts-and-ignores END.  go tool asm drops the statement
// entirely: the AEND Prog is skipped when the program list is flushed, so
// the statements after an END still belong to the same function and the
// encoded body carries no trace of it, whatever operands follow the name
// (the toolchain takes END $0 and END AX alike).  Zero bytes, no effect.
func (e *enc) encodeEnd(ops []Operand) error {
	return nil
}

// encodeAdjsp emits ADJSP $imm: a positive value is SUBQ $imm, SP, a
// negative one ADDQ $-imm, SP, in the imm8 or imm32 form the magnitude
// picks (the same selection subSP and addSP make for the frame).  go tool
// asm refuses ADJSP $0 outright, so a zero value is an error here too; the
// statement's effect on the SP balance is checked by the function-level
// assembly (checkAdjspBalance), as the toolchain's push/pop walk does.
func (e *enc) encodeAdjsp(ops []Operand) error {
	if len(ops) != 1 {
		return fmt.Errorf("ADJSP expects 1 immediate operand, got %d", len(ops))
	}
	imm, ok := ops[0].(Imm)
	if !ok {
		return fmt.Errorf("ADJSP requires an integer immediate")
	}
	switch v := int(imm); {
	case v > 0:
		e.out = append(e.out, subSP(v)...)
	case v < 0:
		e.out = append(e.out, addSP(-v)...)
	default:
		return fmt.Errorf("ADJSP $0 has no encoding")
	}
	return nil
}

// splitSize separates a trailing B/W/L/Q size suffix from the mnemonic.
func splitSize(upper string) (base string, size int) {
	if upper == "" {
		return upper, 0
	}
	switch upper[len(upper)-1] {
	case 'B':
		return upper[:len(upper)-1], 1
	case 'W':
		return upper[:len(upper)-1], 2
	case 'L':
		return upper[:len(upper)-1], 4
	case 'Q':
		return upper[:len(upper)-1], 8
	}
	return upper, 0
}

// encodeVec dispatches a VEX/EVEX mnemonic to the right encoding: KMOVW has
// its own direction-dependent opcodes; KTESTW is always VEX; everything else
// takes EVEX when an operand demands it (a ZMM or K register, or an
// EVEX-only mnemonic) and VEX otherwise.
func (e *enc) encodeVec(upper string, ops []Operand, sfx evexSuffix) error {
	if gs, ok := gatherTable[upper]; ok {
		return e.encodeGather(upper, gs, ops, sfx)
	}
	if ss, ok := scatterTable[upper]; ok {
		return e.encodeScatter(upper, ss, ops, sfx)
	}
	if upper == "KMOVW" || upper == "KMOVQ" || upper == "KMOVB" || upper == "KMOVD" {
		if sfx.any() {
			return fmt.Errorf("%s takes no EVEX suffixes", upper)
		}
		return e.encodeKmov(upper, ops)
	}
	if isKOp(upper) {
		if sfx.any() {
			return fmt.Errorf("%s takes no EVEX suffixes", upper)
		}
		return e.encodeKOp(upper, ops)
	}
	if upper == "KTESTW" || (!evexRequired(upper, ops) && !sfx.evexOnly()) {
		if sfx.any() {
			return fmt.Errorf("%s: the .Z suffix requires an EVEX instruction", upper)
		}
		return e.encodeVex(upper, ops)
	}
	return e.encodeEvex(upper, ops, sfx)
}

// --- instruction components -------------------------------------------------

type instr struct {
	opSize16  bool
	rexW      bool
	rexR      bool
	rexX      bool
	rexB      bool
	rexForced bool // REX needed even with all bits zero (8-bit low registers)
	prefix    byte // legacy 0xF2/0xF3 prefix (0 = none); emitted after 0x66
	opcode    []byte
	modrm     int // -1 if absent
	sib       int // -1 if absent
	disp      []byte
	imm       []byte
	sb        *sbRef // static-symbol displacement in disp, awaiting resolution
}

// sbRef records that an instruction's displacement refers to a static symbol
// rather than holding a literal value.
type sbRef struct {
	name   string
	addend int64
}

func (e *enc) emit(i *instr) error {
	if i.opSize16 {
		e.out = append(e.out, 0x66)
	}
	if i.prefix != 0 {
		e.out = append(e.out, i.prefix)
	}
	rex := byte(0)
	if i.rexW {
		rex |= 0x08
	}
	if i.rexR {
		rex |= 0x04
	}
	if i.rexX {
		rex |= 0x02
	}
	if i.rexB {
		rex |= 0x01
	}
	if rex != 0 || i.rexForced {
		e.out = append(e.out, 0x40|rex)
	}
	e.out = append(e.out, i.opcode...)
	if i.modrm >= 0 {
		e.out = append(e.out, byte(i.modrm))
	}
	if i.sib >= 0 {
		e.out = append(e.out, byte(i.sib))
	}
	if i.sb != nil {
		e.patches = append(e.patches, encPatch{off: len(e.out), name: i.sb.name, addend: i.sb.addend})
	}
	e.out = append(e.out, i.disp...)
	e.out = append(e.out, i.imm...)
	return nil
}

// newInstr starts an instruction with a size-derived REX.W and 0x66 prefix.
func newInstr(opSize int, opcode []byte) *instr {
	return &instr{
		opSize16: opSize == 2,
		rexW:     opSize == 8,
		opcode:   opcode,
		modrm:    -1,
		sib:      -1,
	}
}

// --- ModR/M, SIB, displacement ----------------------------------------------

// setRM fills in the ModR/M (and SIB, displacement, REX bits) for an
// instruction whose reg field holds a real register `reg` and whose r/m field
// holds `rm`.
func setRM(i *instr, reg Reg, rm Operand, opSize int) error {
	return setRMReg(i, reg.idx&7, reg.idx >= 8, reg.needsREX(opSize), rm, opSize)
}

// setRMDigit fills in the ModR/M for an instruction whose reg field is an
// opcode /digit extension (0-7), which carries none of the register REX rules.
func setRMDigit(i *instr, digit int, rm Operand, opSize int) error {
	return setRMReg(i, digit, false, false, rm, opSize)
}

func setRMReg(i *instr, regField int, rexR, regForced bool, rm Operand, opSize int) error {
	i.rexR = rexR
	if regForced {
		i.rexForced = true
	}

	switch r := rm.(type) {
	case Reg:
		i.rexB = r.idx >= 8
		if r.needsREX(opSize) {
			i.rexForced = true
		}
		i.modrm = 0xC0 | regField<<3 | (r.idx & 7)
		return nil
	case Mem:
		return setMem(i, regField, r)
	case sbMem:
		// RIP-relative reference; the displacement is patched once the static
		// symbol's address is known.
		i.modrm = regField<<3 | 0x05 // mod=00, rm=101 → (RIP)+disp32
		i.disp = le32(0)
		i.sb = &sbRef{name: r.name, addend: r.addend}
		return nil
	default:
		return fmt.Errorf("invalid r/m operand %T", rm)
	}
}

func setMem(i *instr, regField int, m Mem) error {
	modrm, sib, disp, xBit, bBit, err := memComponents(regField, m)
	if err != nil {
		return err
	}
	i.modrm = modrm
	i.sib = sib
	i.disp = disp
	i.rexX = xBit == 1
	i.rexB = bBit == 1
	return nil
}

// memComponents computes the ModR/M byte (with the given reg field), the SIB
// byte (-1 if none), the displacement bytes, and the high index/base bits, for
// a memory operand.  It is shared by the REX (scalar) and VEX (vector) paths.
func memComponents(regField int, m Mem) (modrm, sib int, disp []byte, xBit, bBit int, err error) {
	sib = -1
	// A displacement wider than int32 fits no encoding form; truncating it
	// would address a different location, and go tool asm reports "offset
	// too large" for the same operand.
	if m.Disp < -(1<<31) || m.Disp > (1<<31)-1 {
		return 0, -1, nil, 0, 0, fmt.Errorf("displacement %d does not fit in 32 bits", m.Disp)
	}
	// RIP-relative: neither base nor index.
	if !m.HasBase && !m.HasIndex {
		return regField<<3 | 0x05, -1, le32(m.Disp), 0, 0, nil // mod=00, rm=101
	}

	// The SIB scale field only encodes 1/2/4/8; the Go assembler rejects
	// anything else ("bad scale: 16"), so a silent fallback to scale 1 here
	// would mis-assemble the operand instead of reporting it.
	if m.HasIndex && m.Scale != 1 && m.Scale != 2 && m.Scale != 4 && m.Scale != 8 {
		return 0, -1, nil, 0, 0, fmt.Errorf("bad scale: %d", m.Scale)
	}

	needSIB := m.HasIndex || (m.HasBase && m.Base.idx&7 == 4)

	var mod int
	switch {
	case !m.HasBase:
		mod = 0
		disp = le32(m.Disp)
	case m.Base.idx&7 == 5 && m.Disp == 0:
		mod = 1
		disp = []byte{0}
	case m.Disp == 0:
		mod = 0
	case fits8(m.Disp):
		mod = 1
		disp = []byte{byte(int8(m.Disp))}
	default:
		mod = 2
		disp = le32(m.Disp)
	}

	if needSIB {
		idxField := 4 // 100 = no index
		if m.HasIndex {
			idxField = m.Index.idx & 7
			if m.Index.idx >= 8 {
				xBit = 1
			}
		}
		baseField := 5 // 101 = no base (with mod=00 → disp32)
		if m.HasBase {
			baseField = m.Base.idx & 7
			if m.Base.idx >= 8 {
				bBit = 1
			}
		}
		return mod<<6 | regField<<3 | 0x04, scaleBits(m.Scale)<<6 | idxField<<3 | baseField, disp, xBit, bBit, nil
	}

	if m.Base.idx >= 8 {
		bBit = 1
	}
	return mod<<6 | regField<<3 | (m.Base.idx & 7), -1, disp, 0, bBit, nil
}

func scaleBits(scale int) int {
	switch scale {
	case 2:
		return 1
	case 4:
		return 2
	case 8:
		return 3
	default:
		return 0 // scale 1 (or unset)
	}
}

func fits8(v int64) bool { return v >= -128 && v <= 127 }

func le32(v int64) []byte {
	u := uint32(v)
	return []byte{byte(u), byte(u >> 8), byte(u >> 16), byte(u >> 24)}
}

func le16(v int64) []byte {
	u := uint16(v)
	return []byte{byte(u), byte(u >> 8)}
}

func le64(v int64) []byte {
	u := uint64(v)
	b := make([]byte, 8)
	for i := range 8 {
		b[i] = byte(u >> (8 * i))
	}
	return b
}