callgraph_verify/

disasm.rs

//! Recursive-descent function disassembly over real PE bytes (iced-x86).
//!
//! Given a function entry VA, we decode instructions and follow the function's
//! *intra-function* control flow (conditional branches + their fallthrough,
//! unconditional `JMP rel`) until every reachable path terminates at a `RET` /
//! interrupt / undecodable byte. Along the way we harvest, purely from the
//! decoded operands, every fact that lets us catch a rebuild hallucination:
//!
//!   * `direct_calls`  — targets of `CALL rel16/rel32` (the call graph proper)
//!   * `tail_calls`    — `JMP rel` leaving the body (a tail call into another fn)
//!   * `indirect_calls`/`indirect_jumps` — counts of register/memory dispatch
//!   * `vcall_slots`   — displacements of `CALL/JMP [reg+0xNN]` (vtable slots)
//!   * `iat_calls`     — absolute targets of `CALL [0x00a28xxx]` (IAT imports)
//!   * `field_offsets` — displacements of `[reg+0xNN]` accesses (struct fields),
//!     excluding ESP/EBP (those are stack locals/args)
//!   * `globals`       — absolute data addresses referenced as `[0x00bXXXXX]`
//!   * `imm_addrs`     — in-image immediates (`MOV [esi],0x00a7e38c` vtable VAs,
//!     function pointers, string-literal addresses)
//!
//! This is *structural* ground truth: no decompiler in the loop.

use crate::pe::Pe;
use iced_x86::{Decoder, DecoderOptions, FlowControl, Instruction, Mnemonic, OpKind, Register};
use std::collections::{BTreeMap, BTreeSet};

/// Upper bound on how far past the entry we follow control flow. Oblivion's
/// largest functions are a few KB; 64 KB stops a mis-decoded `JMP` from walking
/// the whole section.
const MAX_FN_SPAN: u64 = 0x10000;
/// `[reg+disp]` with disp below this is treated as a struct-field access; larger
/// displacements are almost certainly absolute addresses or noise.
const MAX_FIELD_OFFSET: u64 = 0x8000;

/// The recovered "kind" of a single call argument, harvested from the instruction that set it
/// up (a register load or a `PUSH`). Coarse on purpose — only the distinctions a faithful port
/// must preserve, and only those a short local scan can recover without false positives.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ArgKind {
    /// An immediate constant (`PUSH 0x10`, `MOV ECX,0`, `XOR ECX,ECX`).
    Imm(u64),
    /// A struct-field load `[reg+disp]` off a non-stack base (`PUSH [ESI+0xc]`).
    Field(u64),
    /// An absolute in-image data reference `[0x00bXXXXX]` (a global).
    Global(u64),
    /// Address of a stack local (`LEA ECX,[EBP-8]` / `LEA EAX,[ESP+4]`) — an out-param / temp.
    LocalAddr,
    /// Anything not confidently classified (a register value, computed expression, prior result).
    Other,
}

/// One direct call site with the argument set-up recovered by a short backward scan of its basic
/// block. `ecx`/`edx` are the fastcall/thiscall register args (the receiver lands in `ecx`); `stack`
/// is the pushed args in source (left-to-right) order. `None`/empty means "not recovered in-block".
#[derive(Debug, Clone, Default)]
pub struct CallSite {
    pub target: u64,
    pub ecx: Option<ArgKind>,
    pub edx: Option<ArgKind>,
    pub stack: Vec<ArgKind>,
}

#[derive(Debug, Default)]
pub struct CallGraph {
    pub entry: u64,
    /// Direct `CALL` targets → number of call sites hitting that target.
    pub direct_calls: BTreeMap<u64, u32>,
    /// Per-call-site recovered argument shapes (one entry per direct `CALL` instruction).
    pub call_sites: Vec<CallSite>,
    /// Direct `JMP rel` targets that leave the walked body (tail calls).
    pub tail_calls: BTreeSet<u64>,
    pub indirect_calls: u32,
    pub indirect_jumps: u32,
    /// vtable slot displacements seen in `CALL/JMP [reg+0xNN]`.
    pub vcall_slots: BTreeSet<u64>,
    /// IAT import slot addresses seen in `CALL [0x00a28xxx]`.
    pub iat_calls: BTreeSet<u64>,
    /// Struct-field displacements off a non-stack base register.
    pub field_offsets: BTreeSet<u64>,
    /// Absolute data addresses referenced as memory operands (globals).
    pub globals: BTreeSet<u64>,
    /// In-image immediate values (vtable VAs, fn pointers, string addresses).
    pub imm_addrs: BTreeSet<u64>,
    /// Highest instruction-end VA reached — the function's apparent extent.
    pub max_end: u64,
    pub decode_errors: u32,
    /// First few decode failures as `(ip, raw_bytes)`. Decoders hitting
    /// invalid opcodes is normal at padding/alignment/data boundaries, but a
    /// long run inside what is supposed to be a function usually means the
    /// entry VA is wrong, or the function crosses into embedded data. The
    /// hex bytes here are enough to tell which.
    pub decode_error_samples: Vec<(u64, Vec<u8>)>,
}

