neo-decompiler 0.11.0

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

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

use super::super::super::helpers::{
    collect_call_targets, collect_initslot_offsets, find_manifest_entry_method, offset_as_usize,
    sanitize_identifier,
};
use super::super::call_graph::{
    calla_ldarg_index, calla_target_from_pusha, initslot_arg_count_at, trace_call_arg_source,
    CallArgSource,
};
use super::types::{MethodRef, MethodSpan, MethodTable};

impl MethodTable {
    /// Build a method table using stable method starts plus manifest metadata.
    ///
    /// Stable starts include the script entry, manifest ABI offsets, compiler
    /// `INITSLOT` prologues, and direct internal call targets. This keeps
    /// analysis aligned with inferred helpers without over-splitting detached
    /// tails that are only useful for presentation-time rendering.
    #[must_use]
    pub fn new(instructions: &[Instruction], manifest: Option<&ContractManifest>) -> Self {
        let script_start = instructions.first().map(|ins| ins.offset).unwrap_or(0);
        let script_end = instructions
            .last()
            .map(|ins| ins.offset.saturating_add(1))
            .unwrap_or(script_start);

        // Valid instruction start offsets. A CALLA pointer target (from PUSHA)
        // must land on one to be a real method start; an out-of-range target
        // would otherwise fabricate a phantom span/method. `collect_call_targets`
        // already applies the same filter to CALL/CALL_L targets.
        let known_offsets: BTreeSet<usize> = instructions.iter().map(|ins| ins.offset).collect();

        let mut manifest_index_by_start = BTreeMap::new();
        let entry_manifest = manifest.and_then(|manifest| {
            let entry_method = find_manifest_entry_method(manifest, script_start)?;
            let index = manifest
                .abi
                .methods
                .iter()
                .position(|candidate| std::ptr::eq(candidate, entry_method.0))?;
            Some((entry_method.0, index))
        });

        let mut starts = BTreeMap::new();
        starts.insert(script_start, ());
        for start in collect_initslot_offsets(instructions) {
            starts.insert(start, ());
        }
        for start in collect_call_targets(instructions) {
            starts.insert(start, ());
        }
        let mut callers_by_target: BTreeMap<usize, BTreeSet<usize>> = BTreeMap::new();
        for (index, instruction) in instructions.iter().enumerate() {
            if instruction.opcode == OpCode::CallA {
                if let Some(start) = calla_target_from_pusha(instructions, index)
                    .filter(|start| known_offsets.contains(start))
                {
                    starts.insert(start, ());
                    callers_by_target.entry(start).or_default().insert(index);
                }
                continue;
            }
            if matches!(instruction.opcode, OpCode::Call | OpCode::Call_L) {
                if let Some(target) = direct_call_target(instruction) {
                    callers_by_target.entry(target).or_default().insert(index);
                }
            }
        }

        resolve_calla_argument_targets(instructions, &mut starts, &mut callers_by_target);

        if let Some(manifest) = manifest {
            for (idx, method) in manifest.abi.methods.iter().enumerate() {
                if let Some(start) = offset_as_usize(method.offset) {
                    manifest_index_by_start.insert(start, idx);
                    starts.insert(start, ());
                }
            }
            if let Some((_, index)) = entry_manifest {
                manifest_index_by_start.entry(script_start).or_insert(index);
            }
        }

        let ordered_starts: Vec<usize> = starts.into_keys().collect();
        let mut spans = Vec::new();
        for (position, start) in ordered_starts.iter().copied().enumerate() {
            let end = ordered_starts
                .get(position + 1)
                .copied()
                .unwrap_or(script_end);
            let method = if let Some(manifest) = manifest {
                if let Some(index) = manifest_index_by_start.get(&start).copied() {
                    let manifest_method = &manifest.abi.methods[index];
                    MethodRef {
                        offset: start,
                        name: sanitize_identifier(&manifest_method.name),
                    }
                } else if start == script_start {
                    MethodRef {
                        offset: start,
                        name: entry_manifest
                            .as_ref()
                            .map(|(method, _)| sanitize_identifier(&method.name))
                            .unwrap_or_else(|| "script_entry".to_string()),
                    }
                } else {
                    MethodRef::synthetic(start)
                }
            } else if start == script_start {
                MethodRef {
                    offset: start,
                    name: "script_entry".to_string(),
                }
            } else {
                MethodRef::synthetic(start)
            };

            spans.push(MethodSpan { start, end, method });
        }

        spans.sort_by_key(|span| span.start);

        Self {
            spans,
            manifest_index_by_start,
        }
    }
}

