neo-decompiler 0.8.1

use std::collections::HashSet;
use std::fmt::Write;

use crate::instruction::Instruction;
use crate::manifest::{ContractManifest, ManifestMethod};

use super::super::super::helpers::{
    find_manifest_entry_method, next_inferred_method_offset, offset_as_usize,
};
use super::super::helpers::{
    collect_csharp_parameters, escape_csharp_string, format_csharp_parameters,
    format_manifest_type_csharp, format_method_signature, sanitize_csharp_identifier,
};
use super::body;

pub(super) struct MethodsContext<'a> {
    pub(super) instructions: &'a [Instruction],
    pub(super) inferred_method_starts: &'a [usize],
    pub(super) body_context: body::LiftedBodyContext<'a>,
}

pub(super) fn write_manifest_methods(
    output: &mut String,
    manifest: &ContractManifest,
    context: &MethodsContext<'_>,
    warnings: &mut Vec<String>,
) {
    let entry_method = write_script_entry_if_needed(output, manifest, context, warnings);
    let entry_offset = context
        .instructions
        .first()
        .map(|ins| ins.offset)
        .unwrap_or(0);

    let mut used_signatures: HashSet<(String, String)> = HashSet::new();
    let mut sorted_methods: Vec<&ManifestMethod> = manifest.abi.methods.iter().collect();
    sorted_methods.sort_by_key(|m| m.offset.unwrap_or(i32::MAX));

    let (with_offsets, without_offsets): (Vec<_>, Vec<_>) =
        sorted_methods.into_iter().partition(|m| m.offset.is_some());

    for method in with_offsets.iter() {
        let start = offset_as_usize(method.offset).unwrap_or(0);
        let end = next_inferred_method_offset(context.inferred_method_starts, start)
            .or_else(|| context.instructions.last().map(|i| i.offset + 1))
            .unwrap_or(start);
        let slice: Vec<Instruction> = context
            .instructions
            .iter()
            .filter(|ins| ins.offset >= start && ins.offset < end)
            .cloned()
            .collect();

        let params = collect_csharp_parameters(&method.parameters);
        let signature_key = params
            .iter()
            .map(|param| param.ty.as_str())
            .collect::<Vec<_>>()
            .join(",");
        let param_signature = format_csharp_parameters(&params);
        let base_name = sanitize_csharp_identifier(&method.name);
        let method_name = make_unique_method_name(base_name, &signature_key, &mut used_signatures);
        let return_type = format_manifest_type_csharp(&method.return_type, true);
        let signature = format_method_signature(&method_name, &param_signature, &return_type);

        write_method_attributes(output, &method_name, &method.name, method.safe);
        writeln!(output, "        {signature}").unwrap();
        writeln!(output, "        {{").unwrap();

        if slice.is_empty() {
            writeln!(output, "            throw new NotImplementedException();").unwrap();
        } else {
            let labels: Vec<String> = params.iter().map(|p| p.name.clone()).collect();
            let is_void = return_type == "void";
            body::write_lifted_body(
                output,
                &slice,
                Some(&labels),
                warnings,
                &context.body_context,
                is_void,
            );
        }

        writeln!(output, "        }}").unwrap();
        writeln!(output).unwrap();
    }

    for method in without_offsets {
        let params = collect_csharp_parameters(&method.parameters);
        let signature_key = params
            .iter()
            .map(|param| param.ty.as_str())
            .collect::<Vec<_>>()
            .join(",");
        let param_signature = format_csharp_parameters(&params);
        let base_name = sanitize_csharp_identifier(&method.name);
        let method_name = make_unique_method_name(base_name, &signature_key, &mut used_signatures);
        let return_type = format_manifest_type_csharp(&method.return_type, true);
        let signature = format_method_signature(&method_name, &param_signature, &return_type);

        write_method_attributes(output, &method_name, &method.name, method.safe);
        writeln!(output, "        {signature}").unwrap();
        writeln!(output, "        {{").unwrap();

        if entry_method
            .as_ref()
            .map(|(entry, _)| std::ptr::eq(*entry, method))
            .unwrap_or(false)
        {
            let end = next_inferred_method_offset(context.inferred_method_starts, entry_offset);
            let slice: Vec<Instruction> = match end {
                Some(end) => context
                    .instructions
                    .iter()
                    .filter(|ins| ins.offset >= entry_offset && ins.offset < end)
                    .cloned()
                    .collect(),
                None => context.instructions.to_vec(),
            };
            let slice = if slice.is_empty() {
                context.instructions.to_vec()
            } else {
                slice
            };
            let labels: Vec<String> = params.iter().map(|p| p.name.clone()).collect();
            let is_void = return_type == "void";
            body::write_lifted_body(
                output,
                &slice,
                Some(&labels),
                warnings,
                &context.body_context,
                is_void,
            );
        } else {
            writeln!(output, "            throw new NotImplementedException();").unwrap();
        }

        writeln!(output, "        }}").unwrap();
        writeln!(output).unwrap();
    }
}