impl CallGraph {
    pub fn call_targets(&self) -> BTreeSet<u64> {
        self.direct_calls.keys().copied().collect()
    }

    /// Every in-image address this function references — call targets, tail-call
    /// JMPs, IAT slots, global data refs and baked immediates (vtable VAs / fn
    /// pointers / string addresses). Used to detect inlining: if callee C's whole
    /// footprint also appears in caller F's footprint, the binary pasted C's body
    /// into F, so a Rust call F→C is faithful (un-inlined), not hallucinated.
    pub fn footprint(&self) -> BTreeSet<u64> {
        let mut s: BTreeSet<u64> = self.direct_calls.keys().copied().collect();
        s.extend(&self.tail_calls);
        s.extend(&self.iat_calls);
        s.extend(&self.globals);
        s.extend(&self.imm_addrs);
        s
    }
}

/// Harvest every memory/immediate operand fact from one decoded instruction.
fn scan_operands(pe: &Pe, instr: &Instruction, cg: &mut CallGraph) {
    for i in 0..instr.op_count() {
        match instr.op_kind(i) {
            OpKind::Memory => {
                let base = instr.memory_base();
                let index = instr.memory_index();
                let disp = instr.memory_displacement64();
                if base == Register::None && index == Register::None {
                    // absolute `[0x...]` — a global/static data reference.
                    // Skip LEA: it computes an address without dereferencing, so
                    // the operand isn't a memory reference. (BUG #14)
                    if pe.section_of(disp).is_some() && instr.mnemonic() != Mnemonic::Lea {
                        cg.globals.insert(disp);
                    }
                } else if pe.section_of(disp).is_some()
                    && instr.mnemonic() != Mnemonic::Lea
                {
                    // Indexed `[base + index*scale + disp]` (or `[reg+disp]`)
                    // off a large in-image displacement — a table base, not a
                    // struct field (e.g. FALL.EXE `[eax*4+0x18F01]`). Record it
                    // as a global so bodies may bake the true address without
                    // tripping HALLUCINATED-ADDR. Small offsets miss
                    // `section_of` and fall through to the field branch below;
                    // LEA never dereferences (BUG #14), so it stays excluded.
                    cg.globals.insert(disp);
                } else if base != Register::None
                    && base != Register::ESP
                    && base != Register::EBP
                    && index == Register::None
                    && disp > 0
                    && disp < MAX_FIELD_OFFSET
                    // LEA computes an address (pointer math / &field), it doesn't
                    // dereference a struct field — skip it to avoid e.g. the
                    // `LEA EDX,[EAX+1]` of a strlen loop looking like field 0x1.
                    && instr.mnemonic() != Mnemonic::Lea
                {
                    // `[reg+0xNN]` off a non-stack base — a struct field access.
                    cg.field_offsets.insert(disp);
                }
            }
            OpKind::Immediate8
            | OpKind::Immediate8to16
            | OpKind::Immediate8to32
            | OpKind::Immediate16
            | OpKind::Immediate32 => {
                let v = instr.immediate(i);
                // In-image immediate = a baked address (vtable VA, fn ptr, string).
                if pe.section_of(v).is_some() {
                    cg.imm_addrs.insert(v);
                }
            }
            _ => {}
        }
    }
}

