neo-decompiler 0.11.0

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

use crate::instruction::Instruction;

use super::super::emitter::HighLevelEmitter;

pub(super) struct MethodBodyContext<'a> {
    pub(super) method_labels_by_offset: &'a BTreeMap<usize, String>,
    pub(super) method_arg_counts_by_offset: &'a BTreeMap<usize, usize>,
    pub(super) call_targets_by_offset: &'a BTreeMap<usize, usize>,
    pub(super) calla_targets_by_offset: &'a BTreeMap<usize, usize>,
    pub(super) noreturn_method_offsets: &'a BTreeSet<usize>,
    pub(super) inline_single_use_temps: bool,
    pub(super) emit_trace_comments: bool,
    pub(super) typed_declarations: bool,
    pub(super) slot_types_by_offset: &'a BTreeMap<usize, SlotTypes>,
    pub(super) callt_labels: &'a [String],
    pub(super) callt_param_counts: &'a [usize],
    pub(super) callt_returns_value: &'a [bool],
}

#[derive(Debug, Clone, Default)]
pub(super) struct SlotTypes {
    pub(super) arguments: Vec<&'static str>,
    pub(super) locals: Vec<&'static str>,
    pub(super) statics: Vec<&'static str>,
}

impl SlotTypes {
    pub(super) fn argument_type(&self, index: usize) -> Option<&'static str> {
        self.arguments
            .get(index)
            .copied()
            .filter(|ty| !ty.is_empty())
    }

    fn declaration_type(&self, name: &str) -> Option<&'static str> {
        let (slots, idx) = if let Some(rest) = name.strip_prefix("loc") {
            (&self.locals, parse_slot_index(rest)?)
        } else if let Some(rest) = name.strip_prefix("static") {
            (&self.statics, parse_slot_index(rest)?)
        } else {
            return None;
        };
        slots.get(idx).copied().filter(|ty| !ty.is_empty())
    }
}

pub(super) fn write_method_body(
    output: &mut String,
    instructions: &[Instruction],
    argument_labels: Option<&[String]>,
    warnings: &mut Vec<String>,
    context: &MethodBodyContext<'_>,
    returns_void: bool,
) {
    if instructions.len() > super::super::emitter::MAX_HIGH_LEVEL_METHOD_INSTRUCTIONS {
        let offset = instructions.first().map(|i| i.offset).unwrap_or(0);
        writeln!(
            output,
            "        // method body too large for high-level lifting: {} instructions \
             exceeds the {}-instruction limit; use `disasm` for the full listing",
            instructions.len(),
            super::super::emitter::MAX_HIGH_LEVEL_METHOD_INSTRUCTIONS
        )
        .unwrap();
        warnings.push(format!(
            "high-level: method at 0x{offset:04X} skipped — {} instructions exceeds the \
             high-level lifting limit ({})",
            instructions.len(),
            super::super::emitter::MAX_HIGH_LEVEL_METHOD_INSTRUCTIONS
        ));
        return;
    }
    let mut emitter = HighLevelEmitter::with_program(instructions);
    if let Some(labels) = argument_labels {
        emitter.set_argument_labels(labels);
    }
    emitter.set_inline_single_use_temps(context.inline_single_use_temps);
    emitter.set_emit_trace_comments(context.emit_trace_comments);
    emitter.set_callt_labels(context.callt_labels.to_vec());
    emitter.set_callt_param_counts(context.callt_param_counts.to_vec());
    emitter.set_callt_returns_value(context.callt_returns_value.to_vec());
    emitter.set_method_labels_by_offset(context.method_labels_by_offset);
    emitter.set_method_arg_counts_by_offset(context.method_arg_counts_by_offset);
    emitter.set_call_targets_by_offset(context.call_targets_by_offset);
    emitter.set_calla_targets_by_offset(context.calla_targets_by_offset);
    emitter.set_noreturn_method_offsets(context.noreturn_method_offsets);
    emitter.set_returns_void(returns_void);
    for instruction in instructions {
        emitter.advance_to(instruction.offset);
        emitter.emit_instruction(instruction);
    }
    let result = emitter.finish();
    warnings.extend(result.warnings);
    let statements = result.statements;

    if statements.is_empty() {
        writeln!(output, "        // no instructions decoded").unwrap();
        return;
    }

    let method_start = instructions.first().map(|i| i.offset).unwrap_or(0);
    let slot_types = if context.typed_declarations {
        context
            .slot_types_by_offset
            .get(&method_start)
            .cloned()
            .unwrap_or_default()
    } else {
        SlotTypes::default()
    };

    let mut indent_level = 0usize;
    for line in statements {
        let annotated;
        let trimmed = if context.typed_declarations {
            annotated = annotate_statement_types(line.trim(), &slot_types);
            annotated.trim()
        } else {
            line.trim()
        };
        if trimmed.is_empty() {
            continue;
        }

        if trimmed.starts_with('}') {
            indent_level = indent_level.saturating_sub(1);
        }

        let indent = 8 + indent_level * 4;
        writeln!(output, "{:indent$}{}", "", trimmed, indent = indent).unwrap();

        if trimmed.ends_with('{') {
            indent_level += 1;
        }
    }
}

fn annotate_statement_types(line: &str, types: &SlotTypes) -> String {
    if let Some(rest) = line.strip_prefix("let ") {
        return annotate_let(rest, types)
            .map(|annotated| annotated.to_string())
            .unwrap_or_else(|| line.to_string());
    }
    if let Some(rest) = line.strip_prefix("for (let ") {
        if let Some(annotated) = annotate_let(rest, types) {
            return format!("for ({annotated}");
        }
    }
    line.to_string()
}

fn annotate_let(rest: &str, types: &SlotTypes) -> Option<String> {
    let name = rest
        .split(|c: char| c.is_whitespace() || c == '=')
        .next()
        .unwrap_or("");
    let ty = types.declaration_type(name)?;
    Some(format!("{ty} {rest}"))
}

fn parse_slot_index(rest: &str) -> Option<usize> {
    if rest.is_empty() || !rest.bytes().all(|b| b.is_ascii_digit()) {
        return None;
    }
    rest.parse::<usize>().ok()
}