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

package asm

import "fmt"

// This file implements the system, flag, string and segment families the Go
// assembler carries: the no-operand controls, the string primitives, the
// sign-extension pair, the multi-byte no-ops, the cache controls, MOVBE, the
// compare-exchange doubles, the random source pair, the FS/GS base pair, the
// descriptor-table controls and the LAR/LSL selector reads.  Every encoding
// here is pinned byte for byte against go tool asm through the corpus lines
// in amd64_system_test.go.

// systemNoOperand maps a fixed no-operand mnemonic to its opcode bytes, the
// prefixes spelled out in full.
var systemNoOperand = map[string][]byte{
	"CLC":        {0xF8},
	"STC":        {0xF9},
	"CMC":        {0xF5},
	"CLI":        {0xFA},
	"STI":        {0xFB},
	"HLT":        {0xF4},
	"ICEBP":      {0xF1},
	"XLAT":       {0xD7},
	"LAHF":       {0x9F},
	"SAHF":       {0x9E},
	"PUSHFW":     {0x66, 0x9C},
	"POPFW":      {0x66, 0x9D},
	"IRETW":      {0x66, 0xCF},
	"IRETL":      {0xCF},
	"IRETQ":      {0x48, 0xCF},
	"UD1":        {0x0F, 0xB9},
	"UD2":        {0x0F, 0x0B},
	"CLAC":       {0x0F, 0x01, 0xCA},
	"STAC":       {0x0F, 0x01, 0xCB},
	"CLTS":       {0x0F, 0x06},
	"INVD":       {0x0F, 0x08},
	"WBINVD":     {0x0F, 0x09},
	"SWAPGS":     {0x0F, 0x01, 0xF8},
	"RSM":        {0x0F, 0xAA},
	"MONITOR":    {0x0F, 0x01, 0xC8},
	"MWAIT":      {0x0F, 0x01, 0xC9},
	"RDMSR":      {0x0F, 0x32},
	"WRMSR":      {0x0F, 0x30},
	"RDPMC":      {0x0F, 0x33},
	"RDPKRU":     {0x0F, 0x01, 0xEE},
	"WRPKRU":     {0x0F, 0x01, 0xEF},
	"XSETBV":     {0x0F, 0x01, 0xD1},
	"SYSENTER":   {0x0F, 0x34},
	"SYSENTER64": {0x48, 0x0F, 0x34},
	"SYSEXIT":    {0x0F, 0x35},
	"SYSEXIT64":  {0x48, 0x0F, 0x35},
	"SYSRET":     {0x0F, 0x07},
	"CBW":        {0x66, 0x98},
	"CWDE":       {0x98},
	"CDQE":       {0x48, 0x98},
	"CWD":        {0x66, 0x99},
	"CDQ":        {0x99},
	"CQO":        {0x48, 0x99},
}

// stringOp maps the string-primitive bases to their 32-bit opcode; the byte
// form is one lower, the word spelling carries 0x66 and the quad spelling
// REX.W, exactly the prefix ladder newInstr applies.
var stringOp = map[string]byte{
	"CMPS": 0xA7,
	"INS":  0x6D,
	"LODS": 0xAD,
	"OUTS": 0x6F,
	"SCAS": 0xAF,
}

// nopWidth maps the multi-byte no-op spellings to their operand size.
var nopWidth = map[string]int{
	"NOPW": 2,
	"NOPL": 4,
	"NOPQ": 8,
}

// cacheControl maps the one-memory-operand cache controls to their mandatory
// prefix, opcode group and /digit.
var cacheControl = map[string]struct {
	prefix byte
	op     []byte
	digit  int
}{
	"CLFLUSH":    {0, []byte{0x0F, 0xAE}, 7},
	"CLFLUSHOPT": {0x66, []byte{0x0F, 0xAE}, 7},
	"INVLPG":     {0, []byte{0x0F, 0x01}, 7},
}

// movbeSize maps the MOVBE spellings to their operand size.
var movbeSize = map[string]int{
	"MOVBEW": 2,
	"MOVBEL": 4,
	"MOVBEQ": 8,
}

// randSource maps the random-source bases to their /digit (RDRAND /6,
// RDSEED /7); the destination register rides r/m, mod 11.
var randSource = map[string]int{
	"RDRAND": 6,
	"RDSEED": 7,
}

