//! Derive a function's observable **ABI shape** — its SysV-AMD64 argument footprint — straight from
//! the machine code, so a C++ *prototype* change (an edit to the argument list) becomes an offline
//! diff instead of a silently-stale loader hook.
//!
//! The byte-signature already handles the function's *body* drifting across recompiles: a changed
//! prologue is re-derived and re-validated. What a byte-signature CANNOT see is the argument list
//! changing while the body's opening bytes stay recognisable — the sig still resolves, the offset
//! still points at real code, `validate-live` still passes, yet a loader that calls the function with
//! the OLD prototype now passes the wrong registers. That failure is invisible to every existing gate.
//!
//! This module recovers the one thing that pins the prototype: which argument registers the function
//! reads as inputs. On SysV-AMD64 the first six integer/pointer arguments arrive in RDI, RSI, RDX,
//! RCX, R8, R9 and the first eight floating arguments in XMM0..XMM7, each assigned left-to-right. A
//! register is an *input* exactly when it is live-in at the entry — read on some path before being
//! written. We compute that with a bounded backward liveness over the 14 argument registers, then
//! read off the contiguous integer- and float-argument counts. The result is recompilation-invariant
//! (a rebuild doesn't change which arguments a function takes) and moves precisely when the prototype
//! does — so comparing it across builds flags exactly the prototype changes the byte-sig misses.
//!
//! Known limits (all bias toward UNDER-counting = a missed flag, never a false one): a pure forwarding
//! thunk (`jmp Helper`) reads no arg register of its own, so it shapes as `(0,0)`; an argument used
//! only inside a jump-table (indirect-branch) case isn't followed, so it can be missed. Both stay
//! stable across builds (a thunk stays a thunk), so they don't manufacture false transitions — the
//! diff's `int==0` low-confidence bucket also absorbs the thunk case. `int_args` is the OBSERVABLE
//! footprint = a lower bound on the declared prototype (a constant-returner reads nothing → `int=0`);
//! that too is stable per function, so the cross-build diff still works.

use crate::elf::CodeImage;
use iced_x86::{
    Decoder, DecoderOptions, FlowControl, Instruction, InstructionInfoFactory, Mnemonic, OpAccess,
    OpKind, Register,
};
use std::collections::HashMap;

/// Argument-register slots, in ABI order: 0..6 = RDI,RSI,RDX,RCX,R8,R9; 6..14 = XMM0..XMM7. A `u16`
/// bitmask over these 14 slots is a function's live-in argument set.
const N_INT: usize = 6;
const N_XMM: usize = 8;

/// A function's recovered ABI shape: how many integer/pointer and floating arguments it reads, plus
/// whether it also loads arguments off the stack (a 7th+ integer / 9th+ float argument, or a large
/// by-value struct). The `(int_args, float_args)` pair is the stable cross-build key — `stack_args`
/// is a best-effort extra signal, reported but not used to key the diff.
#[derive(Clone, Copy, PartialEq, Eq, Default, Debug)]
pub struct AbiShape {
    pub int_args: u8,
    pub float_args: u8,
    pub stack_args: bool,
    /// The register class of the return value — a prototype dimension the argument footprint can't see.
    pub ret_class: RetClass,
}

/// How a function returns its result, recovered from the return paths. Complements the argument
/// footprint: a change here (int↔float↔by-value) is a prototype change the arg counts alone miss, and
/// `ByValue` marks the RVO/sret functions that are UNSAFE to blind-call — the caller must pass an
/// output-buffer pointer in RDI, so calling with the object there makes the function WRITE into it
/// (the `CSwapTeams::GetDisplayString` sret trap). Best-effort, with an explicit
/// `Unknown` when the return path doesn't decode — so it only ever adds a signal, never a false one.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Default, Debug)]
pub enum RetClass {
    /// No decodable return path (a forwarding thunk / tail call / undecoded) — no signal.
    #[default]
    Unknown,
    /// No result register written before returning (best-effort void).
    Void,
    /// Scalar / pointer result in RAX (the common case).
    Int,
    /// Floating result in XMM0.
    Float,
    /// Large by-value aggregate (sret): the function writes the result through its incoming RDI output
    /// pointer and returns that pointer. UNSAFE to call with an object in RDI.
    ByValue,
}

