neo-decompiler 0.10.1

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

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum SlotDomain {
    Local(u8),
    Static(u8),
}

pub(super) fn local_load_index(instruction: &Instruction) -> Option<u8> {
    match instruction.opcode {
        OpCode::Ldloc0 => Some(0),
        OpCode::Ldloc1 => Some(1),
        OpCode::Ldloc2 => Some(2),
        OpCode::Ldloc3 => Some(3),
        OpCode::Ldloc4 => Some(4),
        OpCode::Ldloc5 => Some(5),
        OpCode::Ldloc6 => Some(6),
        OpCode::Ldloc => match instruction.operand {
            Some(Operand::U8(index)) => Some(index),
            _ => None,
        },
        _ => None,
    }
}

pub(super) fn static_load_index(instruction: &Instruction) -> Option<u8> {
    match instruction.opcode {
        OpCode::Ldsfld0 => Some(0),
        OpCode::Ldsfld1 => Some(1),
        OpCode::Ldsfld2 => Some(2),
        OpCode::Ldsfld3 => Some(3),
        OpCode::Ldsfld4 => Some(4),
        OpCode::Ldsfld5 => Some(5),
        OpCode::Ldsfld6 => Some(6),
        OpCode::Ldsfld => match instruction.operand {
            Some(Operand::U8(index)) => Some(index),
            _ => None,
        },
        _ => None,
    }
}

pub(super) fn arg_load_index(instruction: &Instruction) -> Option<u8> {
    match instruction.opcode {
        OpCode::Ldarg0 => Some(0),
        OpCode::Ldarg1 => Some(1),
        OpCode::Ldarg2 => Some(2),
        OpCode::Ldarg3 => Some(3),
        OpCode::Ldarg4 => Some(4),
        OpCode::Ldarg5 => Some(5),
        OpCode::Ldarg6 => Some(6),
        OpCode::Ldarg => match instruction.operand {
            Some(Operand::U8(index)) => Some(index),
            _ => None,
        },
        _ => None,
    }
}

fn slot_store_domain(instruction: &Instruction) -> Option<SlotDomain> {
    match instruction.opcode {
        OpCode::Stloc0 => Some(SlotDomain::Local(0)),
        OpCode::Stloc1 => Some(SlotDomain::Local(1)),
        OpCode::Stloc2 => Some(SlotDomain::Local(2)),
        OpCode::Stloc3 => Some(SlotDomain::Local(3)),
        OpCode::Stloc4 => Some(SlotDomain::Local(4)),
        OpCode::Stloc5 => Some(SlotDomain::Local(5)),
        OpCode::Stloc6 => Some(SlotDomain::Local(6)),
        OpCode::Stloc => match instruction.operand {
            Some(Operand::U8(index)) => Some(SlotDomain::Local(index)),
            _ => None,
        },
        OpCode::Stsfld0 => Some(SlotDomain::Static(0)),
        OpCode::Stsfld1 => Some(SlotDomain::Static(1)),
        OpCode::Stsfld2 => Some(SlotDomain::Static(2)),
        OpCode::Stsfld3 => Some(SlotDomain::Static(3)),
        OpCode::Stsfld4 => Some(SlotDomain::Static(4)),
        OpCode::Stsfld5 => Some(SlotDomain::Static(5)),
        OpCode::Stsfld6 => Some(SlotDomain::Static(6)),
        OpCode::Stsfld => match instruction.operand {
            Some(Operand::U8(index)) => Some(SlotDomain::Static(index)),
            _ => None,
        },
        _ => None,
    }
}

pub(super) fn trace_stack_value_producer_before(
    instructions: &[Instruction],
    before_index: usize,
    mut depth: usize,
) -> Option<usize> {
    for index in (0..before_index).rev() {
        let instruction = instructions.get(index)?;
        let (pops, pushes) = stack_effect(instruction)?;
        if depth < pushes {
            return Some(index);
        }
        depth = depth.checked_add(pops)?.checked_sub(pushes)?;
    }
    None
}

fn stack_effect(instruction: &Instruction) -> Option<(usize, usize)> {
    use OpCode::*;
    let opcode = instruction.opcode;
    match opcode {
        Nop => Some((0, 0)),
        PushA | PushNull | PushT | PushF | PushM1 | Push0 | Push1 | Push2 | Push3 | Push4
        | Push5 | Push6 | Push7 | Push8 | Push9 | Push10 | Push11 | Push12 | Push13 | Push14
        | Push15 | Push16 | Pushint8 | Pushint16 | Pushint32 | Pushint64 | Pushint128
        | Pushint256 | Pushdata1 | Pushdata2 | Pushdata4 | Newarray0 | Newmap | Newstruct0
        | Ldloc0 | Ldloc1 | Ldloc2 | Ldloc3 | Ldloc4 | Ldloc5 | Ldloc6 | Ldloc | Ldarg0
        | Ldarg1 | Ldarg2 | Ldarg3 | Ldarg4 | Ldarg5 | Ldarg6 | Ldarg | Ldsfld0 | Ldsfld1
        | Ldsfld2 | Ldsfld3 | Ldsfld4 | Ldsfld5 | Ldsfld6 | Ldsfld => Some((0, 1)),
        Stloc0 | Stloc1 | Stloc2 | Stloc3 | Stloc4 | Stloc5 | Stloc6 | Stloc | Starg0 | Starg1
        | Starg2 | Starg3 | Starg4 | Starg5 | Starg6 | Starg | Stsfld0 | Stsfld1 | Stsfld2
        | Stsfld3 | Stsfld4 | Stsfld5 | Stsfld6 | Stsfld => Some((1, 0)),
        Append => Some((2, 0)),
        Pickitem => Some((2, 1)),
        Dup => Some((1, 2)),
        _ => None,
    }
}

pub(super) fn find_resolution_start_index(
    instructions: &[Instruction],
    before_index: usize,
) -> usize {
    for index in (0..before_index).rev() {
        if let Some(instruction) = instructions.get(index) {
            if is_pointer_resolution_boundary(instruction.opcode) {
                return index + 1;
            }
        }
    }
    0
}

pub(super) fn find_slot_store_before(
    instructions: &[Instruction],
    before_index: usize,
    domain: SlotDomain,
) -> Option<usize> {
    for index in (0..before_index).rev() {
        let instruction = instructions.get(index)?;
        if matches!(domain, SlotDomain::Local(_))
            && is_pointer_resolution_boundary(instruction.opcode)
        {
            return None;
        }
        if slot_store_domain(instruction) == Some(domain) {
            return Some(index);
        }
    }
    None
}

fn is_pointer_resolution_boundary(opcode: OpCode) -> bool {
    matches!(
        opcode,
        OpCode::Ret
            | OpCode::Throw
            | OpCode::Abort
            | OpCode::Abortmsg
            | OpCode::Initslot
            | OpCode::Initsslot
    )
}

pub(super) fn previous_non_nop_index(
    instructions: &[Instruction],
    mut index: usize,
) -> Option<usize> {
    loop {
        let instruction = instructions.get(index)?;
        if instruction.opcode != OpCode::Nop {
            return Some(index);
        }
        index = index.checked_sub(1)?;
    }
}