pub(super) fn write_inferred_methods(
    output: &mut String,
    context: &MethodsContext<'_>,
    manifest: Option<&ContractManifest>,
    warnings: &mut Vec<String>,
) {
    let entry_offset = context.instructions.first().map(|ins| ins.offset);
    let manifest_offsets: HashSet<usize> = manifest
        .map(|manifest| {
            manifest
                .abi
                .methods
                .iter()
                .filter_map(|method| offset_as_usize(method.offset))
                .collect()
        })
        .unwrap_or_default();

    for start in context.inferred_method_starts {
        if Some(*start) == entry_offset || manifest_offsets.contains(start) {
            continue;
        }

        let end = next_inferred_method_offset(context.inferred_method_starts, *start)
            .or_else(|| context.instructions.last().map(|ins| ins.offset + 1))
            .unwrap_or(*start);
        let slice: Vec<Instruction> = context
            .instructions
            .iter()
            .filter(|ins| ins.offset >= *start && ins.offset < end)
            .cloned()
            .collect();
        if slice.is_empty()
            || slice
                .iter()
                .all(|ins| ins.opcode == crate::instruction::OpCode::Nop)
        {
            continue;
        }

        let arg_count = context
            .body_context
            .method_arg_counts_by_offset
            .get(start)
            .copied()
            .unwrap_or(0);
        let params = (0..arg_count)
            .map(|index| format!("dynamic arg{index}"))
            .collect::<Vec<_>>()
            .join(", ");

        // Each inferred helper is preceded by the trailing blank line
        // emitted by the previous method (the synthetic ScriptEntry from
        // `write_fallback_entry`, the last manifest method from
        // `write_manifest_methods`, or the previous iteration of this
        // loop). Emitting another blank line here would double-space the
        // separator.
        writeln!(
            output,
            "        private static dynamic sub_0x{start:04X}({params})"
        )
        .unwrap();
        writeln!(output, "        {{").unwrap();
        let labels = (0..arg_count)
            .map(|index| format!("arg{index}"))
            .collect::<Vec<_>>();
        body::write_lifted_body(
            output,
            &slice,
            (!labels.is_empty()).then_some(labels.as_slice()),
            warnings,
            &context.body_context,
            false,
        );
        writeln!(output, "        }}").unwrap();
        writeln!(output).unwrap();
    }
}

fn write_script_entry_if_needed<'a>(
    output: &mut String,
    manifest: &'a ContractManifest,
    context: &MethodsContext<'_>,
    warnings: &mut Vec<String>,
) -> Option<(&'a ManifestMethod, bool)> {
    let entry_offset = context.instructions.first().map(|ins| ins.offset)?;

    let entry_method = find_manifest_entry_method(manifest, entry_offset);
    if entry_method.is_some() {
        return entry_method;
    }

    let end = next_inferred_method_offset(context.inferred_method_starts, entry_offset);
    let slice: Vec<Instruction> = match end {
        Some(end) => context
            .instructions
            .iter()
            .filter(|ins| ins.offset >= entry_offset && ins.offset < end)
            .cloned()
            .collect(),
        None => context.instructions.to_vec(),
    };

    let slice = if slice.is_empty() {
        context.instructions.to_vec()
    } else {
        slice
    };

    writeln!(
        output,
        "        // warning: manifest entry offset did not match script entry at 0x{entry_offset:04X}; using synthetic ScriptEntry"
    )
    .unwrap();
    // Without a matching manifest entry the return type isn't known.
    // Default to `object` and let the high-level lift preserve any
    // return value (a hardcoded `void` here would make the emitter
    // discard whatever the bytecode's RET was returning).
    let (parameters, argument_labels) = synthetic_entry_arguments(&slice, entry_offset);
    let entry_signature = format_method_signature("ScriptEntry", &parameters, "object");
    writeln!(output, "        {entry_signature}").unwrap();
    writeln!(output, "        {{").unwrap();
    body::write_lifted_body(
        output,
        &slice,
        (!argument_labels.is_empty()).then_some(argument_labels.as_slice()),
        warnings,
        &context.body_context,
        false,
    );
    writeln!(output, "        }}").unwrap();
    writeln!(output).unwrap();
    None
}

