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

package lint

import (
	"fmt"
	"sort"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-devkit/arch"
	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
)

// This file implements register liveness by dataflow over a function's
// control-flow graph, and the checks built on it.  The def/use model is
// deliberately conservative: where an instruction's effect is uncertain it is
// treated as both a use and a def of its register operands, which can only
// suppress a finding, never invent one.

// regEffect is the register-level effect of one instruction.
type regEffect struct {
	def     []string // registers written (killed)
	use     []string // registers read
	saveGPR []string // callee-saved-style: register written to the stack
	restGPR []string // register restored from the stack
}

// analyzeLiveness builds the control-flow graph of a function and computes
// live-in/live-out register sets by iterative backward dataflow.
type liveness struct {
	blocks []*block
	liveIn []map[string]bool
}

type block struct {
	label  string // label that begins this block, if any
	instrs []*ast.Instr
	succ   []int // successor block indices
}

func analyzeLiveness(t *ast.Text, a arch.Arch) *liveness {
	l := &liveness{}
	l.buildCFG(t)
	l.dataflow(a)
	return l
}

// buildCFG splits the function body into basic blocks and wires up successors.
func (l *liveness) buildCFG(t *ast.Text) {
	labelToBlock := map[string]int{}
	var cur *block
	flush := func() {
		if cur != nil && len(cur.instrs) > 0 {
			l.blocks = append(l.blocks, cur)
		}
		cur = nil
	}
	startBlock := func(lbl string) {
		flush()
		cur = &block{label: lbl}
	}

	startBlock("")
	for _, s := range t.Body {
		switch st := s.(type) {
		case *ast.Label:
			// A label begins a new block and is a jump target.
			startBlock(st.Name.Text)
			labelToBlock[st.Name.Text] = len(l.blocks) // index once flushed
		case *ast.Instr:
			if cur == nil {
				startBlock("")
			}
			cur.instrs = append(cur.instrs, st)
			if terminates(st) || isConditionalBranch(st) {
				startBlock("")
			}
		}
	}
	flush()

	// Fix up label->block indices (labels were recorded before the following
	// block was appended) and build successor edges.
	for i, b := range l.blocks {
		if b.label != "" {
			labelToBlock[b.label] = i
		}
	}
	for i, b := range l.blocks {
		if len(b.instrs) == 0 {
			if i+1 < len(l.blocks) {
				b.succ = append(b.succ, i+1)
			}
			continue
		}
		last := b.instrs[len(b.instrs)-1]
		mnem := strings.ToUpper(last.Mnemonic.Text)
		switch {
		case mnem == "RET" || mnem == "UNDEF":
			// No successors.
		case isConditionalBranch(last):
			if tgt, ok := branchTarget(last); ok {
				if j, found := labelToBlock[tgt]; found {
					b.succ = append(b.succ, j)
				}
			}
			if i+1 < len(l.blocks) {
				b.succ = append(b.succ, i+1) // fall-through
			}
		case isUnconditionalBranchAny(mnem):
			if tgt, ok := branchTarget(last); ok {
				if j, found := labelToBlock[tgt]; found {
					b.succ = append(b.succ, j)
				}
			}
		default:
			if i+1 < len(l.blocks) {
				b.succ = append(b.succ, i+1)
			}
		}
	}
}

// dataflow runs the standard backward liveness iteration to a fixed point.
func (l *liveness) dataflow(a arch.Arch) {
	n := len(l.blocks)
	l.liveIn = make([]map[string]bool, n)
	liveOut := make([]map[string]bool, n)
	use := make([]map[string]bool, n)
	def := make([]map[string]bool, n)
	for i, b := range l.blocks {
		use[i], def[i] = blockUseDef(b, a)
		l.liveIn[i] = map[string]bool{}
		liveOut[i] = map[string]bool{}
	}
	for changed := true; changed; {
		changed = false
		for i := n - 1; i >= 0; i-- {
			out := map[string]bool{}
			for _, s := range l.blocks[i].succ {
				for r := range l.liveIn[s] {
					out[r] = true
				}
			}
			if !sameSet(out, liveOut[i]) {
				liveOut[i] = out
				changed = true
			}
			// in = use ∪ (out − def)
			in := map[string]bool{}
			for r := range use[i] {
				in[r] = true
			}
			for r := range out {
				if !def[i][r] {
					in[r] = true
				}
			}
			if !sameSet(in, l.liveIn[i]) {
				l.liveIn[i] = in
				changed = true
			}
		}
	}
}