impl RetClass {
    pub fn describe(self) -> &'static str {
        match self {
            RetClass::Unknown => "ret=?",
            RetClass::Void => "ret=void",
            RetClass::Int => "ret=int",
            RetClass::Float => "ret=float",
            RetClass::ByValue => "ret=byval",
        }
    }

    /// Inverse of [`Self::describe`] — parse the token back. `None` for an unrecognised string.
    pub fn from_describe(s: &str) -> Option<RetClass> {
        Some(match s {
            "ret=?" => RetClass::Unknown,
            "ret=void" => RetClass::Void,
            "ret=int" => RetClass::Int,
            "ret=float" => RetClass::Float,
            "ret=byval" => RetClass::ByValue,
            _ => return None,
        })
    }
}

// Order + on-disk form are pinned to `describe()`: `AbiSig` stores a `RetClass`, is serialized into the
// model, and `mode_abi` sorts `AbiSig` with a tie-break "to the larger by Ord". Both must serialize and
// order by the describe() token EXACTLY, or the consensus bytes move — so serialization is the
// describe() token and Ord is that token's lexical order, not the enum's declaration order.
impl Ord for RetClass {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.describe().cmp(other.describe())
    }
}
impl PartialOrd for RetClass {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}
impl serde::Serialize for RetClass {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        s.serialize_str(self.describe())
    }
}
impl<'de> serde::Deserialize<'de> for RetClass {
    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        let s = String::deserialize(d)?;
        RetClass::from_describe(&s)
            .ok_or_else(|| serde::de::Error::custom(format!("bad ret class {s:?}")))
    }
}

impl AbiShape {
    /// The recompilation-invariant identity used to compare shapes across builds. Deliberately omits
    /// `stack_args` (a frame-layout heuristic that a rebuild could flip) so a diff never false-flags
    /// on it. Return class is compared separately (`ret_class`) — a return-type change is reported
    /// distinctly from an argument-list change.
    pub fn key(&self) -> (u8, u8) {
        (self.int_args, self.float_args)
    }

    /// Safe to blind-call with only a `this` pointer: at most one integer arg (the implicit `this`), no
    /// float args, no stack args, AND not an sret return. The live oracle invokes ONLY such methods —
    /// anything else needs arguments it doesn't have, so calling it would pass garbage. Both live callers
    /// gate on exactly this shape, so it lives here (one safety-critical definition) rather than inline.
    ///
    /// The `ByValue` guard is the real teeth: an sret/RVO function receives a hidden output-buffer pointer in
    /// its FIRST integer register (RDI) and returns it, so a `this`-less sret (`int_args == 1` = just that
    /// pointer) would otherwise look "this-only" — and blind-calling it with the live object as RDI makes the
    /// callee WRITE its return value THROUGH the object = memory corruption. That is precisely what
    /// `RetClass::ByValue` marks, so consult it here rather than trust the naming convention alone.
    pub fn is_this_only(&self) -> bool {
        self.int_args <= 1
            && self.float_args == 0
            && !self.stack_args
            && self.ret_class != RetClass::ByValue
    }
}

/// The argument-register slot a register belongs to, or `None` if it isn't one. Sub-registers fold to
/// their slot (EDI/DI/DIL -> RDI's slot); YMM/ZMM 0..7 fold to the matching XMM slot.
fn arg_slot(r: Register) -> Option<usize> {
    match r.full_register() {
        Register::RDI => return Some(0),
        Register::RSI => return Some(1),
        Register::RDX => return Some(2),
        Register::RCX => return Some(3),
        Register::R8 => return Some(4),
        Register::R9 => return Some(5),
        _ => {}
    }
    // Vector argument registers: XMM/YMM/ZMM 0..7 all map to the same 8 float slots.
    let base = if r.is_xmm() {
        Register::XMM0
    } else if r.is_ymm() {
        Register::YMM0
    } else if r.is_zmm() {
        Register::ZMM0
    } else {
        return None;
    };
    let i = (r as u32).wrapping_sub(base as u32) as usize;
    (i < N_XMM).then_some(N_INT + i)
}

