neo-decompiler 0.10.1

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

use crate::instruction::Instruction;

use super::super::super::{MethodRef, MethodTable};
use super::super::{CallEdge, CallTarget};
use super::source::{
    calla_ldarg_index, initslot_arg_count_at, trace_call_arg_source, CallArgSource,
};

/// Second pass over call edges: resolve CALLA targets that load their function
/// pointer from an argument slot (LDARG N) by tracing back through callers.
pub(in crate::decompiler::analysis::call_graph) fn resolve_ldarg_calla_targets(
    instructions: &[Instruction],
    edges: &mut [CallEdge],
    table: &MethodTable,
    methods: &mut BTreeMap<usize, MethodRef>,
) {
    let offset_to_index: BTreeMap<usize, usize> = instructions
        .iter()
        .enumerate()
        .map(|(i, instr)| (instr.offset, i))
        .collect();

    let sites = collect_ldarg_calla_sites(instructions, edges, methods, &offset_to_index);
    if sites.is_empty() {
        return;
    }

    loop {
        let callers_by_target = collect_callers_by_target(edges);
        let mut progress = false;
        for (edge_idx, arg_idx, method_offset) in &sites {
            if !matches!(edges[*edge_idx].target, CallTarget::Indirect { .. }) {
                continue;
            }

            let mut visited = BTreeSet::new();
            let resolved = resolve_argument_target_recursive(
                instructions,
                &offset_to_index,
                &callers_by_target,
                methods,
                *method_offset,
                *arg_idx,
                &mut visited,
            );

            if let Some(target) = resolved.filter(|target| offset_to_index.contains_key(target)) {
                let callee = table.resolve_internal_target(target);
                methods.insert(callee.offset, callee.clone());
                edges[*edge_idx].target = CallTarget::Internal { method: callee };
                progress = true;
            }
        }

        if !progress {
            break;
        }
    }
}

fn collect_ldarg_calla_sites(
    instructions: &[Instruction],
    edges: &[CallEdge],
    methods: &BTreeMap<usize, MethodRef>,
    offset_to_index: &BTreeMap<usize, usize>,
) -> Vec<(usize, u8, usize)> {
    // Collect unresolved CALLA sites preceded by LDARG. `edge.caller.offset`
    // can be stale for helpers discovered during the first pass, so use the
    // methods map to find the true containing method for each CALLA.
    let mut sites = Vec::new();
    for (edge_idx, edge) in edges.iter().enumerate() {
        if edge.opcode != "CALLA" || !matches!(edge.target, CallTarget::Indirect { .. }) {
            continue;
        }
        let Some(&calla_idx) = offset_to_index.get(&edge.call_offset) else {
            continue;
        };
        if let Some(arg_idx) = calla_ldarg_index(instructions, calla_idx) {
            let actual_method_offset = methods
                .range(..=edge.call_offset)
                .next_back()
                .map(|(&offset, _)| offset)
                .unwrap_or(edge.caller.offset);
            sites.push((edge_idx, arg_idx, actual_method_offset));
        }
    }
    sites
}

fn collect_callers_by_target(edges: &[CallEdge]) -> BTreeMap<usize, Vec<usize>> {
    let mut callers_by_target: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
    for edge in edges {
        if let CallTarget::Internal { method } = &edge.target {
            if edge.opcode == "CALL" || edge.opcode == "CALL_L" || edge.opcode == "CALLA" {
                callers_by_target
                    .entry(method.offset)
                    .or_default()
                    .push(edge.call_offset);
            }
        }
    }
    callers_by_target
}

fn resolve_argument_target_recursive(
    instructions: &[Instruction],
    offset_to_index: &BTreeMap<usize, usize>,
    callers_by_target: &BTreeMap<usize, Vec<usize>>,
    methods: &BTreeMap<usize, MethodRef>,
    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_offset in call_sites {
        let &call_idx = offset_to_index.get(&call_offset)?;
        match trace_call_arg_source(instructions, call_idx, arg_index, callee_arg_count) {
            Some(CallArgSource::Target(target)) => return Some(target),
            Some(CallArgSource::PassThrough(next_arg)) => {
                let caller_method_offset = methods
                    .range(..=call_offset)
                    .next_back()
                    .map(|(&offset, _)| offset)
                    .unwrap_or(call_offset);
                if let Some(target) = resolve_argument_target_recursive(
                    instructions,
                    offset_to_index,
                    callers_by_target,
                    methods,
                    caller_method_offset,
                    next_arg,
                    visited,
                ) {
                    return Some(target);
                }
            }
            None => {}
        }
    }

    None
}