/// Classify operand `i` of `instr` as an argument kind (immediate / field / global / other).
fn classify_operand(pe: &Pe, instr: &Instruction, i: u32) -> ArgKind {
    match instr.op_kind(i) {
        OpKind::Immediate8
        | OpKind::Immediate8to16
        | OpKind::Immediate8to32
        | OpKind::Immediate16
        | OpKind::Immediate32 => ArgKind::Imm(instr.immediate(i)),
        OpKind::Memory => {
            let base = instr.memory_base();
            let index = instr.memory_index();
            let disp = instr.memory_displacement64();
            if base == Register::None && index == Register::None {
                if pe.section_of(disp).is_some() {
                    ArgKind::Global(disp)
                } else {
                    ArgKind::Other
                }
            } else if base != Register::ESP
                && base != Register::EBP
                && base != Register::None
                && index == Register::None
                && disp > 0
                && disp < MAX_FIELD_OFFSET
            {
                ArgKind::Field(disp)
            } else {
                ArgKind::Other
            }
        }
        _ => ArgKind::Other,
    }
}

/// Classify the *source* of a register-setup instruction (`MOV reg,src` / `LEA reg,src` /
/// `XOR reg,reg`). Returns the arg kind the register holds going into the following `CALL`.
fn classify_reg_src(pe: &Pe, instr: &Instruction) -> ArgKind {
    match instr.mnemonic() {
        Mnemonic::Lea => {
            let base = instr.memory_base();
            if base == Register::EBP || base == Register::ESP {
                ArgKind::LocalAddr
            } else {
                ArgKind::Other // lea of a field address — not distinguished in v1
            }
        }
        Mnemonic::Xor => ArgKind::Imm(0), // only called when op0==op1 (zero idiom)
        Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx => classify_operand(pe, instr, 1),
        _ => ArgKind::Other,
    }
}

/// True if `instr` writes register `reg` as its destination (op0) via a recognized setup form.
fn writes_reg(instr: &Instruction, reg: Register) -> bool {
    if instr.op0_kind() != OpKind::Register || instr.op0_register() != reg {
        return false;
    }
    match instr.mnemonic() {
        Mnemonic::Mov | Mnemonic::Lea | Mnemonic::Movzx | Mnemonic::Movsx => true,
        // XOR reg,reg is the zero idiom only when both operands are the same register.
        Mnemonic::Xor => instr.op1_kind() == OpKind::Register && instr.op1_register() == reg,
        _ => false,
    }
}

/// Recover the argument shape of a `CALL` from the instructions preceding it in its basic block
/// (`buf` excludes the CALL itself). Scans backward: the first `PUSH`es encountered are the
/// left-to-right stack args (cdecl/stdcall push right-to-left, so backward order = source order),
/// and the first writes to ECX/EDX are the fastcall/thiscall register args.
pub(crate) fn recover_call_site(pe: &Pe, buf: &[Instruction], target: u64) -> CallSite {
    let mut cs = CallSite {
        target,
        ..Default::default()
    };
    for instr in buf.iter().rev().take(64) {
        // Stop at a prologue / frame boundary: an explicit ESP adjust (`sub esp,N`) or any
        // FS-relative access (SEH frame set-up). Argument pushes never precede these, so this
        // keeps callee-saved / SEH-cookie pushes out of the recovered argument list.
        if is_frame_boundary(instr) {
            break;
        }
        if instr.mnemonic() == Mnemonic::Push {
            cs.stack.push(classify_operand(pe, instr, 0));
        } else if cs.ecx.is_none() && writes_reg(instr, Register::ECX) {
            cs.ecx = Some(classify_reg_src(pe, instr));
        } else if cs.edx.is_none() && writes_reg(instr, Register::EDX) {
            cs.edx = Some(classify_reg_src(pe, instr));
        }
    }
    cs
}

/// A prologue / stack-frame boundary that argument set-up never precedes: an `ADD/SUB ESP,imm`
/// or any instruction touching FS-relative memory (the SEH frame).
fn is_frame_boundary(instr: &Instruction) -> bool {
    if matches!(instr.mnemonic(), Mnemonic::Add | Mnemonic::Sub)
        && instr.op0_kind() == OpKind::Register
        && instr.op0_register() == Register::ESP
    {
        return true;
    }
    (0..instr.op_count())
        .any(|i| instr.op_kind(i) == OpKind::Memory && instr.memory_segment() == Register::FS)
}