// blockUseDef computes the registers used before definition (use) and the
// registers defined (def) within a basic block.
func blockUseDef(b *block, a arch.Arch) (use, def map[string]bool) {
	use = map[string]bool{}
	def = map[string]bool{}
	for _, in := range b.instrs {
		eff := instrEffect(in, a)
		for _, r := range eff.use {
			if !def[r] {
				use[r] = true
			}
		}
		for _, r := range eff.def {
			def[r] = true
		}
	}
	return use, def
}

// terminates reports whether an instruction ends basic-block flow unconditionally.
func terminates(in *ast.Instr) bool {
	m := strings.ToUpper(in.Mnemonic.Text)
	return m == "RET" || m == "UNDEF" || isUnconditionalBranchAny(m)
}

func isConditionalBranch(in *ast.Instr) bool {
	m := strings.ToUpper(in.Mnemonic.Text)
	// Conditional jumps/branches, but not the unconditional ones.
	if isUnconditionalBranchAny(m) || m == "RET" || m == "UNDEF" || m == "CALL" {
		return false
	}
	return strings.HasPrefix(m, "J") || strings.HasPrefix(m, "B") ||
		strings.HasPrefix(m, "CBZ") || strings.HasPrefix(m, "CBNZ") ||
		strings.HasPrefix(m, "TBZ") || strings.HasPrefix(m, "TBNZ") ||
		strings.HasPrefix(m, "BEQ") || strings.HasPrefix(m, "BNE")
}

// isUnconditionalBranchAny is an arch-agnostic unconditional-branch test.
func isUnconditionalBranchAny(m string) bool {
	switch m {
	case "JMP", "J", "JR", "B", "BR", "JIRL":
		return true
	}
	return false
}

// branchTarget returns the local-label target of a branch, if it is one.
func branchTarget(in *ast.Instr) (string, bool) {
	for _, op := range in.Operands {
		if op.Kind == ast.OpAddr && op.Addr.Sym != nil && op.Addr.Sym.Pseudo == "" &&
			op.Addr.Base == "" && op.Addr.Sym.Name != "" {
			return op.Addr.Sym.Name, true
		}
	}
	return "", false
}

// instrEffect returns the register-level effect of one instruction.
func instrEffect(in *ast.Instr, a arch.Arch) regEffect {
	var eff regEffect
	mnem := strings.ToUpper(in.Mnemonic.Text)

	// PUSH/POP move a register to/from the stack.
	if strings.HasPrefix(mnem, "PUSH") {
		for _, op := range in.Operands {
			if r := gprName(op, a); r != "" {
				eff.use = append(eff.use, r)
				eff.saveGPR = append(eff.saveGPR, r)
			}
		}
		return eff
	}
	if strings.HasPrefix(mnem, "POP") {
		for _, op := range in.Operands {
			if r := gprName(op, a); r != "" {
				eff.def = append(eff.def, r)
				eff.restGPR = append(eff.restGPR, r)
			}
		}
		return eff
	}

	compare := isCompare(mnem)
	dstIdx := dstIndex(in, a)

	for i, op := range in.Operands {
		r := gprName(op, a)
		if r != "" {
			if i == dstIdx && !compare {
				eff.def = append(eff.def, r)
				// Arithmetic also reads its destination.
				eff.use = append(eff.use, r)
			} else {
				eff.use = append(eff.use, r)
			}
		}
		// Detect saves/restores through the stack frame.
		if isStackAddr(op) {
			// The other operand (the register) is being saved or restored.
			for j, other := range in.Operands {
				if j == i {
					continue
				}
				if rr := gprName(other, a); rr != "" {
					if j == dstIdx && !compare {
						eff.restGPR = append(eff.restGPR, rr) // reg loaded from stack
					} else {
						eff.saveGPR = append(eff.saveGPR, rr) // reg stored to stack
					}
				}
			}
		}
	}
	return eff
}

// dstIndex returns the operand index of the destination register: last for the
// Plan 9 (amd64) spelling, first for arm64/riscv64/loong64.
func dstIndex(in *ast.Instr, a arch.Arch) int {
	if a == arch.AMD64 {
		return len(in.Operands) - 1
	}
	return 0
}

// isCompare reports whether the mnemonic only reads its operands (setting flags).
func isCompare(m string) bool {
	return strings.HasPrefix(m, "CMP") || strings.HasPrefix(m, "TEST") ||
		strings.HasPrefix(m, "CMN") || strings.HasPrefix(m, "TST") ||
		m == "FCMP" || m == "FCMPE"
}

