neo-decompiler 0.11.0

use std::collections::{BTreeSet, HashMap};

use crate::instruction::{Instruction, OpCode, Operand};
use crate::manifest::ContractManifest;

use super::control_flow::{collect_control_flow_edges, relative_target};
use super::manifest::offset_as_usize;

/// Build a sorted list of inferred method starts.
///
/// Sources include:
/// - script entry offset;
/// - manifest ABI offsets (when present);
/// - every `INITSLOT` instruction offset (compiler-emitted method prologues);
/// - offsets immediately following terminating instructions (`RET`, `THROW`,
///   `ABORT`, `ABORTMSG`) when control-flow indicates a detached chunk
///   (incoming from outside the baseline method, or no same-baseline incoming
///   to the remaining tail range).
#[must_use]
pub(in super::super::super) fn inferred_method_starts(
    instructions: &[Instruction],
    manifest: Option<&ContractManifest>,
) -> Vec<usize> {
    let mut starts = BTreeSet::new();
    if let Some(entry) = instructions.first() {
        starts.insert(entry.offset);
    }

    if let Some(manifest) = manifest {
        starts.extend(
            manifest
                .abi
                .methods
                .iter()
                .filter_map(|method| offset_as_usize(method.offset)),
        );
    }

    starts.extend(collect_initslot_offsets(instructions));
    starts.extend(collect_call_targets(instructions));
    let baseline_starts = starts.clone();
    starts.extend(collect_post_ret_method_offsets(
        instructions,
        &baseline_starts,
    ));
    starts.into_iter().collect()
}

/// Return the next known method start after `start`.
#[must_use]
pub(in super::super::super) fn next_inferred_method_offset(
    starts: &[usize],
    start: usize,
) -> Option<usize> {
    starts.iter().copied().find(|offset| *offset > start)
}

/// Return the argument count declared by an `INITSLOT` prologue at `start`.
#[must_use]
pub(in super::super::super) fn initslot_argument_count_at(
    instructions: &[Instruction],
    start: usize,
) -> Option<usize> {
    instructions
        .iter()
        .find(|ins| ins.offset == start && ins.opcode == OpCode::Initslot)
        .and_then(|ins| match &ins.operand {
            Some(Operand::Bytes(bytes)) if bytes.len() >= 2 => Some(bytes[1] as usize),
            _ => None,
        })
}

/// Collect offsets where the NEF script starts a new method (`INITSLOT`).
pub(in super::super::super) fn collect_initslot_offsets(
    instructions: &[Instruction],
) -> Vec<usize> {
    let mut offsets = instructions
        .iter()
        .filter(|ins| matches!(ins.opcode, OpCode::Initslot))
        .map(|ins| ins.offset)
        .collect::<Vec<_>>();
    offsets.sort_unstable();
    offsets.dedup();
    offsets
}

/// Collect targets of internal `CALL` / `CALL_L` instructions.
///
/// Each CALL target is a method entry point that may lack an `INITSLOT`
/// prologue (e.g. simple helpers that use no locals/arguments). Adding
/// these as baseline method starts prevents their bodies from being
/// inlined into the caller's method body.
pub(in super::super::super) fn collect_call_targets(instructions: &[Instruction]) -> Vec<usize> {
    let known_offsets: BTreeSet<usize> = instructions.iter().map(|ins| ins.offset).collect();
    let mut targets = Vec::new();
    for instruction in instructions {
        if matches!(instruction.opcode, OpCode::Call | OpCode::Call_L) {
            if let Some(target) = relative_target(instruction, &known_offsets) {
                targets.push(target);
            }
        }
    }
    targets.sort_unstable();
    targets.dedup();
    targets
}

fn collect_post_ret_method_offsets(
    instructions: &[Instruction],
    baseline_starts: &BTreeSet<usize>,
) -> Vec<usize> {
    let known_offsets: BTreeSet<usize> = instructions.iter().map(|ins| ins.offset).collect();
    let control_flow_edges = collect_control_flow_edges(instructions, &known_offsets);

    // Resolve the baseline method `[start, end)` that contains `offset`.
    let method_range = |offset: usize| -> (usize, usize) {
        let start = baseline_starts
            .range(..=offset)
            .next_back()
            .copied()
            .unwrap_or(offset);
        let end = baseline_starts
            .range((start + 1)..)
            .next()
            .copied()
            .unwrap_or(usize::MAX);
        (start, end)
    };

    let mut edges_by_target: HashMap<usize, Vec<usize>> = HashMap::new();
    // max_forward_in_method: baseline method start -> the furthest forward-edge
    // target that stays inside that method. Precomputing this once turns the
    // "is there an in-method edge past `next`?" test into an O(1) lookup instead
    // of an O(method length) scan per terminator. Without it the whole pass is
    // O(n²) and a crafted in-cap NEF (many crossing jumps in one method) hangs
    // the decompiler.
    let mut max_forward_in_method: HashMap<usize, usize> = HashMap::new();
    for &(source, target) in &control_flow_edges {
        edges_by_target.entry(target).or_default().push(source);
        let (m_start, m_end) = method_range(source);
        if target < m_end {
            max_forward_in_method
                .entry(m_start)
                .and_modify(|t| *t = (*t).max(target))
                .or_insert(target);
        }
    }

    let mut starts = instructions
        .windows(2)
        .filter_map(|pair| {
            let current = &pair[0];
            if !matches!(
                current.opcode,
                OpCode::Ret | OpCode::Throw | OpCode::Abort | OpCode::Abortmsg
            ) {
                return None;
            }
            let next = &pair[1];
            let (method_start, method_end) = method_range(current.offset);

            let incoming = edges_by_target.get(&next.offset);
            let has_incoming_from_same_baseline_method = incoming.is_some_and(|sources| {
                sources.iter().any(|&s| s >= method_start && s < method_end)
            });
            let has_incoming_from_other_baseline_method = incoming.is_some_and(|sources| {
                sources.iter().any(|&s| s < method_start || s >= method_end)
            });
            // Equivalent to scanning every in-method edge for a target in
            // (next.offset, method_end): the per-method maximum forward target
            // (already bounded to < method_end) exceeds next.offset iff one exists.
            let has_same_baseline_incoming_later_in_range = max_forward_in_method
                .get(&method_start)
                .is_some_and(|&max_target| max_target > next.offset);

            let detached_tail_after_terminator = has_incoming_from_other_baseline_method
                || (!has_incoming_from_same_baseline_method
                    && !has_same_baseline_incoming_later_in_range);
            detached_tail_after_terminator.then_some(next.offset)
        })
        .collect::<Vec<_>>();
    starts.sort_unstable();
    starts.dedup();
    starts
}