/// If `addr` is a one-instruction thunk (`JMP target`), return `target`. Callers often invoke a
/// function through such a thunk (a tail-jump stub, or `/OPT:ICF`-folded forwarder), so the byte
/// call edge points at the thunk, not the real callee — following it one level lets the call-graph
/// check see the true target.
pub fn resolve_thunk(pe: &Pe, addr: u64) -> Option<u64> {
    let bytes = pe.read(addr, 16)?;
    let mut dec = Decoder::with_ip(32, bytes, addr, DecoderOptions::NONE);
    if !dec.can_decode() {
        return None;
    }
    let instr = dec.decode();
    if instr.flow_control() == FlowControl::UnconditionalBranch {
        let t = instr.near_branch_target();
        if t != 0 {
            return Some(t);
        }
    }
    None
}

/// Does `instr` *read* `reg` (as a source register or a memory base/index)? Excludes the
/// destination register operand (op0) so a pure write like `MOV ECX,EAX` doesn't count as a read.
fn reads_reg(instr: &Instruction, reg: Register) -> bool {
    for i in 0..instr.op_count() {
        match instr.op_kind(i) {
            OpKind::Memory => {
                if instr.memory_base() == reg || instr.memory_index() == reg {
                    return true;
                }
            }
            OpKind::Register if i > 0 && instr.op_register(i) == reg => {
                return true;
            }
            _ => {}
        }
    }
    false
}

/// Heuristically determine whether the function at `addr` takes register parameters in ECX / EDX
/// (i.e. is thiscall/fastcall). Scans the entry block for a READ of the register before any write:
/// thiscall reads ECX (the `this`), fastcall also reads EDX, cdecl/stdcall read neither. Used to
/// align a caller's recovered register args to the callee's actual parameters (so a stray EDX in a
/// caller is not mistaken for an argument to a thiscall callee). Conservative: under-detection just
/// means a position is skipped, never a false mismatch.
pub fn reg_params(pe: &Pe, addr: u64) -> (bool, bool) {
    let Some(bytes) = pe.read(addr, 32 * 32) else {
        return (false, false);
    };
    let mut dec = Decoder::with_ip(32, bytes, addr, DecoderOptions::NONE);
    let mut instr = Instruction::default();
    let (mut ecx_written, mut edx_written) = (false, false);
    let (mut ecx_param, mut edx_param) = (false, false);
    let mut n = 0;
    while dec.can_decode() && n < 24 {
        dec.decode_out(&mut instr);
        if instr.is_invalid() {
            break;
        }
        if !ecx_written && reads_reg(&instr, Register::ECX) {
            ecx_param = true;
        }
        if !edx_written && reads_reg(&instr, Register::EDX) {
            edx_param = true;
        }
        if instr.op0_kind() == OpKind::Register {
            match instr.op0_register() {
                Register::ECX => ecx_written = true,
                Register::EDX => edx_written = true,
                _ => {}
            }
        }
        // A leaf that immediately forwards its incoming registers to another call is still
        // thiscall/fastcall even though iced-x86 cannot represent the implicit register use on
        // the CALL itself. Preserve untouched ECX/EDX on the first instruction so constructors
        // such as BSScissorTriShape (which forwards both registers to its base ctor) align stack
        // arguments correctly. Later calls retain the caller-saved rule below.
        if n == 0
            && matches!(instr.flow_control(), FlowControl::Call)
            && instr.near_branch_target() != instr.next_ip()
        {
            if !ecx_written {
                ecx_param = true;
            }
            if !edx_written {
                edx_param = true;
            }
        }
        // Per the x86 cdecl/fastcall/thiscall conventions, ECX and EDX are caller-saved:
        // any CALL in the entry block clobbers them, so a later READ of ECX/EDX is not
        // evidence that the callee takes that register as a parameter. Mark them written
        // across the CALL — without this, a stray `MOV EAX, [ECX]` after a prior CALL
        // would falsely set `ecx_param = true` and pollute ARG-SHAPE diffs.
        if matches!(instr.flow_control(), FlowControl::Call) {
            ecx_written = true;
            edx_written = true;
        }
        if !matches!(instr.flow_control(), FlowControl::Next | FlowControl::Call) {
            break;
        }
        n += 1;
    }
    (ecx_param, edx_param)
}