/// The destination register of an instruction that decodes as read-writing that register but whose
/// RESULT does not depend on the register's prior value — so it is a def, not a genuine input read.
/// Three cases:
///  (1) same-register zeroing idioms (`xor r,r`, `pxor x,x`, `vxorps x,x,x`);
///  (2) same-register all-ones idioms (`pcmpeqd x,x`);
///  (3) legacy scalar-SSE writes (`cvtsi2sd`, `sqrtsd`, `movsd`-reg, …) whose low lanes are FULLY
///      written and whose read-write access only models the preserved UPPER lanes — a decode artifact.
/// Case (3) is essential: without it a float argument's count would move with body codegen (a bare
/// `cvtsi2sd xmm0,rax` decodes as reading xmm0; a dependency-broken `xorps xmm0,xmm0; cvtsi2sd …` does
/// not), which is exactly the recompilation drift the shape must be immune to. Genuine RMW arithmetic
/// (`addss`/`mulss`/…) is NOT listed — those really do read their destination, so their read is kept.
/// VEX forms take their merge lanes from an explicit source operand, so their destination is a pure
/// Write and never decodes as a false read.
fn false_read_dst(insn: &Instruction) -> Option<Register> {
    let self_idiom = matches!(
        insn.mnemonic(),
        Mnemonic::Xor
            | Mnemonic::Sub
            | Mnemonic::Sbb
            | Mnemonic::Pxor
            | Mnemonic::Xorps
            | Mnemonic::Xorpd
            | Mnemonic::Vpxor
            | Mnemonic::Vxorps
            | Mnemonic::Vxorpd
            | Mnemonic::Pcmpeqb
            | Mnemonic::Pcmpeqw
            | Mnemonic::Pcmpeqd
            | Mnemonic::Pcmpeqq
    );
    if self_idiom {
        match insn.op_count() {
            // legacy 2-operand: `xor r, r` / `pcmpeqd x, x`
            2 if insn.op0_kind() == OpKind::Register
                && insn.op1_kind() == OpKind::Register
                && insn.op0_register() == insn.op1_register() =>
            {
                return Some(insn.op0_register());
            }
            // VEX 3-operand: `vpxor dst, src, src`
            3 if insn.op1_kind() == OpKind::Register
                && insn.op2_kind() == OpKind::Register
                && insn.op1_register() == insn.op2_register() =>
            {
                return Some(insn.op0_register());
            }
            _ => {}
        }
    }
    // Legacy scalar-SSE merge-only writes: low lanes fully written, old value doesn't feed the result.
    let merge_only = matches!(
        insn.mnemonic(),
        Mnemonic::Movss
            | Mnemonic::Movsd
            | Mnemonic::Cvtsi2ss
            | Mnemonic::Cvtsi2sd
            | Mnemonic::Cvtss2sd
            | Mnemonic::Cvtsd2ss
            | Mnemonic::Sqrtss
            | Mnemonic::Sqrtsd
            | Mnemonic::Roundss
            | Mnemonic::Roundsd
            | Mnemonic::Rcpss
            | Mnemonic::Rsqrtss
    );
    (merge_only && insn.op0_kind() == OpKind::Register).then(|| insn.op0_register())
}

