spimdisasm 2.0.0-alpha.1

/* SPDX-FileCopyrightText: © 2024-2025 Decompollaborate */
/* SPDX-License-Identifier: MIT */

use core::fmt;

use rabbitizer::{Instruction, InstructionDisplayFlags};

#[cfg(feature = "pyo3")]
use pyo3::prelude::*;

use crate::{
    addresses::{Size, Vram},
    collections::addended_ordered_map::FindSettings,
    context::Context,
    metadata::{GeneratedBy, LabelType, SegmentMetadata, SymbolMetadata},
    relocation::RelocationInfo,
    symbols::{
        display::sym_common_display::WordComment, processed::FunctionSymProcessed,
        trait_symbol::RomSymbol, RomSymbolProcessed, Symbol,
    },
};

use super::{
    sym_display_error::SymDisplayError, InternalSymDisplSettings, SymCommonDisplaySettings,
};

#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
#[cfg_attr(feature = "pyo3", pyclass(module = "spimdisasm"))]
pub struct FunctionDisplaySettings {
    common: SymCommonDisplaySettings,

    display_flags: InstructionDisplayFlags,

    asm_label_indentation: u8,

    _gp_rel_hack: bool,
}

impl FunctionDisplaySettings {
    pub fn new(display_flags: InstructionDisplayFlags) -> Self {
        Self {
            common: SymCommonDisplaySettings::new(),
            display_flags,
            asm_label_indentation: 2,
            _gp_rel_hack: false,
        }
    }

    pub fn set_rom_comment_width(&mut self, rom_comment_width: u8) {
        self.common.set_rom_comment_width(rom_comment_width);
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
#[must_use]
pub struct FunctionDisplay<'ctx, 'sym, 'flg> {
    context: &'ctx Context,
    sym: &'sym FunctionSymProcessed,
    settings: &'flg FunctionDisplaySettings,

    owned_segment: &'ctx SegmentMetadata,
    metadata: &'ctx SymbolMetadata,

    internal_settings: InternalSymDisplSettings,
}

impl<'ctx, 'sym, 'flg> FunctionDisplay<'ctx, 'sym, 'flg> {
    pub(crate) fn new(
        context: &'ctx Context,
        sym: &'sym FunctionSymProcessed,
        settings: &'flg FunctionDisplaySettings,
        internal_settings: InternalSymDisplSettings,
    ) -> Result<Self, SymDisplayError> {
        let owned_segment = context.find_owned_segment(sym.parent_segment_info())?;
        let find_settings = FindSettings::new(false);
        let metadata = owned_segment
            .find_symbol(sym.vram_range().start(), find_settings)
            .ok_or(SymDisplayError::SelfSymNotFound())?;

        Ok(Self {
            context,
            sym,
            settings,
            owned_segment,
            metadata,
            internal_settings,
        })
    }

    #[must_use]
    pub(crate) fn sym(&self) -> &'sym FunctionSymProcessed {
        self.sym
    }

    #[must_use]
    pub(crate) fn settings_common(&self) -> &'flg SymCommonDisplaySettings {
        &self.settings.common
    }
}