/// Walk the function at `entry` and return its facts, sourced purely from bytes.
pub fn analyze(pe: &Pe, entry: u64) -> Result<CallGraph, String> {
    if !pe.is_executable_va(entry) {
        return Err(format!(
            "{entry:#010x} is not inside any PE section of this image\n{}",
            pe.describe_sections(entry)
        ));
    }
    let mut cg = CallGraph {
        entry,
        ..Default::default()
    };
    let lo = entry;
    let hi = entry.saturating_add(MAX_FN_SPAN);

    let mut work = vec![entry];
    let mut seen: BTreeSet<u64> = BTreeSet::new();
    // Per-branch-target seed: when we encounter `Jcc target` / `JMP target` /
    // indirect-jump case, the args set up before the branch (PUSHes, MOV-to-ECX,
    // etc.) also apply to the target block's first CALL — the predecessor and
    // target share the same register/stack state at the branch instruction. The
    // outer loop resets `block_buf` when it pops a new `start`, so without this
    // seed map the target's calls would be processed with an empty `block_buf`
    // and `recover_call_site` would undercount arguments. (BUG #8 fix.)
    //
    // Multiple predecessors may push the same target — `entry().or_insert_with`
    // keeps the FIRST push, since the args are equivalent along either path up
    // to the branch instruction itself (the differing branches only affect the
    // value of flags, which `recover_call_site` doesn't read).
    let mut branch_seed: std::collections::HashMap<u64, Vec<Instruction>> =
        std::collections::HashMap::new();

    while let Some(start) = work.pop() {
        if !seen.insert(start) {
            continue;
        }
        let Some(bytes) = pe.read(start, 32 * 64) else {
            continue;
        };
        let mut dec = Decoder::with_ip(32, bytes, start, DecoderOptions::NONE);
        let mut instr = Instruction::default();
        // Straight-line instruction buffer for the current block, used to recover the argument
        // set-up preceding each CALL. Reset at block boundaries (branches) and after each call.
        // For non-entry `start`s, seed from the predecessor branch (see `branch_seed`).
        let mut block_buf: Vec<Instruction> = branch_seed.remove(&start).unwrap_or_default();

        loop {
            if !dec.can_decode() {
                break;
            }
            let ip = dec.ip();
            if ip != start && seen.contains(&ip) {
                break; // block already covered by another path
            }
            if ip > start {
                seen.insert(ip);
            }
            dec.decode_out(&mut instr);
            if instr.is_invalid() {
                cg.decode_errors += 1;
                // Keep up to 5 samples so the user can tell at a glance
                // whether the failures are padding (single 0x00/0xCC/0x90),
                // an aligned jump table (4 distinct bytes), or something
                // structural (random bytes that decode to nothing — entry
                // VA is probably wrong).
                if cg.decode_error_samples.len() < 5 {
                    let bytes = pe.read(ip, 8).unwrap_or(&[]).to_vec();
                    cg.decode_error_samples.push((ip, bytes));
                }
                break;
            }
            cg.max_end = cg.max_end.max(instr.next_ip());
            scan_operands(pe, &instr, &mut cg);
            block_buf.push(instr);

            match instr.flow_control() {
                FlowControl::Call => {
                    let t = instr.near_branch_target();
                    if t != 0 {
                        *cg.direct_calls.entry(t).or_insert(0) += 1;
                        let prior = &block_buf[..block_buf.len() - 1]; // exclude the CALL itself
                        let cs = recover_call_site(pe, prior, t);
                        cg.call_sites.push(cs);
                    }
                    block_buf.clear(); // args after a call start fresh
                }
                FlowControl::IndirectCall => {
                    cg.indirect_calls += 1;
                    record_indirect_target(&instr, pe, &mut cg);
                    block_buf.clear(); // an indirect call's args must not leak into the next call
                }
                FlowControl::ConditionalBranch => {
                    let t = instr.near_branch_target();
                    if (lo..hi).contains(&t) {
                        branch_seed.entry(t).or_insert_with(|| block_buf.clone());
                        work.push(t);
                    } else {
                        cg.tail_calls.insert(t);
                    }
                    // Don't clear `block_buf` here: the args set up before the
                    // conditional apply to BOTH the taken and the fallthrough
                    // paths (they share register/stack state at the branch).
                    // Clearing would make the fallthrough's first CALL recover
                    // an empty argument set. The buffer is reset by the next
                    // CALL we encounter (Call/IndirectCall arms above clear
                    // post-call) or by an explicit boundary.
                }
                FlowControl::UnconditionalBranch => {
                    let t = instr.near_branch_target();
                    if (lo..hi).contains(&t) {
                        // Single-JMP entry stub: a function whose entire
                        // entry block is one unconditional JMP (sibling-call
                        // `JMP tail_target`, ICF forwarder, `jmp` over no
                        // bytes) transfers control elsewhere instead of
                        // containing its own body. Following it would merge
                        // the callee's body (and globals) into the caller,
                        // hiding the tail edge the gate checks. Scoped to
                        // exactly one instruction so a `JMP` that skips an
                        // inline jump table after a prologue (multi-instr
                        // entry block) keeps the follow behavior below.
                        if start == lo && block_buf.len() == 1 {
                            cg.tail_calls.insert(t);
                        } else {
                            branch_seed.entry(t).or_insert_with(|| block_buf.clone());
                            work.push(t);
                        }
                    } else {
                        cg.tail_calls.insert(t);
                    }
                    break;
                }
                FlowControl::IndirectBranch => {
                    cg.indirect_jumps += 1;
                    record_indirect_target(&instr, pe, &mut cg);
                    let mut cases: Vec<u64> = Vec::new();
                    // Indirect-jump target recovery, by operand shape:
                    // * SIB-indexed `[base+idx*s+TABLE]` (MSVC-style): the
                    //   existing follow_jump_table read below. Unchanged.
                    // * Unindexed `[reg+TABLE]` with NO base register, i.e.
                    //   an absolute `[disp]` (IAT thunks, import jumps):
                    //   NEVER followed. The disp names data, not code, and
                    //   the fallthrough past a non-returning jump is dead
                    //   code -- decoding it attributes neighboring
                    //   functions' calls here (verified: 9KB/20-call blowup
                    //   on noivilbo 0x006ff480, whose IAT tail-jmps sit
                    //   after cmp/jcc).
                    // * Unindexed `[reg+TABLE]` WITH a base register: Watcom
                    //   switch lowering (`shl reg,2` then `jmp [reg+TABLE]`,
                    //   e.g. FALL.EXE 0x602c0). The table entries are
                    //   pre-fixup bytes on LE, so the read below yields
                    //   nothing; instead continue decoding at the fallthrough
                    //   address -- the case bodies tile the bytes right after
                    //   the dispatch (verified: 602c0's five hidden 7d6ae
                    //   calls, previously mis-reported as HALLUCINATED).
                    //   Gated on switch evidence (a compare/test plus a
                    //   conditional branch in this block) so vtable tail-jumps
                    //   like `jmp [ecx+0x44]` keep today's behavior.
                    let has_index = instr.memory_index() != Register::None;
                    // The base must be produced locally (e.g. the scaled
                    // index): an incoming register like a vtable `this`
                    // is never a switch table.
                    let mut base_written = false;
                    if instr.memory_base() != Register::None {
                        // CMP/TEST only write flags, not their operands --
                        // without this carve-out the evidence check itself
                        // (`cmp reg,...`) would trivially satisfy "written".
                        for b in block_buf.iter() {
                            match b.mnemonic() {
                                Mnemonic::Cmp | Mnemonic::Test => continue,
                                _ => {}
                            }
                            if b.op0_kind() == OpKind::Register
                                && b.op0_register() == instr.memory_base()
                            {
                                base_written = true;
                                break;
                            }
                        }
                    }
                    let unindexed_switch =
                        !has_index && base_written && is_switch_dispatch(&block_buf);
                    follow_jump_table(&instr, pe, lo, hi, unindexed_switch, &mut cases);
                    // GCC -O0 switch lowering keeps the scale in a separate
                    // `shl`, adds the table base as an immediate, loads the
                    // entry, then jumps through a register:
                    //   shl $2,%eax; add $TABLE,%eax; mov (%eax),%eax; jmp *%eax
                    // The JMP itself has no memory operand, so the readers
                    // above see nothing (MinGW `c_callgraph_od` classify:
                    // 1 indirect jump, 0 cases, 0 calls). Recover the table
                    // base from the `add` and read it like any other table,
                    // gated on the same switch evidence so a vtable
                    // `mov eax,[ecx+slot]; jmp eax` (no bound check, `this`
                    // never `add`ed with an in-image base) stays untouched.
                    if cases.is_empty()
                        && instr.op_count() > 0
                        && instr.op_kind(0) == OpKind::Register
                        && is_switch_dispatch(&block_buf)
                    {
                        if let Some(table) =
                            register_table_base(&block_buf, pe, instr.op_register(0))
                        {
                            read_table_entries(pe, table, lo, hi, &mut cases);
                        }
                    }
                    for c in &cases {
                        branch_seed.entry(*c).or_insert_with(|| block_buf.clone());
                        work.push(*c);
                    }
                    if unindexed_switch {
                        let ft = instr.next_ip();
                        if (lo..hi).contains(&ft) {
                            branch_seed.entry(ft).or_insert_with(|| block_buf.clone());
                            work.push(ft);
                        }
                    }
                    break;
                }
                FlowControl::Return | FlowControl::Exception => break,
                // NOTE: `Interrupt` (`int n`) is deliberately NOT terminal:
                // DOS `int 21h` / DPMI `int 31h` are API calls that return
                // normally, so decoding must continue at `next_ip()`.
                // Treating them like `ret` truncated listings and hid
                // calls reachable only via post-`int` fallthrough.
                FlowControl::Next | FlowControl::Interrupt | FlowControl::XbeginXabortXend => {}
            }
        }
    }
    Ok(cg)
}