/// Build the `object arg0, object arg1, ...` parameter list and the
/// matching `arg0`, `arg1`, ... label list for a synthetic ScriptEntry.
/// `slice` is the bytecode slice covering the entry method, used to
/// recover the INITSLOT-declared argument count. Returns `(parameters,
/// labels)` where `labels` is empty when the method takes no
/// arguments (the caller can short-circuit `argument_labels` with
/// `then_some(...)`).
fn synthetic_entry_arguments(slice: &[Instruction], entry_offset: usize) -> (String, Vec<String>) {
    let arg_count =
        crate::decompiler::helpers::initslot_argument_count_at(slice, entry_offset).unwrap_or(0);
    let parameters = (0..arg_count)
        .map(|i| format!("object arg{i}"))
        .collect::<Vec<_>>()
        .join(", ");
    let labels: Vec<String> = (0..arg_count).map(|i| format!("arg{i}")).collect();
    (parameters, labels)
}

pub(super) fn write_fallback_entry(
    output: &mut String,
    context: &MethodsContext<'_>,
    warnings: &mut Vec<String>,
) {
    let entry_offset = context
        .instructions
        .first()
        .map(|ins| ins.offset)
        .unwrap_or(0);
    let end = next_inferred_method_offset(context.inferred_method_starts, entry_offset)
        .or_else(|| context.instructions.last().map(|i| i.offset + 1))
        .unwrap_or(entry_offset);
    let slice: Vec<Instruction> = context
        .instructions
        .iter()
        .filter(|ins| ins.offset >= entry_offset && ins.offset < end)
        .cloned()
        .collect();
    let slice = if slice.is_empty() {
        context.instructions.to_vec()
    } else {
        slice
    };

    let entry_method_name = "ScriptEntry".to_string();
    // No manifest is available, so the return type isn't known.
    // Default to `object` (widest type) and let the high-level
    // lift preserve any return value — a hardcoded `void` here
    // would make the emitter discard whatever the bytecode's
    // RET was returning.
    let (parameters, argument_labels) = synthetic_entry_arguments(&slice, entry_offset);
    let entry_signature = format_method_signature(&entry_method_name, &parameters, "object");
    writeln!(output, "        {entry_signature}").unwrap();
    writeln!(output, "        {{").unwrap();
    body::write_lifted_body(
        output,
        &slice,
        (!argument_labels.is_empty()).then_some(argument_labels.as_slice()),
        warnings,
        &context.body_context,
        false,
    );
    writeln!(output, "        }}").unwrap();
    // Trailing blank line for consistency with manifest-driven
    // method emission (each manifest method ends with one); without
    // this the close-brace of a synthetic ScriptEntry sits flush
    // against the class close-brace.
    writeln!(output).unwrap();
}

fn write_method_attributes(output: &mut String, method_name: &str, raw_name: &str, is_safe: bool) {
    if method_name != raw_name {
        writeln!(
            output,
            "        [DisplayName(\"{}\")]",
            escape_csharp_string(raw_name)
        )
        .unwrap();
    }
    if is_safe {
        writeln!(output, "        [Safe]").unwrap();
    }
}

fn make_unique_method_name(
    base: String,
    signature: &str,
    used: &mut HashSet<(String, String)>,
) -> String {
    let mut candidate = base.clone();
    let mut index = 1usize;
    while !used.insert((candidate.clone(), signature.to_string())) {
        candidate = format!("{base}_{index}");
        index += 1;
    }
    candidate
}