Mulberry/tools/python/annotate.py

'''
alessatool/annotate

mark up assembly with line numbers.
can also be used to create full-tu scratches in decomp.me.

example:

alessatool annotate \
    --asm-path silent-hill-2/config/SLUS_202.28/asm/Effect2/hh_class_water_03.s \
    --tu

run `alessatool annotate --help` for more information.
'''

from sys import stdout, stdin
from pathlib import Path
from subprocess import run
from dataclasses import dataclass
from constants import *
from struct import unpack

@dataclass
class AnnotationArgs:
    vram_start: int | None
    vram_end: int | None
    elf_path: Path
    asm_path: Path | None
    out_path: Path
    addr2line_path: Path
    line_file_path: Path | None
    tu: bool
    encoding: str
    stdout: bool
    verbose: bool

def annotate_asm(args: AnnotationArgs):
    if args.asm_path:
        with open(args.asm_path, "r") as asm_file:
            asm_contents = asm_file.read()
    else:
        asm_contents = stdin.read()

    asm_lines = asm_contents.splitlines()
    asm_line_index = 0

    vram_start = args.vram_start
    vram_end = args.vram_end

    if vram_start is None or vram_end is None:
        vram_start, vram_end = find_vram_bounds(asm_lines)

    line_file_path = get_line_file_path(args)

    if line_file_path is None or not line_file_path.exists():
        addresses = (f"0x{v:X}" for v in range(vram_start, vram_end, 0x4))
        proc = run([args.addr2line_path, "-e", args.elf_path, *addresses], capture_output=True, encoding=args.encoding)
        addr2line_output_lines = proc.stdout.splitlines()
    else:
        # parse the binary line number file.

        # the format is a list of u_shorts, one per line number.
        # there should be one line number per vram address, and each vram
        # address should be exactly 4 bytes apart, mirroring how the addr2line
        # output is formatted

        with open(line_file_path, "rb") as line_file:
            line_data = line_file.read()
            line_numbers = unpack(f"<{len(line_data) // 2}H", line_data)
            compile_unit = args.asm_path.with_suffix(".c").name
            addr2line_output_lines = list(map(lambda n : to_addr2line_format(compile_unit, n), line_numbers))

    main_tu_name = None
    prev_tu_name = None
    prev_line_number = -1
    function_count = 0
    is_in_function_label = False
    is_in_text_section = False
    current_vram_addr = vram_start
    annotated_asm_lines = []

    for addr_index in range(0, len(addr2line_output_lines) - 1):
        line = addr2line_output_lines[addr_index]

        if line.startswith("?"):
            current_vram_addr += 0x4
            continue

        separator_index = line.rfind(":")
        current_line_number = int(line[separator_index+1:])

        if current_line_number == prev_line_number:
            current_vram_addr += 0x4
            continue

        vram_addr_str = f"{current_vram_addr:X}"

        current_tu_name = line[0:separator_index]
        if current_tu_name != main_tu_name and current_tu_name != prev_tu_name:
            if main_tu_name:
                print(f"[warn] alessatool/annotate: address at {vram_addr_str} belongs to {current_tu_name}")
            else:
                main_tu_name = current_tu_name
        prev_tu_name = current_tu_name

        while True:
            if asm_line_index >= len(asm_lines):
                raise AssertionError(
                    f"address 0x{current_vram_addr:X} not found in asm"
                )

            asm_line = asm_lines[asm_line_index]

            if asm_line.startswith(SECTION_DIRECTIVE):
                is_in_text_section = asm_line.startswith(TEXT_SECTION_DIRECTIVE)

            if not is_in_text_section:
                annotated_asm_lines.append(asm_line)
                asm_line_index += 1
                continue

            if line_has_vram_addr(asm_line, vram_addr_str):
                break

            should_append_asm_line = True
            if args.tu:
                asm_line_trimmed = asm_line.strip()

                # track when we go in & out of function symbols
                if asm_line_trimmed.startswith(FUNCTION_SYMBOL_LABEL):
                    function_count += 1
                    is_in_function_label = True
                    annotated_asm_lines.append(f"{UNIQUE_TEXT_SECTION_DIRECTIVE}{function_count}")
                    annotated_asm_lines.append("")
                elif asm_line_trimmed.startswith(END_FUNCTION_SYMBOL_LABEL):
                    is_in_function_label = False

                # remove `nop`s
                if not is_in_function_label and asm_line_trimmed.endswith("nop"):
                    should_append_asm_line = False
                
                # remove `macro.inc` include directive
                elif asm_line_trimmed == INCLUDE_MACRO_INC_DIRECTIVE:
                    should_append_asm_line = False

            if should_append_asm_line:
                annotated_asm_lines.append(asm_line)

            asm_line_index += 1

        annotated_asm_lines.append(f"\t.loc 1 {current_line_number}")    
        annotated_asm_lines.append(asm_line)
        asm_line_index += 1

        prev_line_number = current_line_number
        current_vram_addr += 0x4
    
    while asm_line_index < len(asm_lines) - 1:
        annotated_asm_lines.append(asm_lines[asm_line_index])
        asm_line_index += 1

    assert main_tu_name, "no valid compilation unit found"
    annotated_asm_lines = [
        ".section .debug",
        ".previous",
        ".text",
        f".file 1 \"{main_tu_name}\"",
        *annotated_asm_lines
    ]

    append_final_new_line(annotated_asm_lines)
    annotated_asm_contents = "\n".join(annotated_asm_lines)

    if not args.stdout and args.out_path:
        with open(args.out_path, "w") as out_file:
            out_file.write(annotated_asm_contents)
        if args.verbose:
            print(f"alessatool/annotate: wrote asm to {args.out_path}")
    else:
        stdout.write(annotated_asm_contents)

def get_line_file_path(args: AnnotationArgs):
    if args.line_file_path is not None:
        return args.line_file_path

    if args.elf_path.name == SH2_SERIAL and args.asm_path and "Event/stage" in args.asm_path.as_posix():
        return Path(f"{TOOLS}/alessatool/dwarf") / Path(args.asm_path.name).with_suffix(".line")

    return None

def to_addr2line_format(compile_unit: str, line_number: int):
    return f"{compile_unit}:{line_number}"

def find_vram_bounds(asm_lines: list[str]):
    start = None
    passed_first_label = False
    end = None

    for line in asm_lines:
        # assumes first address comment is on line after first label
        if "label" in line and not passed_first_label:
            passed_first_label = True
            continue

        if passed_first_label and start is None:
            start = get_vram_addr_from_line(line)

        # assumes .text section comes first
        if line.startswith(SECTION_DIRECTIVE) and not line.startswith(TEXT_SECTION_DIRECTIVE):
            break

        vram_addr = get_vram_addr_from_line(line)
        if vram_addr is None:
            continue

        end = vram_addr

    return (start, end)

def get_vram_addr_from_line(line: str):
    if "/*" not in line:
        return None

    _before_comment, after_comment = line.split("/* ", 1)
    _file_addr, vram_addr, _rest = after_comment.split(" ", 2)

    try: 
        return int(vram_addr, 0x10)
    except:
        return None

def line_has_vram_addr(line: str, addr_str: str) -> bool:
    if addr_str not in line or "*/" not in line:
        return False

    return line.index("*/") > line.index(addr_str)

def append_final_new_line(lines: list[str]):
    if lines[-1] != "":
        lines.append("")