papermario/tools/splat/util/symbols.py

from dataclasses import dataclass
from typing import Dict, List, Optional, Set, TYPE_CHECKING

import spimdisasm
import tqdm
from intervaltree import Interval, IntervalTree

# circular import
if TYPE_CHECKING:
    from segtypes.segment import Segment

from util import log, options

all_symbols: List["Symbol"] = []
all_symbols_dict: Dict[int, List["Symbol"]] = {}
all_symbols_ranges = IntervalTree()
ignored_addresses: Set[int] = set()
to_mark_as_defined: Set[str] = set()

# Initialize a spimdisasm context, used to store symbols and functions
spim_context = spimdisasm.common.Context()

TRUEY_VALS = ["true", "on", "yes", "y"]
FALSEY_VALS = ["false", "off", "no", "n"]


def is_truey(str: str) -> bool:
    return str.lower() in TRUEY_VALS


def is_falsey(str: str) -> bool:
    return str.lower() in FALSEY_VALS


def add_symbol(sym: "Symbol"):
    all_symbols.append(sym)
    if sym.vram_start is not None:
        if sym.vram_start not in all_symbols_dict:
            all_symbols_dict[sym.vram_start] = []
        all_symbols_dict[sym.vram_start].append(sym)

    # For larger symbols, add their ranges to interval trees for faster lookup
    if sym.size > 4:
        all_symbols_ranges.addi(sym.vram_start, sym.vram_end, sym)


def initialize(all_segments: "List[Segment]"):
    global all_symbols
    global all_symbols_dict
    global all_symbols_ranges

    all_symbols = []
    all_symbols_dict = {}
    all_symbols_ranges = IntervalTree()

    def get_seg_for_name(name: str) -> Optional["Segment"]:
        for segment in all_segments:
            if segment.name == name:
                return segment
        return None

    # Manual list of func name / addrs
    for path in options.opts.symbol_addrs_paths:
        if path.exists():
            with open(path) as f:
                sym_addrs_lines = f.readlines()
                for line_num, line in enumerate(
                    tqdm.tqdm(sym_addrs_lines, desc=f"Loading symbols ({path.stem})")
                ):
                    line = line.strip()
                    if not line == "" and not line.startswith("//"):
                        comment_loc = line.find("//")
                        line_main = line
                        line_ext = ""

                        if comment_loc != -1:
                            line_ext = line[comment_loc + 2 :].strip()
                            line_main = line[:comment_loc].strip()

                        try:
                            line_split = line_main.split("=")
                            name = line_split[0].strip()
                            addr = int(line_split[1].strip()[:-1], 0)
                        except:
                            log.parsing_error_preamble(path, line_num, line)
                            log.write("Line should be of the form")
                            log.write(
                                "<function_name> = <address> // attr0:val0 attr1:val1 [...]"
                            )
                            log.write("with <address> in hex preceded by 0x, or dec")
                            log.write("")
                            raise

                        sym = Symbol(addr, given_name=name)

                        ignore_sym = False
                        if line_ext:
                            for info in line_ext.split(" "):
                                if ":" in info:
                                    if info.count(":") > 1:
                                        log.parsing_error_preamble(path, line_num, line)
                                        log.write(f"Too many ':'s in '{info}'")
                                        log.error("")

                                    attr_name, attr_val = info.split(":")
                                    if attr_name == "":
                                        log.parsing_error_preamble(path, line_num, line)
                                        log.write(
                                            f"Missing attribute name in '{info}', is there extra whitespace?"
                                        )
                                        log.error("")
                                    if attr_val == "":
                                        log.parsing_error_preamble(path, line_num, line)
                                        log.write(
                                            f"Missing attribute value in '{info}', is there extra whitespace?"
                                        )
                                        log.error("")

                                    # Non-Boolean attributes
                                    try:
                                        if attr_name == "type":
                                            type = attr_val
                                            sym.type = type
                                            continue
                                        if attr_name == "size":
                                            size = int(attr_val, 0)
                                            sym.given_size = size
                                            continue
                                        if attr_name == "rom":
                                            rom_addr = int(attr_val, 0)
                                            sym.rom = rom_addr
                                            continue
                                        if attr_name == "segment":
                                            seg = get_seg_for_name(attr_val)
                                            if seg is None:
                                                log.parsing_error_preamble(
                                                    path, line_num, line
                                                )
                                                log.write(
                                                    f"Cannot find segment '{attr_val}'"
                                                )
                                                log.error("")
                                            else:
                                                # Add segment to symbol
                                                sym.segment = seg
                                            continue
                                    except:
                                        log.parsing_error_preamble(path, line_num, line)
                                        log.write(
                                            f"value of attribute '{attr_name}' could not be read:"
                                        )
                                        log.write("")
                                        raise

                                    # Boolean attributes
                                    tf_val = (
                                        True
                                        if is_truey(attr_val)
                                        else False
                                        if is_falsey(attr_val)
                                        else None
                                    )
                                    if tf_val is None:
                                        log.parsing_error_preamble(path, line_num, line)
                                        log.write(
                                            f"Invalid Boolean value '{attr_val}' for attribute '{attr_name}', should be one of"
                                        )
                                        log.write([*TRUEY_VALS, *FALSEY_VALS])
                                        log.error("")
                                    else:
                                        if attr_name == "dead":
                                            sym.dead = tf_val
                                            continue
                                        if attr_name == "defined":
                                            sym.defined = tf_val
                                            continue
                                        if attr_name == "extract":
                                            sym.extract = tf_val
                                            continue
                                        if attr_name == "ignore":
                                            ignore_sym = tf_val
                                            continue
                                        if attr_name == "force_migration":
                                            sym.force_migration = tf_val
                                            continue
                                        if attr_name == "force_not_migration":
                                            sym.force_not_migration = tf_val
                                            continue
                        if ignore_sym:
                            ignored_addresses.add(sym.vram_start)
                            ignore_sym = False
                            continue

                        if sym.segment:
                            sym.segment.add_symbol(sym)

                        sym.user_declared = True
                        add_symbol(sym)