fn resolve_calla_argument_targets(
    instructions: &[Instruction],
    starts: &mut BTreeMap<usize, ()>,
    callers_by_target: &mut BTreeMap<usize, BTreeSet<usize>>,
) {
    loop {
        // Sorted Vec of method-start offsets; rebuilt per fixpoint iter so
        // newly-inserted starts from the previous round are visible.
        let method_starts: Vec<usize> = starts.keys().copied().collect();
        let mut progress = false;

        for (index, instruction) in instructions.iter().enumerate() {
            if instruction.opcode != OpCode::CallA {
                continue;
            }
            let Some(arg_index) = calla_ldarg_index(instructions, index) else {
                continue;
            };
            let Some(method_offset) = largest_le(&method_starts, instruction.offset) else {
                continue;
            };
            let mut visited = BTreeSet::new();
            if let Some(start) = resolve_argument_target_for_method(
                instructions,
                callers_by_target,
                &method_starts,
                method_offset,
                arg_index,
                &mut visited,
            ) {
                let mut changed = starts.insert(start, ()).is_none();
                let callers = callers_by_target.entry(start).or_default();
                if callers.insert(index) {
                    changed = true;
                }
                if changed {
                    progress = true;
                }
            }
        }

        if !progress {
            break;
        }
    }
}

fn resolve_argument_target_for_method(
    instructions: &[Instruction],
    callers_by_target: &BTreeMap<usize, BTreeSet<usize>>,
    method_starts: &[usize],
    method_offset: usize,
    arg_index: u8,
    visited: &mut BTreeSet<(usize, u8)>,
) -> Option<usize> {
    if !visited.insert((method_offset, arg_index)) {
        return None;
    }

    let call_sites = callers_by_target.get(&method_offset)?;
    let callee_arg_count =
        initslot_arg_count_at(instructions, method_offset).unwrap_or(arg_index as usize + 1);

    for &call_index in call_sites {
        let call_offset = instructions.get(call_index)?.offset;
        match trace_call_arg_source(instructions, call_index, arg_index, callee_arg_count) {
            Some(CallArgSource::Target(target)) => return Some(target),
            Some(CallArgSource::PassThrough(next_arg)) => {
                let caller_method_offset =
                    largest_le(method_starts, call_offset).unwrap_or(call_offset);
                if let Some(target) = resolve_argument_target_for_method(
                    instructions,
                    callers_by_target,
                    method_starts,
                    caller_method_offset,
                    next_arg,
                    visited,
                ) {
                    return Some(target);
                }
            }
            None => {}
        }
    }

    None
}

/// Returns the largest element in `sorted` that is `<= target`.
/// `sorted` MUST be sorted ascending. O(log n).
fn largest_le(sorted: &[usize], target: usize) -> Option<usize> {
    let pos = sorted.partition_point(|&x| x <= target);
    if pos > 0 {
        Some(sorted[pos - 1])
    } else {
        None
    }
}

fn direct_call_target(instruction: &Instruction) -> Option<usize> {
    let delta = match instruction.operand {
        Some(Operand::Jump(value)) => value as isize,
        Some(Operand::Jump32(value)) => value as isize,
        _ => return None,
    };
    instruction.offset.checked_add_signed(delta)
}