/// One instruction's effect on the argument registers: `(use_mask, def_mask, reads_a_stack_arg)`.
/// `use` = arg slots read (Read/CondRead/ReadWrite) minus false-read destinations; `def` = arg slots
/// fully written (>=32-bit Write/ReadWrite, plus false-read destinations, which kill upward liveness).
/// Shared by `decode_region` and the tests so a test can never mirror-drift from the real logic.
fn insn_effect(factory: &mut InstructionInfoFactory, insn: &Instruction) -> (u16, u16, bool) {
    let (mut use_m, mut def_m) = (0u16, 0u16);
    let mut stack = false;
    let info = factory.info(insn);
    for ur in info.used_registers() {
        let Some(slot) = arg_slot(ur.register()) else {
            continue;
        };
        if matches!(
            ur.access(),
            OpAccess::Read | OpAccess::CondRead | OpAccess::ReadWrite | OpAccess::ReadCondWrite
        ) {
            use_m |= 1 << slot;
        }
        // A def kills upward liveness only for a full-width write (>=32-bit writes clear the upper
        // bits; an 8/16-bit partial write leaves the register partly live, so it doesn't kill).
        if matches!(ur.access(), OpAccess::Write | OpAccess::ReadWrite) && ur.register().size() >= 4
        {
            def_m |= 1 << slot;
        }
    }
    // Stack argument: a read of `[rbp + disp]` above the saved frame (return addr at +8, first stack
    // arg at +16). Best-effort — frame-pointer-omitted stack args aren't caught.
    for um in info.used_memory() {
        if um.base() == Register::RBP
            && matches!(
                um.access(),
                OpAccess::Read | OpAccess::CondRead | OpAccess::ReadWrite
            )
            && (16..0x1000).contains(&(um.displacement() as i64))
        {
            stack = true;
        }
    }
    // A false-read destination (zeroing/all-ones idiom, or a scalar-SSE merge-only write) is a def,
    // not a use — even though it decodes as read-write of that register.
    if let Some(dst) = false_read_dst(insn)
        && let Some(slot) = arg_slot(dst)
    {
        use_m &= !(1 << slot);
        def_m |= 1 << slot;
    }
    (use_m, def_m, stack)
}

/// How an instruction writes RAX — the return-value register. `FromRdi` is the sret / return-this tell
/// (`mov rax, rdi` / `lea rax, [rdi]`): RAX takes the incoming output pointer.
#[derive(Clone, Copy, PartialEq, Eq)]
enum RaxWrite {
    None,
    FromRdi,
    Other,
}

/// One instruction's effect on the RETURN registers: does it write RAX (and from RDI?), does it write
/// XMM0, and does it store through RDI (the sret output-write). Used to classify the return value.
fn result_effect(
    factory: &mut InstructionInfoFactory,
    insn: &Instruction,
) -> (RaxWrite, bool, bool) {
    // The sret / return-this tell: a full 64-bit RAX <- RDI copy.
    let rax_from_rdi = match insn.mnemonic() {
        Mnemonic::Mov => {
            insn.op0_kind() == OpKind::Register
                && insn.op0_register() == Register::RAX
                && insn.op1_kind() == OpKind::Register
                && insn.op1_register() == Register::RDI
        }
        Mnemonic::Lea => {
            insn.op0_kind() == OpKind::Register
                && insn.op0_register() == Register::RAX
                && insn.memory_base() == Register::RDI
                && insn.memory_index() == Register::None
                && insn.memory_displacement64() == 0
        }
        _ => false,
    };
    let info = factory.info(insn);
    let (mut rax_written, mut xmm0_written) = (false, false);
    for ur in info.used_registers() {
        if !matches!(ur.access(), OpAccess::Write | OpAccess::ReadWrite) {
            continue;
        }
        // RAX written as a full (>=32-bit) value = a scalar/pointer return candidate.
        if ur.register().full_register() == Register::RAX && ur.register().size() >= 4 {
            rax_written = true;
        }
        // XMM0/YMM0/ZMM0 (the float-return slot) written = a float return candidate.
        if arg_slot(ur.register()) == Some(N_INT) {
            xmm0_written = true;
        }
    }
    // A store through RDI as base = the function uses RDI as an output buffer (the sret write).
    let stores_rdi = info.used_memory().iter().any(|um| {
        um.base() == Register::RDI && matches!(um.access(), OpAccess::Write | OpAccess::ReadWrite)
    });
    let rax = if rax_from_rdi {
        RaxWrite::FromRdi
    } else if rax_written {
        RaxWrite::Other
    } else {
        RaxWrite::None
    };
    (rax, xmm0_written, stores_rdi)
}