/// Addresses a tail-call target references (its globals, baked immediates,
/// and IAT slots), unioned over every `JMP`-out edge of `cg`.
///
/// A tail-call `JMP G` transfers control to G, so G's data references execute
/// as part of THIS function. The call-graph diff (`diff_body_calls`) already
/// counts a tail target's *callees* as acceptable spellings; this is the
/// address-side mirror for the HALLUCINATED-ADDR check: a port that
/// duplicates a shared tail body (e.g. a predicate tail-jumped into from two
/// origins) bakes the tail's addresses faithfully. Only ever ADDS acceptance
/// — a target that fails to analyze is skipped, never an error.
pub fn tail_referenced(pe: &Pe, cg: &CallGraph) -> BTreeSet<u64> {
    let mut out = BTreeSet::new();
    for t in &cg.tail_calls {
        if let Ok(tcg) = analyze(pe, *t) {
            out.extend(&tcg.globals);
            out.extend(&tcg.imm_addrs);
            out.extend(&tcg.iat_calls);
        }
    }
    out
}

/// Follow an MSVC-style `switch` jump table: `JMP dword ptr [idx*4 + TABLE_VA]`.
/// `TABLE_VA` (the operand displacement, an in-image .rdata address) holds an
/// array of u32 case addresses. We read entries until one falls outside the
/// function `[lo, hi)` span -- switch cases always point back into the function,
/// so that's a reliable terminator. Recovers the case bodies (and their `CALL`s)
/// that a plain recursive descent, unable to follow an indirect jump, would miss.
///
/// `allow_unindexed` covers Watcom-style lowering where the scale lives in a
/// separate `shl` and the dispatch is a plain `[reg+TABLE]` with no SIB index
/// (e.g. FALL.EXE `0x602c0`: `jmp dword ptr cs:[eax+0x502D8]`). Only set it
/// when the block already looks like a switch (see `is_switch_dispatch`):
/// without that evidence a vtable tail-jump would read arbitrary words as
/// "cases". On LE binaries the table entries are pre-fixup bytes, so the
/// read yields no in-range targets and the fallthrough arm below carries
/// the recovery instead.
pub(crate) fn follow_jump_table(
    instr: &Instruction,
    pe: &Pe,
    lo: u64,
    hi: u64,
    allow_unindexed: bool,
    work: &mut Vec<u64>,
) {
    for i in 0..instr.op_count() {
        if instr.op_kind(i) != OpKind::Memory {
            continue;
        }
        let disp = instr.memory_displacement64();
        if instr.memory_index() == Register::None && !allow_unindexed {
            return;
        }
        if pe.section_of(disp).is_none() {
            return;
        }
        read_table_entries(pe, disp, lo, hi, work);
        return;
    }
}