def initialize_spim_context(all_segments: "List[Segment]") -> None:
    global_vrom_start = None
    global_vrom_end = None
    global_vram_start = None
    global_vram_end = None
    overlay_segments: Set[spimdisasm.common.SymbolsSegment] = set()

    spim_context.bannedSymbols |= ignored_addresses

    from segtypes.common.code import CommonSegCode

    for segment in all_segments:
        if not isinstance(segment, CommonSegCode):
            # We only care about the VRAMs of code segments
            continue

        if segment.special_vram_segment:
            # Special segments which should not be accounted in the global VRAM calculation, like N64's IPL3
            continue

        if (
            not isinstance(segment.vram_start, int)
            or not isinstance(segment.vram_end, int)
            or not isinstance(segment.rom_start, int)
            or not isinstance(segment.rom_end, int)
        ):
            continue

        ram_id = segment.get_exclusive_ram_id()
        if ram_id is None:
            if global_vram_start is None:
                global_vram_start = segment.vram_start
            elif segment.vram_start < global_vram_start:
                global_vram_start = segment.vram_start

            if global_vram_end is None:
                global_vram_end = segment.vram_end
            elif global_vram_end < segment.vram_end:
                global_vram_end = segment.vram_end

            if global_vrom_start is None:
                global_vrom_start = segment.rom_start
            elif segment.rom_start < global_vrom_start:
                global_vrom_start = segment.rom_start

            if global_vrom_end is None:
                global_vrom_end = segment.rom_end
            elif global_vrom_end < segment.rom_end:
                global_vrom_end = segment.rom_end

        else:
            spim_segment = spim_context.addOverlaySegment(
                ram_id,
                segment.rom_start,
                segment.rom_end,
                segment.vram_start,
                segment.vram_end,
            )
            # Add the segment-specific symbols first
            for symbols_list in segment.seg_symbols.values():
                for sym in symbols_list:
                    add_symbol_to_spim_segment(spim_segment, sym)

            overlay_segments.add(spim_segment)

    if (
        global_vram_start is not None
        and global_vram_end is not None
        and global_vrom_start is not None
        and global_vrom_end is not None
    ):
        spim_context.changeGlobalSegmentRanges(
            global_vrom_start, global_vrom_end, global_vram_start, global_vram_end
        )

        # Check the vram range of the global segment does not overlap with any overlay segment
        for ovl_segment in overlay_segments:
            assert (
                ovl_segment.vramStart <= ovl_segment.vramEnd
            ), f"{ovl_segment.vramStart:X} {ovl_segment.vramEnd:X}"
            if (
                ovl_segment.vramEnd > global_vram_start
                and global_vram_end > ovl_segment.vramStart
            ):
                log.write(
                    f"Warning: the vram range ([0x{ovl_segment.vramStart:X}, 0x{ovl_segment.vramEnd:X}]) of the non-global segment at rom address 0x{ovl_segment.vromStart:X} overlaps with the global vram range ([0x{global_vram_start:X}, 0x{global_vram_end:X}])",
                    status="warn",
                )

    # pass the global symbols to spimdisasm
    for segment in all_segments:
        if not isinstance(segment, CommonSegCode):
            # We only care about the VRAMs of code segments
            continue

        ram_id = segment.get_exclusive_ram_id()
        if ram_id is not None:
            continue

        for symbols_list in segment.seg_symbols.values():
            for sym in symbols_list:
                add_symbol_to_spim_segment(spim_context.globalSegment, sym)