/// Classify the return value from the decoded region: for each `ret`, find the nearest preceding result
/// write and read off its register class, then take a consensus. `ByValue` (sret) requires BOTH the
/// return-the-RDI-pointer pattern AND a store through RDI — so a plain `return this` (returns RDI but
/// doesn't write through it) reads as an `Int` pointer return, not a by-value trap.
fn derive_ret_class(list: &[Insn], stores_rdi: bool) -> RetClass {
    let mut votes: Vec<RetClass> = Vec::new();
    for (r, insn) in list.iter().enumerate() {
        if !insn.is_ret {
            continue;
        }
        let mut cls = RetClass::Void; // no result write found before the return
        for k in (r.saturating_sub(64)..r).rev() {
            match list[k].rax {
                RaxWrite::FromRdi => {
                    cls = if stores_rdi {
                        RetClass::ByValue
                    } else {
                        RetClass::Int
                    };
                    break;
                }
                RaxWrite::Other => {
                    cls = RetClass::Int;
                    break;
                }
                RaxWrite::None if list[k].xmm0 => {
                    cls = RetClass::Float;
                    break;
                }
                RaxWrite::None => {}
            }
        }
        votes.push(cls);
    }
    if votes.is_empty() {
        return RetClass::Unknown; // no return path (thunk / tail call)
    }
    if votes.contains(&RetClass::ByValue) {
        return RetClass::ByValue; // sret is definitive wherever it appears
    }
    match (
        votes.contains(&RetClass::Int),
        votes.contains(&RetClass::Float),
    ) {
        (true, true) => RetClass::Unknown, // paths disagree on the return register — ambiguous
        (true, false) => RetClass::Int,
        (false, true) => RetClass::Float,
        (false, false) => RetClass::Void,
    }
}

/// One decoded instruction's argument-register effect + return effect + in-function successors.
struct Insn {
    ip: u64,
    use_m: u16,     // arg slots read before this instruction can write them (a use)
    def_m: u16,     // arg slots fully written (kills upward liveness)
    rax: RaxWrite,  // how it writes RAX (the return register)
    xmm0: bool,     // whether it writes XMM0 (the float-return register)
    is_ret: bool,   // whether it returns
    succ: Vec<u64>, // successor instruction addresses inside the analysed region
}