/// Read a u32 VA jump table at `table_va`, pushing entries while they point
/// back into the function `[lo, hi)` (the terminator shared by every table
/// shape: switch cases always target the function's own bodies).
fn read_table_entries(pe: &Pe, table_va: u64, lo: u64, hi: u64, work: &mut Vec<u64>) {
    let mut k = 0u64;
    while k < 1024 {
        let Some(b) = pe.read(table_va + k * 4, 4) else {
            break;
        };
        if b.len() < 4 {
            break;
        }
        let target = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as u64;
        if !(lo..hi).contains(&target) {
            break; // first entry outside the function ends the table
        }
        work.push(target);
        k += 1;
    }
}

/// Table base for a register-indirect switch dispatch (`jmp *%reg`).
/// Finds the entry load `mov %reg,[%base]` nearest the jump, then the
/// `add %base,$TABLE` / `lea %base,[...]` / `mov %base,$TABLE` that
/// established it with an in-image table address. Returns `None` when the
/// jump register comes straight from memory with no locally-built table
/// (e.g. a vtable `mov %eax,[%ecx+slot]; jmp *%eax`, whose `%ecx` is an
/// incoming `this` never `add`ed with an in-image base).
fn register_table_base(block: &[Instruction], pe: &Pe, jump_reg: Register) -> Option<u64> {
    // Nearest load into the jump register: `mov %reg,[%base]`.
    let mut load_base: Option<Register> = None;
    for b in block.iter().rev() {
        if matches!(
            b.mnemonic(),
            Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx
        ) && b.op0_kind() == OpKind::Register
            && b.op0_register() == jump_reg
            && b.op1_kind() == OpKind::Memory
        {
            load_base = Some(b.memory_base());
            break;
        }
    }
    let base = load_base?;
    // The table address built into the base register further up.
    for b in block.iter().rev() {
        match b.mnemonic() {
            Mnemonic::Add => {
                if b.op0_kind() == OpKind::Register && b.op0_register() == base {
                    for i in 0..b.op_count() {
                        if matches!(
                            b.op_kind(i),
                            OpKind::Immediate8
                                | OpKind::Immediate8to16
                                | OpKind::Immediate8to32
                                | OpKind::Immediate16
                                | OpKind::Immediate32
                        ) && pe.section_of(b.immediate(i)).is_some()
                        {
                            return Some(b.immediate(i));
                        }
                    }
                }
            }
            Mnemonic::Lea => {
                if b.op0_kind() == OpKind::Register
                    && b.op0_register() == base
                    && b.op1_kind() == OpKind::Memory
                    && pe.section_of(b.memory_displacement64()).is_some()
                {
                    return Some(b.memory_displacement64());
                }
            }
            Mnemonic::Mov | Mnemonic::Movzx | Mnemonic::Movsx => {
                if b.op0_kind() == OpKind::Register && b.op0_register() == base {
                    for i in 0..b.op_count() {
                        if matches!(
                            b.op_kind(i),
                            OpKind::Immediate8
                                | OpKind::Immediate8to16
                                | OpKind::Immediate8to32
                                | OpKind::Immediate16
                                | OpKind::Immediate32
                        ) && pe.section_of(b.immediate(i)).is_some()
                        {
                            return Some(b.immediate(i));
                        }
                    }
                }
            }
            _ => {}
        }
    }
    None
}