// fsGsBase maps the FS/GS base accessors to their /digit in the F3-prefixed
// 0F AE group; the L and Q spellings exist.
var fsGsBase = map[string]int{
	"RDFSBASE": 0,
	"RDGSBASE": 1,
	"WRFSBASE": 2,
	"WRGSBASE": 3,
}

// descTable maps the descriptor-table accesses to their /digit in 0F 01;
// each takes one memory operand alone.
var descTable = map[string]int{
	"LGDT": 2,
	"LIDT": 3,
	"SGDT": 0,
	"SIDT": 1,
}

// sysRmEntry is one 0F 00/01 register-or-memory access.  sized marks the
// members whose trailing width letter (SLDTW, STRQ, SMSWL) carries the width
// prefix ladder; the rest are fixed-width single names.
type sysRmEntry struct {
	group byte
	digit int
	sized bool
}

// sysRm maps the system register accesses LLDT/LTR/VERR/VERW/SLDT/STR (group
// 0F 00), LMSW/SMSW (0F 01).
var sysRm = map[string]sysRmEntry{
	"LLDT": {0x00, 2, false},
	"LTR":  {0x00, 3, false},
	"VERR": {0x00, 4, false},
	"VERW": {0x00, 5, false},
	"SLDT": {0x00, 0, true},
	"STR":  {0x00, 1, true},
	"LMSW": {0x01, 6, false},
	"SMSW": {0x01, 4, true},
}

// selectorRead maps the selector reads LAR and LSL to their opcodes; both
// load the destination register from an r/m selector, width prefixes per the
// suffix.
var selectorRead = map[string]byte{
	"LAR": 0x02,
	"LSL": 0x03,
}

// farSegLoad maps the far-segment loads to their opcodes; memory source
// alone, destination register, width prefixes per the suffix.
var farSegLoad = map[string]byte{
	"LFS": 0xB4,
	"LGS": 0xB5,
	"LSS": 0xB2,
}

// encodeSystem encodes the system, flag, string and segment families.  It
// reports whether the mnemonic belongs to the family.
func (e *enc) encodeSystem(upper string, ops []Operand) (bool, error) {
	if op, ok := systemNoOperand[upper]; ok {
		if len(ops) != 0 {
			return true, fmt.Errorf("%s takes no operands, got %d", upper, len(ops))
		}
		return true, e.emit(&instr{opcode: append([]byte(nil), op...), modrm: -1, sib: -1})
	}
	// The string primitives carry a B/W/L/Q suffix only; CMPSD and friends
	// are the SSE compare family's names and must reach their own dispatch.
	if b, size := splitSize(upper); size != 0 {
		switch upper[len(upper)-1] {
		case 'B', 'W', 'L', 'Q':
			if op32, ok := stringOp[b]; ok {
				return true, e.encodeSystemString(upper, op32, ops)
			}
		}
	}
	if size, ok := nopWidth[upper]; ok {
		if len(ops) != 1 {
			return true, fmt.Errorf("%s expects 1 operand, got %d", upper, len(ops))
		}
		i := newInstr(size, []byte{0x0F, 0x1F})
		if err := setRMDigit(i, 0, ops[0], size); err != nil {
			return true, err
		}
		return true, e.emit(i)
	}
	if m, ok := cacheControl[upper]; ok {
		if len(ops) != 1 {
			return true, fmt.Errorf("%s expects 1 memory operand, got %d", upper, len(ops))
		}
		if !isX86Mem(ops[0]) {
			return true, fmt.Errorf("%s requires a memory operand", upper)
		}
		i := &instr{prefix: m.prefix, opcode: m.op, modrm: -1, sib: -1}
		if err := setRMDigit(i, m.digit, ops[0], 8); err != nil {
			return true, err
		}
		return true, e.emit(i)
	}
	if _, ok := movbeSize[upper]; ok {
		return true, e.encodeSystemMovbe(upper, ops)
	}
	if digit, ok := randSource[base(upper, 6)]; ok {
		return true, e.encodeSystemRand(upper, digit, ops)
	}
	if digit, ok := fsGsBase[base(upper, 8)]; ok {
		return true, e.encodeSystemFsGsBase(upper, digit, ops)
	}
	if digit, ok := descTable[upper]; ok {
		if len(ops) != 1 {
			return true, fmt.Errorf("%s expects 1 memory operand, got %d", upper, len(ops))
		}
		if !isX86Mem(ops[0]) {
			return true, fmt.Errorf("%s requires a memory operand", upper)
		}
		i := &instr{opcode: []byte{0x0F, 0x01}, modrm: -1, sib: -1}
		if err := setRMDigit(i, digit, ops[0], 8); err != nil {
			return true, err
		}
		return true, e.emit(i)
	}
	if m, ok := sysRm[upper]; ok {
		return true, e.encodeSystemRm(upper, m, ops)
	}
	if b, size := splitSize(upper); size != 0 {
		if m, ok := sysRm[b]; ok && m.sized {
			return true, e.encodeSystemRm(upper, m, ops)
		}
	}
	if op, ok := selectorRead[base(upper, 3)]; ok {
		return true, e.encodeSystemSelectorRead(upper, op, ops)
	}
	if op, ok := farSegLoad[base(upper, 3)]; ok {
		return true, e.encodeSystemFarLoad(upper, op, ops)
	}
	if upper == "CMPXCHG8B" || upper == "CMPXCHG16B" {
		if len(ops) != 1 {
			return true, fmt.Errorf("%s expects 1 memory operand, got %d", upper, len(ops))
		}
		if !isX86Mem(ops[0]) {
			return true, fmt.Errorf("%s requires a memory operand", upper)
		}
		i := newInstr(0, []byte{0x0F, 0xC7})
		i.rexW = upper == "CMPXCHG16B"
		if err := setRMDigit(i, 1, ops[0], 8); err != nil {
			return true, err
		}
		return true, e.emit(i)
	}
	return false, nil
}