/// Decode the function at `entry` into its bounded reachable instructions, following the same
/// conditional/unconditional control flow as the fingerprinter. Each instruction records which
/// argument registers it uses/defs (for liveness) and its in-region successors (for the CFG). Also
/// returns whether a stack argument was read anywhere in the region.
fn decode_region(img: &CodeImage, entry: u64) -> Option<(Vec<Insn>, bool, bool)> {
    const MAX_SPAN: usize = 96 * 1024;
    const MAX_INSNS: usize = 8000;
    let code = img.code_at(entry)?;
    let cap = code.len().min(MAX_SPAN);
    let in_span = |t: u64| t >= entry && ((t - entry) as usize) < cap;

    let mut factory = InstructionInfoFactory::new();
    let mut recs: HashMap<u64, Insn> = HashMap::new();
    let mut stack_args = false;
    let mut stores_rdi = false;
    let mut work = vec![entry];
    let mut insn = Instruction::default();
    while let Some(start) = work.pop() {
        if recs.contains_key(&start) || !in_span(start) {
            continue;
        }
        let off = (start - entry) as usize;
        let mut dec = Decoder::with_ip(64, &code[off..], start, DecoderOptions::NONE);
        if !dec.can_decode() {
            continue;
        }
        dec.decode_out(&mut insn);
        if insn.is_invalid() || insn.len() == 0 || recs.len() >= MAX_INSNS {
            continue;
        }

        let (use_m, def_m, stack_hit) = insn_effect(&mut factory, &insn);
        stack_args |= stack_hit;
        let (rax, xmm0, rdi_store) = result_effect(&mut factory, &insn);
        stores_rdi |= rdi_store;
        let is_ret = insn.flow_control() == FlowControl::Return;

        let next = start + insn.len() as u64;
        let mut succ = Vec::new();
        match insn.flow_control() {
            FlowControl::Return | FlowControl::IndirectBranch => {}
            FlowControl::UnconditionalBranch => {
                let t = insn.near_branch_target();
                if in_span(t) {
                    succ.push(t); // in-function jump; else it's a tail call (no in-region successor)
                }
            }
            FlowControl::ConditionalBranch => {
                succ.push(next);
                let t = insn.near_branch_target();
                if in_span(t) {
                    succ.push(t);
                }
            }
            _ => succ.push(next), // fall-through (incl. call/indirect-call: the call reads no arg regs)
        }
        for &s in &succ {
            if !recs.contains_key(&s) {
                work.push(s);
            }
        }
        recs.insert(
            start,
            Insn {
                ip: start,
                use_m,
                def_m,
                rax,
                xmm0,
                is_ret,
                succ,
            },
        );
    }
    if recs.is_empty() {
        return None;
    }
    let mut list: Vec<Insn> = recs.into_values().collect();
    list.sort_by_key(|i| i.ip);
    Some((list, stack_args, stores_rdi))
}