def add_symbol_to_spim_segment(
    segment: spimdisasm.common.SymbolsSegment, sym: "Symbol"
) -> spimdisasm.common.ContextSymbol:
    if sym.type == "func":
        context_sym = segment.addFunction(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    elif sym.type == "jtbl":
        context_sym = segment.addJumpTable(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    elif sym.type == "jtbl_label":
        context_sym = segment.addJumpTableLabel(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    elif sym.type == "label":
        context_sym = segment.addBranchLabel(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    else:
        context_sym = segment.addSymbol(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
        if sym.type is not None:
            context_sym.type = sym.type

    if sym.user_declared:
        context_sym.isUserDeclared = True
    if sym.defined:
        context_sym.isDefined = True
    if sym.rom is not None:
        context_sym.vromAddress = sym.rom
    if sym.given_size is not None:
        context_sym.size = sym.size
    if sym.force_migration:
        context_sym.forceMigration = True
    if sym.force_not_migration:
        context_sym.forceNotMigration = True
    context_sym.setNameGetCallbackIfUnset(lambda _: sym.name)

    return context_sym


def add_symbol_to_spim_section(
    section: spimdisasm.mips.sections.SectionBase, sym: "Symbol"
) -> spimdisasm.common.ContextSymbol:
    if sym.type == "func":
        context_sym = section.addFunction(
            sym.vram_start, isAutogenerated=not sym.user_declared, symbolVrom=sym.rom
        )
    elif sym.type == "jtbl":
        context_sym = section.addJumpTable(
            sym.vram_start, isAutogenerated=not sym.user_declared, symbolVrom=sym.rom
        )
    elif sym.type == "jtbl_label":
        context_sym = section.addJumpTableLabel(
            sym.vram_start, isAutogenerated=not sym.user_declared, symbolVrom=sym.rom
        )
    elif sym.type == "label":
        context_sym = section.addBranchLabel(
            sym.vram_start, isAutogenerated=not sym.user_declared, symbolVrom=sym.rom
        )
    else:
        context_sym = section.addSymbol(
            sym.vram_start, isAutogenerated=not sym.user_declared, symbolVrom=sym.rom
        )
        if sym.type is not None:
            context_sym.type = sym.type

    if sym.user_declared:
        context_sym.isUserDeclared = True
    if sym.defined:
        context_sym.isDefined = True
    if sym.rom is not None:
        context_sym.vromAddress = sym.rom
    if sym.given_size is not None:
        context_sym.size = sym.size
    if sym.force_migration:
        context_sym.forceMigration = True
    if sym.force_not_migration:
        context_sym.forceNotMigration = True
    context_sym.setNameGetCallbackIfUnset(lambda _: sym.name)

    return context_sym


def create_symbol_from_spim_symbol(
    segment: "Segment", context_sym: spimdisasm.common.ContextSymbol
) -> "Symbol":
    in_segment = False

    sym_type = None
    if context_sym.type == spimdisasm.common.SymbolSpecialType.jumptable:
        in_segment = True
        sym_type = "jtbl"
    elif context_sym.type == spimdisasm.common.SymbolSpecialType.function:
        sym_type = "func"
    elif context_sym.type == spimdisasm.common.SymbolSpecialType.branchlabel:
        in_segment = True
        sym_type = "label"
    elif context_sym.type == spimdisasm.common.SymbolSpecialType.jumptablelabel:
        in_segment = True
        sym_type = "jtbl_label"

    if not in_segment:
        if (
            context_sym.overlayCategory is None
            and segment.get_exclusive_ram_id() is None
        ):
            in_segment = segment.contains_vram(context_sym.vram)
        elif context_sym.overlayCategory == segment.get_exclusive_ram_id():
            if context_sym.vromAddress is not None:
                in_segment = segment.contains_rom(context_sym.vromAddress)
            else:
                in_segment = segment.contains_vram(context_sym.vram)

    sym = segment.create_symbol(
        context_sym.vram, in_segment, type=sym_type, reference=True
    )

    if sym.given_name is None and context_sym.name is not None:
        sym.given_name = context_sym.name

    # To keep the symbol name in sync between splat and spimdisasm
    context_sym.setNameGetCallback(lambda _: sym.name)

    if context_sym.size is not None:
        sym.given_size = context_sym.getSize()
    if context_sym.vromAddress is not None:
        sym.rom = context_sym.getVrom()
    if context_sym.isDefined:
        sym.defined = True
    if context_sym.referenceCounter > 0:
        sym.referenced = True

    return sym


def mark_c_funcs_as_defined():
    for symbol in all_symbols:
        if len(to_mark_as_defined) == 0:
            return
        sym_name = symbol.name
        if sym_name in to_mark_as_defined:
            symbol.defined = True
            to_mark_as_defined.remove(sym_name)


@dataclass
class Symbol:
    vram_start: int

    given_name: Optional[str] = None
    rom: Optional[int] = None
    type: Optional[str] = None
    given_size: Optional[int] = None
    segment: Optional["Segment"] = None

    defined: bool = False
    referenced: bool = False
    dead: bool = False
    extract: bool = True
    user_declared: bool = False

    force_migration: bool = False
    force_not_migration: bool = False

    _generated_default_name: Optional[str] = None
    _last_type: Optional[str] = None

    def __str__(self):
        return self.name

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Symbol):
            return False
        return self.vram_start == other.vram_start and self.segment == other.segment

    # https://stackoverflow.com/a/56915493/6292472
    def __hash__(self):
        return hash((self.vram_start, self.segment))

    def format_name(self, format: str) -> str:
        ret = format

        ret = ret.replace("$VRAM", f"{self.vram_start:08X}")

        if "$ROM" in ret:
            if not isinstance(self.rom, int):
                log.error(
                    f"Attempting to rom-name a symbol with no ROM address: {self.vram_start:X} typed {self.type}"
                )
            ret = ret.replace("$ROM", f"{self.rom:X}")

        if "$SEG" in ret:
            if self.segment is None:
                # This probably is fine - we can't expect every symbol to have a segment. Fall back to just the ram address
                return f"{self.vram_start:X}"
            assert self.segment is not None
            ret = ret.replace("$SEG", self.segment.name)

        return ret

    @property
    def default_name(self) -> str:
        if self._generated_default_name is not None:
            if self.type == self._last_type:
                return self._generated_default_name

        if self.segment:
            if isinstance(self.rom, int):
                suffix = self.format_name(self.segment.symbol_name_format)
            else:
                suffix = self.format_name(self.segment.symbol_name_format_no_rom)
        else:
            if isinstance(self.rom, int):
                suffix = self.format_name(options.opts.symbol_name_format)
            else:
                suffix = self.format_name(options.opts.symbol_name_format_no_rom)

        if self.type == "func":
            prefix = "func"
        elif self.type == "jtbl":
            prefix = "jtbl"
        elif self.type in {"jtbl_label", "label"}:
            return f".L{suffix}"
        else:
            prefix = "D"

        self._last_type = self.type
        self._generated_default_name = f"{prefix}_{suffix}"
        return self._generated_default_name

    @property
    def rom_end(self):
        return None if not self.rom else self.rom + self.size

    @property
    def vram_end(self):
        return self.vram_start + self.size

    @property
    def name(self) -> str:
        return self.given_name if self.given_name else self.default_name

    @property
    def size(self) -> int:
        if self.given_size is not None:
            return self.given_size
        return 4

    def contains_vram(self, offset):
        return offset >= self.vram_start and offset < self.vram_end

    def contains_rom(self, offset):
        return offset >= self.rom and offset < self.rom_end