neo-decompiler 0.11.0

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

use super::super::pointers::{pusha_absolute_target, resolve_slot_pointer_target};
use super::super::slots::{
    arg_load_index, find_slot_store_before, local_load_index, previous_non_nop_index,
    static_load_index, SlotDomain,
};

/// Check if a CALLA instruction ultimately loads its pointer from an argument
/// slot and return that argument index.
pub(in crate::decompiler::analysis) fn calla_ldarg_index(
    instructions: &[Instruction],
    calla_index: usize,
) -> Option<u8> {
    let producer_index = previous_non_nop_index(instructions, calla_index.checked_sub(1)?)?;
    trace_argument_index_from_value_source(instructions, producer_index)
}

/// Extract the argument count from an INITSLOT instruction at the given method offset.
pub(in crate::decompiler::analysis) fn initslot_arg_count_at(
    instructions: &[Instruction],
    method_offset: usize,
) -> Option<usize> {
    instructions
        .iter()
        .find(|i| i.offset == method_offset && i.opcode == OpCode::Initslot)
        .and_then(|i| match &i.operand {
            Some(Operand::Bytes(bytes)) if bytes.len() >= 2 => Some(bytes[1] as usize),
            _ => None,
        })
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(in crate::decompiler::analysis) enum CallArgSource {
    Target(usize),
    PassThrough(u8),
}

/// Trace backwards from a CALL instruction to find the source of the
/// `arg_index`-th argument (0-indexed).
///
/// Neo VM pops arguments top-first: the top of stack becomes arg0, the next
/// item becomes arg1, etc. So `arg0` is the last item pushed (0 items to skip)
/// and `arg N` requires skipping N single-push instructions.
pub(in crate::decompiler::analysis) fn trace_call_arg_source(
    instructions: &[Instruction],
    call_index: usize,
    arg_index: u8,
    callee_arg_count: usize,
) -> Option<CallArgSource> {
    if (arg_index as usize) >= callee_arg_count {
        return None;
    }
    let call_instruction = instructions.get(call_index)?;
    let skip_count = if call_instruction.opcode == OpCode::CallA {
        arg_index as usize + 1
    } else {
        arg_index as usize
    };

    let mut cursor = call_index.checked_sub(1)?;
    let mut remaining = skip_count;

    loop {
        let instr = instructions.get(cursor)?;
        if instr.opcode == OpCode::Nop {
            cursor = cursor.checked_sub(1)?;
            continue;
        }

        if remaining == 0 {
            return trace_argument_source_at(instructions, cursor);
        }

        remaining -= 1;
        cursor = cursor.checked_sub(1)?;
    }
}

fn trace_argument_source_at(instructions: &[Instruction], cursor: usize) -> Option<CallArgSource> {
    let instr = instructions.get(cursor)?;
    if instr.opcode == OpCode::PushA {
        return pusha_absolute_target(instr).map(CallArgSource::Target);
    }
    if let Some(slot) = local_load_index(instr) {
        return resolve_slot_pointer_target(instructions, cursor, SlotDomain::Local(slot))
            .map(CallArgSource::Target)
            .or_else(|| {
                trace_argument_index_from_value_source(instructions, cursor)
                    .map(CallArgSource::PassThrough)
            });
    }
    if let Some(slot) = static_load_index(instr) {
        return resolve_slot_pointer_target(instructions, cursor, SlotDomain::Static(slot))
            .map(CallArgSource::Target)
            .or_else(|| {
                trace_argument_index_from_value_source(instructions, cursor)
                    .map(CallArgSource::PassThrough)
            });
    }
    arg_load_index(instr).map(CallArgSource::PassThrough)
}

fn trace_argument_index_from_value_source(
    instructions: &[Instruction],
    mut source_index: usize,
) -> Option<u8> {
    loop {
        let instruction = instructions.get(source_index)?;
        if instruction.opcode == OpCode::Dup {
            source_index = previous_non_nop_index(instructions, source_index.checked_sub(1)?)?;
            continue;
        }
        if let Some(arg_index) = arg_load_index(instruction) {
            return Some(arg_index);
        }

        let domain = if let Some(slot) = local_load_index(instruction) {
            SlotDomain::Local(slot)
        } else {
            SlotDomain::Static(static_load_index(instruction)?)
        };

        let store_index = find_slot_store_before(instructions, source_index, domain)?;
        source_index = previous_non_nop_index(instructions, store_index.checked_sub(1)?)?;
    }
}