/// Switch-dispatch evidence: a compare/test plus a conditional branch in
/// the current straight-line block (e.g. `cmp edx,6; ja default` ahead of
/// `jmp [reg+TABLE]`). Used to license fallthrough decoding past an
/// indirect jump only where a switch table is plausible, so vtable
/// tail-jumps (no bound check) never gain fallthrough coverage.
fn is_switch_dispatch(block: &[Instruction]) -> bool {
    let mut has_cmp = false;
    let mut has_jcc = false;
    for i in block {
        match i.mnemonic() {
            Mnemonic::Cmp | Mnemonic::Test => has_cmp = true,
            _ => {}
        }
        if matches!(i.flow_control(), FlowControl::ConditionalBranch) {
            has_jcc = true;
        }
    }
    has_cmp && has_jcc
}

/// Classify an indirect `CALL`/`JMP [...]` target: `[reg+disp]` is a vtable slot
/// (record the displacement); `[abs]` in the IAT range is an import (record addr).
fn record_indirect_target(instr: &Instruction, pe: &Pe, cg: &mut CallGraph) {
    // The branch target operand is the (single) memory operand, if any.
    for i in 0..instr.op_count() {
        if instr.op_kind(i) == OpKind::Memory {
            let base = instr.memory_base();
            let disp = instr.memory_displacement64();
            if base == Register::None && instr.memory_index() == Register::None {
                if pe.section_of(disp).is_some() {
                    cg.iat_calls.insert(disp); // CALL [0x00a28xxx]
                }
            } else if base != Register::None && disp < MAX_FIELD_OFFSET {
                cg.vcall_slots.insert(disp); // CALL [reg+slot]
            }
            return;
        }
    }
}