gasm-sdk/lint/liveness.go

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

package lint

import (
	"slices"
	"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
		}
	}
	// A trailing label whose block was never flushed (no instruction follows
	// it) still holds the index the next block would have taken, which is one
	// past the end.  Drop those entries so a branch to such a label wires no
	// edge instead of indexing past the live sets in the dataflow.
	for name, v := range labelToBlock {
		if v >= len(l.blocks) {
			delete(labelToBlock, name)
		}
	}
	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.  The
// last bare-symbol operand is taken because the Plan 9 branch target sits in
// the destination position, after any operand registers: CBZ R0, done or
// BEQ R1, R2, done.  Scanning forward would pick R0 up as the target.
func branchTarget(in *ast.Instr) (string, bool) {
	for _, op := range slices.Backward(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)

	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: in Plan 9
// notation the destination is the last operand on every architecture Go
// supports (amd64, arm64, riscv64 and loong64 alike).
func dstIndex(in *ast.Instr) int {
	return len(in.Operands) - 1
}

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

// goroutineAlias maps the assembler's architectural alias for the goroutine
// register to its numeric name, verified against go tool asm of the Go 1.27
// toolchain: `MOVQ AX, g` encodes the same bytes as R14 on amd64, and on
// arm64, riscv64 and loong64 the alias is the ONLY spelling the toolchain
// accepts (a numeric R28, X27 or R22 operand is rejected), so a kernel can
// clobber the goroutine register through `g` alone.  The name is
// case-sensitive: only lowercase `g` assembles.
var goroutineAlias = map[arch.Arch]string{
	arch.AMD64:   "R14",
	arch.ARM64:   "R28",
	arch.RISCV:   "X27",
	arch.LOONG64: "R22",
}

// 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)
	}
	// The goroutine alias is not an arch-table register; resolve it so a
	// clobber written through `g` is audited like the numeric register.
	if n, ok := goroutineAlias[a]; ok && name == "g" {
		return n
	}
	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
}

// goFixedGPRs returns the general-purpose registers the Go ABI designates as
// fixed across calls, the ones hand-written assembly must not permanently
// clobber.  This follows cmd/compile/abi-internal.md, not the platform ABI:
// Go's stack-based ABI0 (which hand-written assembly uses) has no System V
// style callee-saved registers, so clobbering the argument and scratch
// registers (amd64 BX, R12, R13, R15, …) is legal.
//
// Two groups are returned.  always holds registers whose loss is never safe.
// runtime holds registers that survive an ABI0 leaf only because the
// transition machinery restores them (on amd64 the g pointer is reloaded
// from TLS): clobbering them is safe exactly in NOSPLIT functions that make
// no calls, which is how the runtime's own assembly uses them.
func goFixedGPRs(a arch.Arch) (always, runtime map[string]bool) {
	switch a {
	case arch.AMD64:
		// BP maintains the frame chain; R14 holds the current goroutine.
		// R15 is scratch except in dynamically linked binaries, so it is not
		// flagged.
		return gprSet("BP"), gprSet("R14")
	case arch.ARM64:
		// R18 is reserved for the OS on some platforms, R28 holds the current
		// goroutine, R29 is the frame pointer.
		return gprSet("R18", "R28", "R29"), nil
	case arch.RISCV:
		// X27 holds the current goroutine.
		return gprSet("X27"), nil
	case arch.LOONG64:
		// R22 holds the current goroutine.
		return gprSet("R22"), nil
	}
	return nil, nil
}

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

// clobberedGoFixed returns the Go-ABI-fixed registers a function writes
// without also saving and restoring them.  The first result lists registers
// whose loss is never safe; the second lists the goroutine-pointer class,
// whose loss is reported only when reachesRuntime is true (a non-NOSPLIT
// function, or one that makes calls, the ABI0 transition machinery restores
// the g pointer only on such paths).
func clobberedGoFixed(l *liveness, a arch.Arch, reachesRuntime bool) (always, runtime []string) {
	alwaysSet, runtimeSet := goFixedGPRs(a)
	if len(alwaysSet) == 0 && len(runtimeSet) == 0 {
		return nil, 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
			}
		}
	}
	clobbered := func(set map[string]bool) []string {
		var out []string
		for r := range set {
			if def[r] && !(saved[r] && restored[r]) {
				out = append(out, r)
			}
		}
		slices.Sort(out)
		return out
	}
	always = clobbered(alwaysSet)
	if reachesRuntime {
		runtime = clobbered(runtimeSet)
	}
	return always, runtime
}