// base returns the first n characters of an upper-case mnemonic, or the empty
// string when the mnemonic is shorter: the safe head lookup for the families
// whose width suffix rides the tail (RDRANDW, RDFSBASEQ, LARW).
func base(upper string, n int) string {
	if len(upper) <= n {
		return ""
	}
	return upper[:n]
}

// encodeSystemString encodes a string primitive: no operands, the width
// suffix picks the byte form, the 0x66 prefix or REX.W.  Only the B/W/L/Q
// suffixes belong to the family: CMPSD and friends are the SSE compare
// family's names and must reach their own dispatch.
func (e *enc) encodeSystemString(upper string, op32 byte, ops []Operand) error {
	switch upper[len(upper)-1] {
	case 'B', 'W', 'L', 'Q':
	default:
		return fmt.Errorf("unsupported instruction %q", upper)
	}
	b, size := splitSize(upper)
	if _, ok := stringOp[b]; !ok || size == 0 {
		return fmt.Errorf("unsupported instruction %q", upper)
	}
	if len(ops) != 0 {
		return fmt.Errorf("%s takes no operands, got %d", upper, len(ops))
	}
	// The byte spelling is the 32-bit opcode minus one; the W and Q forms
	// ride newInstr's prefix ladder, the L form the bare opcode.
	op := op32
	if size == 1 {
		op--
	}
	return e.emit(newInstr(size, []byte{op}))
}

// encodeSystemMovbe encodes MOVBE: a register source stores (F1, reg = the
// register, r/m = memory), a register destination loads (F0, same fields).
func (e *enc) encodeSystemMovbe(mnem string, ops []Operand) error {
	size := movbeSize[mnem]
	if len(ops) != 2 {
		return fmt.Errorf("%s expects 2 operands, got %d", mnem, len(ops))
	}
	srcReg, srcIsReg := ops[0].(Reg)
	dstReg, dstIsReg := ops[1].(Reg)
	var op byte
	var reg Reg
	var rm Operand
	switch {
	case srcIsReg && isX86Mem(ops[1]):
		op, reg, rm = 0xF1, srcReg, ops[1] // store
	case dstIsReg && isX86Mem(ops[0]):
		op, reg, rm = 0xF0, dstReg, ops[0] // load
	default:
		return fmt.Errorf("%s takes one register and one memory operand", mnem)
	}
	i := newInstr(size, []byte{0x0F, 0x38, op})
	if err := setRM(i, reg, rm, size); err != nil {
		return err
	}
	return e.emit(i)
}