// gprName returns the canonical general-purpose register name of an operand, or
// "" if the operand is not a bare GPR reference.
func gprName(op *ast.Operand, a arch.Arch) string {
	if op == nil || op.Kind != ast.OpAddr || op.Addr.Sym == nil {
		return ""
	}
	if op.Addr.Base != "" || op.Addr.Sym.Pseudo != "" || op.Addr.Sym.Name == "" {
		return ""
	}
	name := op.Addr.Sym.Name
	if r, ok := arch.ForArch(a).Register(name); ok && (r.Class == arch.GPR || r.Class == arch.GPRSub) {
		return canonicalGPR(name)
	}
	return ""
}

// canonicalGPR maps a sized sub-register to its base GPR (amd64 only).
func canonicalGPR(name string) string {
	upper := strings.ToUpper(name)
	// Named 8/16/32-bit forms of the classic registers.
	switch upper {
	case "AL", "AH", "AX":
		return "AX"
	case "BL", "BH", "BX":
		return "BX"
	case "CL", "CH", "CX":
		return "CX"
	case "DL", "DH", "DX":
		return "DX"
	case "SIL":
		return "SI"
	case "DIL":
		return "DI"
	case "BPL":
		return "BP"
	case "SPL":
		return "SP"
	}
	// Numbered sub-registers R8B/R8W/R8D → R8.
	if len(upper) >= 3 && upper[0] == 'R' {
		switch upper[len(upper)-1] {
		case 'B', 'W', 'D':
			return upper[:len(upper)-1]
		}
	}
	return upper
}

// isStackAddr reports whether an operand addresses the stack frame
// (base SP, or an FP/SP-relative symbol).
func isStackAddr(op *ast.Operand) bool {
	if op == nil || op.Kind != ast.OpAddr {
		return false
	}
	if op.Addr.Base == "SP" {
		return true
	}
	if op.Addr.Sym != nil && (op.Addr.Sym.Pseudo == "SP" || op.Addr.Sym.Pseudo == "FP") {
		return true
	}
	return false
}

func sameSet(a, b map[string]bool) bool {
	if len(a) != len(b) {
		return false
	}
	for k := range a {
		if !b[k] {
			return false
		}
	}
	return true
}

// calleeSavedGPRs returns the general-purpose registers an assembly function
// must preserve for its caller, using the register names the assembler accepts
// for each architecture.
func calleeSavedGPRs(a arch.Arch) map[string]bool {
	switch a {
	case arch.AMD64:
		return gprSet("BX", "BP", "R12", "R13", "R14", "R15")
	case arch.ARM64:
		names := []string{"R29", "R30"} // FP, LR
		for i := 19; i <= 28; i++ {
			names = append(names, fmt.Sprintf("R%d", i))
		}
		return gprSet(names...)
	case arch.RISCV:
		// RA (X1) and the S registers (X8, X9, X18–X27) are callee-saved.
		names := []string{"X1", "RA", "X8", "X9", "S0", "S1", "FP"}
		for i := 18; i <= 27; i++ {
			names = append(names, fmt.Sprintf("X%d", i))
		}
		for i := 2; i <= 11; i++ {
			names = append(names, fmt.Sprintf("S%d", i))
		}
		return gprSet(names...)
	case arch.LOONG64:
		// RA (R1), FP (R22) and S0–S8 (R23–R31) are callee-saved.
		names := []string{"R1", "RA", "R22", "FP"}
		for i := 23; i <= 31; i++ {
			names = append(names, fmt.Sprintf("R%d", i))
		}
		for i := 0; i <= 8; i++ {
			names = append(names, fmt.Sprintf("S%d", i))
		}
		return gprSet(names...)
	}
	return nil
}

func gprSet(names ...string) map[string]bool {
	m := make(map[string]bool, len(names))
	for _, n := range names {
		m[n] = true
	}
	return m
}

// clobberedCalleeSaved returns the callee-saved registers a function writes
// without also saving and restoring them — i.e. registers whose caller-owned
// value is lost across the call.  It walks the blocks of the liveness analysis
// (so the control-flow graph is what supplies the instruction set) and
// aggregates each instruction's register effects.
func clobberedCalleeSaved(l *liveness, a arch.Arch) []string {
	callee := calleeSavedGPRs(a)
	if len(callee) == 0 {
		return nil
	}
	def := map[string]bool{}
	saved := map[string]bool{}
	restored := map[string]bool{}
	for _, b := range l.blocks {
		for _, in := range b.instrs {
			eff := instrEffect(in, a)
			for _, r := range eff.def {
				def[r] = true
			}
			for _, r := range eff.saveGPR {
				saved[r] = true
			}
			for _, r := range eff.restGPR {
				restored[r] = true
			}
		}
	}
	var out []string
	for r := range callee {
		if def[r] && !(saved[r] && restored[r]) {
			out = append(out, r)
		}
	}
	sort.Strings(out)
	return out
}