impl FunctionDisplay<'_, '_, '_> {
    fn display_label(&self, f: &mut fmt::Formatter<'_>, current_vram: Vram) -> fmt::Result {
        if let Some(label) = self.owned_segment.find_label(current_vram) {
            // Only emit a label if it is referenced
            if label.reference_counter() == 0 && label.generated_by() == GeneratedBy::Autogenerated
            {
                return Ok(());
            }

            if self.settings.asm_label_indentation > 0 {
                write!(
                    f,
                    "{:width$}",
                    " ",
                    width = self.settings.asm_label_indentation as usize
                )?;
            }

            let use_macro = match label.label_type() {
                LabelType::Branch => false,
                LabelType::Jumptable => !self.internal_settings.migrate(),
                LabelType::GccExceptTable => true,
                LabelType::AlternativeEntry => true,
            };

            // TODO:
            /*
            if not use_macro:
                if common.GlobalConfig.ASM_GLOBALIZE_TEXT_LABELS_REFERENCED_BY_NON_JUMPTABLE:
                    # Check if any non-jumptable symbol references this label
                    for otherSym in labelSym.referenceSymbols:
                        if otherSym.getTypeSpecial() != common.SymbolSpecialType.jumptable:
                            use_macro = True
                            break
            */

            if use_macro {
                // label = labelSym.getreferenceSymbols()

                self.settings
                    .common
                    .display_label(f, self.context.global_config(), label)?;
            } else {
                write!(
                    f,
                    "{}:{}",
                    label.display_name(),
                    self.settings.common.line_end()
                )?;
            }
        }

        Ok(())
    }

    fn display_instruction(
        &self,
        f: &mut fmt::Formatter<'_>,
        instr: &Instruction,
        prev_instr_had_delay_slot: bool,
    ) -> fmt::Result {
        let arr_bytes = self
            .context
            .global_config()
            .endian()
            .bytes_from_word(instr.word());
        let vram = instr.vram();
        let rom = self.sym.rom_vram_range().rom_from_vram(vram);
        self.settings
            .common
            .display_asm_comment(f, rom, vram, WordComment::U32(arr_bytes))?;

        // TODO: why an extra space?
        write!(f, " ")?;

        let extra_ljust = if prev_instr_had_delay_slot {
            write!(f, " ")?;
            -1
        } else {
            0
        };

        let imm_override = self.get_reloc(instr).and_then(|x| {
            x.display(
                self.context,
                self.sym.parent_segment_info(),
                true,
                self.metadata.compiler(),
                self.internal_settings,
            )
        });

        let instr_display = instr.display(&self.settings.display_flags, imm_override, extra_ljust);

        #[cfg(feature = "pyo3")]
        let instr_display = {
            use alloc::string::ToString;

            // TODO: hack for splat diffs
            let mut temp = instr_display.to_string().replace("$s8", "$fp");
            if !self.settings.display_flags.named_gpr() {
                temp = temp.replace("$zero", "$0").replace("$ra", "$31");
            }
            temp
        };

        write!(f, "{instr_display}")
    }

    fn get_reloc(&self, instr: &Instruction) -> Option<&RelocationInfo> {
        let index = (instr.vram() - self.sym.vram_range().start()).inner() / 4;
        self.sym.relocs()[index as usize]
            .as_ref()
            .filter(|x| !x.reloc_type().is_none())
    }

    fn display_end_of_line_comment(
        &self,
        f: &mut fmt::Formatter<'_>,
        instr: &Instruction,
    ) -> fmt::Result {
        let vram = instr.vram();
        let rom = self.sym.rom_vram_range().rom_from_vram(vram).unwrap();

        if self.sym.handwritten_instrs().contains(&rom) {
            write!(f, " /* handwritten instruction */")?;
        }

        Ok(())
    }
}

impl fmt::Display for FunctionDisplay<'_, '_, '_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if !self.sym.handwritten_instrs().is_empty() {
            write!(
                f,
                "/* Handwritten function */{}",
                self.settings.common.line_end()
            )?;
        }

        self.settings
            .common
            .display_sym_property_comments(f, self.metadata, self.owned_segment)?;
        self.settings
            .common
            .display_sym_prev_alignment(f, self.metadata)?;
        self.settings.common.display_symbol_name(
            f,
            self.context.global_config(),
            self.metadata,
            false,
            self.metadata.section_type(),
        )?;

        let mut size = Size::new(0);
        let symbol_size = if let Some(s) = self.metadata.size() {
            let new_size = if let Some(padding) = self.metadata.trailing_padding_size() {
                Size::new(s.inner().saturating_sub(padding.inner()))
            } else {
                s
            };

            Some(new_size)
        } else {
            None
        };

        let mut prev_instr_had_delay_slot = false;
        for instr in self.sym.instructions() {
            let current_vram = instr.vram();
            self.display_label(f, current_vram)?;
            self.display_instruction(f, instr, prev_instr_had_delay_slot)?;

            self.display_end_of_line_comment(f, instr)?;
            write!(f, "{}", self.settings.common.line_end())?;

            prev_instr_had_delay_slot = instr.opcode().has_delay_slot();

            size += Size::new(4);
            if Some(size) == symbol_size {
                self.settings.common.display_sym_end(
                    f,
                    self.context.global_config(),
                    self.metadata,
                )?;
            }
        }

        Ok(())
    }
}

#[cfg(feature = "pyo3")]
pub(crate) mod python_bindings {
    use super::*;

    #[pymethods]
    impl FunctionDisplaySettings {
        #[new]
        pub fn py_new(display_flags: InstructionDisplayFlags) -> Self {
            Self::new(display_flags)
        }

        #[pyo3(name = "set_rom_comment_width")]
        pub fn py_set_rom_comment_width(&mut self, rom_comment_width: u8) {
            self.set_rom_comment_width(rom_comment_width);
        }
    }
}