// encodeSystemRand encodes RDRAND/RDSEED: the single register operand rides
// r/m under the /digit, mod 11, with the width prefix the suffix picks.
func (e *enc) encodeSystemRand(mnem string, digit int, ops []Operand) error {
	b, size := splitSize(mnem)
	if _, ok := randSource[b]; !ok || size == 0 {
		return fmt.Errorf("unsupported instruction %q", mnem)
	}
	if len(ops) != 1 {
		return fmt.Errorf("%s expects 1 register operand, got %d", mnem, len(ops))
	}
	dstReg, ok := ops[0].(Reg)
	if !ok || dstReg.isVec() {
		return fmt.Errorf("%s destination must be a general register", mnem)
	}
	i := newInstr(size, []byte{0x0F, 0xC7})
	if err := setRMDigit(i, digit, dstReg, size); err != nil {
		return err
	}
	return e.emit(i)
}

// encodeSystemFsGsBase encodes the FS/GS base accessors: F3-prefixed 0F AE
// under the /digit, the register in r/m; the Q spellings add REX.W.
func (e *enc) encodeSystemFsGsBase(mnem string, digit int, ops []Operand) error {
	b, size := splitSize(mnem)
	if _, ok := fsGsBase[b]; !ok || (size != 4 && size != 8) {
		return fmt.Errorf("unsupported instruction %q", mnem)
	}
	if len(ops) != 1 {
		return fmt.Errorf("%s expects 1 register operand, got %d", mnem, len(ops))
	}
	dstReg, ok := ops[0].(Reg)
	if !ok || dstReg.isVec() {
		return fmt.Errorf("%s destination must be a general register", mnem)
	}
	i := newInstr(0, []byte{0x0F, 0xAE})
	i.prefix = 0xF3
	i.rexW = size == 8
	if err := setRMDigit(i, digit, dstReg, 8); err != nil {
		return err
	}
	return e.emit(i)
}

// encodeSystemRm encodes a 0F 00/01 r/m access: the operand is a register or
// memory; the sized members carry the width prefix ladder the suffix fixes.
func (e *enc) encodeSystemRm(mnem string, m sysRmEntry, ops []Operand) error {
	size := 0
	if m.sized {
		_, size = splitSize(mnem)
		if size == 0 {
			return fmt.Errorf("unsupported instruction %q", mnem)
		}
	}
	if len(ops) != 1 {
		return fmt.Errorf("%s expects 1 operand, got %d", mnem, len(ops))
	}
	i := newInstr(size, []byte{0x0F, m.group})
	if err := setRMDigit(i, m.digit, ops[0], size); err != nil {
		return err
	}
	return e.emit(i)
}

// encodeSystemSelectorRead encodes LAR/LSL: the destination register loads
// from an r/m selector, width prefixes per the suffix.
func (e *enc) encodeSystemSelectorRead(mnem string, op byte, ops []Operand) error {
	b, size := splitSize(mnem)
	if _, ok := selectorRead[b]; !ok || size == 0 {
		return fmt.Errorf("unsupported instruction %q", mnem)
	}
	if len(ops) != 2 {
		return fmt.Errorf("%s expects 2 operands, got %d", mnem, len(ops))
	}
	dstReg, ok := ops[1].(Reg)
	if !ok || dstReg.isVec() {
		return fmt.Errorf("%s destination must be a general register", mnem)
	}
	i := newInstr(size, []byte{0x0F, op})
	if err := setRM(i, dstReg, ops[0], size); err != nil {
		return err
	}
	return e.emit(i)
}

// encodeSystemFarLoad encodes LFS/LGS/LSS: the destination register loads a
// far pointer from memory, width prefixes per the suffix.
func (e *enc) encodeSystemFarLoad(mnem string, op byte, ops []Operand) error {
	b, size := splitSize(mnem)
	if _, ok := farSegLoad[b]; !ok || size == 0 {
		return fmt.Errorf("unsupported instruction %q", mnem)
	}
	if len(ops) != 2 {
		return fmt.Errorf("%s expects 2 operands, got %d", mnem, len(ops))
	}
	if !isX86Mem(ops[0]) {
		return fmt.Errorf("%s requires a memory source", mnem)
	}
	dstReg, ok := ops[1].(Reg)
	if !ok || dstReg.isVec() {
		return fmt.Errorf("%s destination must be a general register", mnem)
	}
	i := newInstr(size, []byte{0x0F, op})
	if err := setRM(i, dstReg, ops[0], size); err != nil {
		return err
	}
	return e.emit(i)
}