/// Recover the ABI shape of the function at `entry`, or `None` if it doesn't decode. Runs a bounded
/// backward liveness over the 14 argument registers to find the entry's live-in set, then reads off
/// the contiguous integer- and float-argument counts (a later argument register being live-in implies
/// the earlier ones are arguments too — the SysV assignment is left-to-right and gap-free).
pub fn abi_shape(img: &CodeImage, entry: u64) -> Option<AbiShape> {
    let (list, stack_args, stores_rdi) = decode_region(img, entry)?;
    let n = list.len();
    let idx: HashMap<u64, usize> = list.iter().enumerate().map(|(i, r)| (r.ip, i)).collect();
    let succ: Vec<Vec<usize>> = list
        .iter()
        .map(|r| r.succ.iter().filter_map(|s| idx.get(s).copied()).collect())
        .collect();

    // live_in[i] = use[i] | (live_out[i] & !def[i]); live_out[i] = union of successors' live_in.
    // Iterating in reverse index order converges in a couple of passes on a mostly-forward CFG; the
    // pass cap bounds the pathological (deeply nested loops) case. Deterministic regardless.
    let mut live_in = vec![0u16; n];
    for _ in 0..64 {
        let mut changed = false;
        for i in (0..n).rev() {
            let mut out = 0u16;
            for &s in &succ[i] {
                out |= live_in[s];
            }
            let v = list[i].use_m | (out & !list[i].def_m);
            if v != live_in[i] {
                live_in[i] = v;
                changed = true;
            }
        }
        if !changed {
            break;
        }
    }

    let entry_live = idx.get(&entry).map(|&i| live_in[i]).unwrap_or(live_in[0]);
    let int_args = (0..N_INT)
        .rev()
        .find(|&s| entry_live & (1 << s) != 0)
        .map_or(0, |s| s + 1) as u8;
    let float_args = (0..N_XMM)
        .rev()
        .find(|&s| entry_live & (1 << (N_INT + s)) != 0)
        .map_or(0, |s| s + 1) as u8;
    // `stack_args` only makes sense once the register slots are full; a stray rbp read below that is a
    // spill, not an argument.
    let stack_args = stack_args && (int_args as usize == N_INT || float_args as usize == N_XMM);
    Some(AbiShape {
        int_args,
        float_args,
        stack_args,
        ret_class: derive_ret_class(&list, stores_rdi),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    // Decode a tiny hand-assembled straight-line function and recover its shape through the REAL
    // per-instruction helper (`insn_effect`) + the real liveness formula — so a test can't pass while
    // the production path is wrong. (A single-successor chain; the fixpoint isn't exercised here.)
    fn shape_of(bytes: &[u8]) -> AbiShape {
        let entry = 0x1000u64;
        let mut factory = InstructionInfoFactory::new();
        let mut recs: Vec<Insn> = Vec::new();
        let mut stores_rdi = false;
        let mut dec = Decoder::with_ip(64, bytes, entry, DecoderOptions::NONE);
        let mut insn = Instruction::default();
        while dec.can_decode() {
            dec.decode_out(&mut insn);
            if insn.is_invalid() || insn.len() == 0 {
                break;
            }
            let (use_m, def_m, _) = insn_effect(&mut factory, &insn);
            let (rax, xmm0, rdi_store) = result_effect(&mut factory, &insn);
            stores_rdi |= rdi_store;
            let ip = insn.ip();
            let stop = insn.flow_control() == FlowControl::Return;
            let next = ip + insn.len() as u64;
            recs.push(Insn {
                ip,
                use_m,
                def_m,
                rax,
                xmm0,
                is_ret: stop,
                succ: if stop { vec![] } else { vec![next] },
            });
            if stop {
                break;
            }
        }
        // straight-line liveness (single-successor chain)
        let n = recs.len();
        let mut live = vec![0u16; n];
        for i in (0..n).rev() {
            let out = recs[i].succ.first().map_or(0, |_| live[i + 1]);
            live[i] = recs[i].use_m | (out & !recs[i].def_m);
        }
        let el = live[0];
        let int_args = (0..N_INT)
            .rev()
            .find(|&s| el & (1 << s) != 0)
            .map_or(0, |s| s + 1) as u8;
        let float_args = (0..N_XMM)
            .rev()
            .find(|&s| el & (1 << (N_INT + s)) != 0)
            .map_or(0, |s| s + 1) as u8;
        AbiShape {
            int_args,
            float_args,
            stack_args: false,
            ret_class: derive_ret_class(&recs, stores_rdi),
        }
    }

    #[test]
    fn mov_rax_rdi_is_one_int_arg() {
        // mov rax, rdi ; ret
        assert_eq!(shape_of(&[0x48, 0x89, 0xF8, 0xC3]).key(), (1, 0));
    }

    #[test]
    fn deref_this_is_one_int_arg() {
        // mov rax, [rdi+8] ; ret  (rdi read as a memory base = `this` pointer)
        assert_eq!(shape_of(&[0x48, 0x8B, 0x47, 0x08, 0xC3]).key(), (1, 0));
    }

    #[test]
    fn xor_eax_eax_is_zero_args() {
        // xor eax, eax ; ret  (zeroing idiom is a def, not a use; eax isn't an arg reg anyway)
        assert_eq!(shape_of(&[0x31, 0xC0, 0xC3]).key(), (0, 0));
    }

    #[test]
    fn reading_rdx_fills_earlier_int_args() {
        // mov rax, rdx ; ret  — rdx (slot 2) live-in => contiguity fills rdi, rsi => 3 int args
        assert_eq!(shape_of(&[0x48, 0x89, 0xD0, 0xC3]).key(), (3, 0));
    }

    #[test]
    fn addss_reads_two_float_args() {
        // addss xmm0, xmm1 ; ret  — xmm0 (read-write) + xmm1 (read) live-in => 2 float args
        assert_eq!(shape_of(&[0xF3, 0x0F, 0x58, 0xC1, 0xC3]).key(), (0, 2));
    }

    #[test]
    fn pxor_self_is_not_a_float_arg() {
        // pxor xmm0, xmm0 ; ret  — zeroing idiom, xmm0 is a def not an input
        assert_eq!(shape_of(&[0x66, 0x0F, 0xEF, 0xC0, 0xC3]).key(), (0, 0));
    }

    #[test]
    fn pcmpeqd_self_all_ones_is_not_a_float_arg() {
        // pcmpeqd xmm2, xmm2 ; ret  — all-ones idiom (result independent of xmm2's prior value)
        assert_eq!(shape_of(&[0x66, 0x0F, 0x76, 0xD2, 0xC3]).key(), (0, 0));
    }

    // --- scalar-SSE merge-only writes must NOT be counted as float arguments (regression guards) ---

    #[test]
    fn cvtsi2sd_dest_is_not_a_float_arg() {
        // cvtsi2sd xmm0, edi ; ret  — converts an INT arg (edi) to double; xmm0 is a merge-write dest,
        // its read-write access is an upper-lane artifact, not a float argument => (1 int, 0 float).
        assert_eq!(shape_of(&[0xF2, 0x0F, 0x2A, 0xC7, 0xC3]).key(), (1, 0));
    }

    #[test]
    fn cvtsi2sd_shape_is_codegen_invariant() {
        // The SAME function must shape identically whether or not the compiler emits a
        // dependency-breaking `xorps xmm0,xmm0` before the convert.
        let bare = shape_of(&[0xF2, 0x0F, 0x2A, 0xC7, 0xC3]); // cvtsi2sd xmm0,edi ; ret
        let broken = shape_of(&[0x0F, 0x57, 0xC0, 0xF2, 0x0F, 0x2A, 0xC7, 0xC3]); // xorps xmm0,xmm0 ; …
        assert_eq!(bare.key(), (1, 0));
        assert_eq!(broken.key(), bare.key());
    }

    #[test]
    fn movsd_reg_dest_does_not_inflate_float_args() {
        // movsd xmm2, xmm0 ; ret  — copies float arg xmm0 into scratch xmm2; only xmm0 is an argument.
        assert_eq!(shape_of(&[0xF2, 0x0F, 0x10, 0xD0, 0xC3]).key(), (0, 1));
    }

    #[test]
    fn sqrtsd_scratch_dest_not_counted() {
        // sqrtsd xmm3, xmm1 ; ret  — result into scratch xmm3 from float arg xmm1; xmm1 (slot 1) fills
        // xmm0 by contiguity => 2 float args, NOT 4 (xmm3 the merge-dest must not inflate the count).
        assert_eq!(shape_of(&[0xF2, 0x0F, 0x51, 0xD9, 0xC3]).key(), (0, 2));
    }

    // --- return class ---

    #[test]
    fn scalar_return_is_int() {
        // mov eax, edi ; ret  — returns a scalar in RAX.
        assert_eq!(shape_of(&[0x89, 0xF8, 0xC3]).ret_class, RetClass::Int);
    }

    #[test]
    fn float_return_is_float() {
        // movsd xmm0, xmm1 ; ret  — the return register is XMM0.
        assert_eq!(
            shape_of(&[0xF2, 0x0F, 0x10, 0xC1, 0xC3]).ret_class,
            RetClass::Float
        );
    }

    #[test]
    fn return_this_pointer_is_int_not_byval() {
        // mov rax, rdi ; ret  — returns the RDI pointer but never writes THROUGH it, so it's a plain
        // pointer return (`return this`), not an sret trap.
        assert_eq!(shape_of(&[0x48, 0x89, 0xF8, 0xC3]).ret_class, RetClass::Int);
    }

    #[test]
    fn sret_write_through_rdi_is_byval() {
        // mov [rdi], rsi ; mov rax, rdi ; ret  — writes the result through the incoming RDI output
        // pointer AND returns it => the by-value (sret) shape that is unsafe to blind-call.
        assert_eq!(
            shape_of(&[0x48, 0x89, 0x37, 0x48, 0x89, 0xF8, 0xC3]).ret_class,
            RetClass::ByValue
        );
    }

    #[test]
    fn bare_ret_is_void() {
        // ret  — no result register written before returning.
        assert_eq!(shape_of(&[0xC3]).ret_class, RetClass::Void);
    }
}
