app.py

"""N64LLMDecompile - HuggingFace ZeroGPU Space

Model: fine-tuned LLM4Decompile-v2 9B (ghidra-refine, merged). Input is Ghidra
pseudo-C; output is bare C. Completion-style prompt.

Flow:
  Paste Ghidra pseudo-C -> LLM -> C  (+ optional match rating if you also upload
                                      the target object to score against)

The ONLY thing you must edit is MODEL_ID below (your HF username). Everything else
runs as-is; scoring degrades gracefully if its pieces aren't present yet.
"""

from __future__ import annotations

import os
import re
import shlex
import subprocess
import tempfile
from pathlib import Path
from typing import Any

import gradio as gr

try:
    import diff as asm_differ  # type: ignore
    ASM_DIFFER_OK = True
except Exception as exc:  # noqa: BLE001
    asm_differ = None  # type: ignore
    ASM_DIFFER_OK = False
    ASM_DIFFER_IMPORT_ERROR = str(exc)

try:
    import spaces  # type: ignore
    GPU = spaces.GPU
except Exception:
    def GPU(fn=None, **_kwargs):
        if fn is None:
            return lambda f: f
        return fn


# --------------------------------------------------------------------------- #
# Config  --  EDIT THIS LINE: put your HF username before the slash
# --------------------------------------------------------------------------- #

MODEL_ID = "MatthewReingold/llm4decompile-9b-v2-n64-finetune"

MAX_NEW_TOKENS = 2048

COMPILERS_DIR = (Path(__file__).parent / "compilers").resolve()
COMPILER_PATHS = {
    "ido5.3":      str(COMPILERS_DIR / "ido5.3" / "cc"),
    "ido7.1":      str(COMPILERS_DIR / "ido7.1" / "cc"),
    "gcc2.7.2kmc": str(COMPILERS_DIR / "gcc2.7.2kmc" / "gcc"),
    "gcc2.8.1pm":  str(COMPILERS_DIR / "gcc2.8.1pm" / "gcc"),
}
COMPILERS = list(COMPILER_PATHS)

FLAG_HINTS = {
    "ido5.3":      "-O2 -mips2 -Xfullwarn -signed -32 ...",
    "ido7.1":      "-O2 -mips2 -Xfullwarn -signed -32 ...",
    "gcc2.7.2kmc": "-O2 -mips3 -G0 ...",
    "gcc2.8.1pm":  "-O2 -mips3 -G0 ...",
}

OBJDUMP = "mips-linux-gnu-objdump"
COMPILE_TIMEOUT_SECONDS = 30
OBJDUMP_TIMEOUT_SECONDS = 30
MAX_ASM_LINES = 10000
LANGUAGE_ID = "MIPS:BE:32:default"


def score_label(score: int) -> str:
    if score == 0:
        return "Perfect match"
    if score <= 10:
        return "Near-perfect"
    if score <= 50:
        return "Close"
    if score <= 100:
        return "Partial"
    if score <= 500:
        return "Rough"
    return "Far off"


# --------------------------------------------------------------------------- #
# Compile + score  (from eval.py, adapted for one upload at a time)
# --------------------------------------------------------------------------- #

def strip_line_comments(source_code: str) -> str:
    return re.sub(r"//[^\n]*", "", source_code)


def write_translation_unit(path: Path, context: str, source_code: str) -> None:
    pieces = []
    if context.strip():
        pieces.append(strip_line_comments(context.rstrip()))
    pieces.append(source_code.rstrip())
    path.write_text("\n\n".join(pieces) + "\n", encoding="utf-8")


def make_empty_stub_source(name: str) -> str:
    if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name or ""):
        name = "stub"
    return f"void {name}(void) {{\n}}\n"


def run_command(command: list[str], timeout_seconds: int,
                env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        command, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True, timeout=timeout_seconds, env=env,
    )


def compiler_command(compiler_name: str, flags: str,
                     source_path: Path, output_path: Path) -> list[str] | None:
    compiler_path = COMPILER_PATHS.get(compiler_name)
    if compiler_path is None:
        return None
    compiler = os.path.expanduser(compiler_path)
    return [compiler, *shlex.split(flags or ""), "-c", str(source_path), "-o", str(output_path)]


def compile_source(compiler_name: str, flags: str, context: str, work_dir: Path,
                   source_code: str, output_name: str) -> tuple[Path | None, dict | None]:
    source_path = work_dir / f"{output_name}.c"
    output_path = work_dir / f"{output_name}.o"
    write_translation_unit(source_path, context, source_code)

    command = compiler_command(compiler_name, flags, source_path, output_path)
    if command is None:
        return None, {"type": "unsupported_compiler", "compiler": compiler_name}

    env = os.environ.copy()
    compiler_dir = os.path.dirname(os.path.expanduser(command[0]))
    env["PATH"] = compiler_dir + ":" + env.get("PATH", "")

    try:
        completed = run_command(command, COMPILE_TIMEOUT_SECONDS, env=env)
    except FileNotFoundError as exc:
        return None, {"type": "compiler_not_found", "command": command, "error": str(exc)}
    except subprocess.TimeoutExpired:
        return None, {"type": "compile_timeout", "timeout_seconds": COMPILE_TIMEOUT_SECONDS}

    if completed.returncode != 0 or not output_path.exists():
        return None, {
            "type": "compile_failed",
            "returncode": completed.returncode,
            "stdout": completed.stdout,
            "stderr": completed.stderr,
        }
    return output_path, None


def make_diff_config() -> Any:
    return asm_differ.Config(
        arch=asm_differ.get_arch("mips"),
        diff_obj=True, file=None, make=False, source_old_binutils=False,
        diff_section=".text", inlines=False,
        max_function_size_lines=MAX_ASM_LINES,
        max_function_size_bytes=MAX_ASM_LINES * 4,
        formatter=asm_differ.PythonFormatter(arch_str="mips"),
        diff_mode=asm_differ.DiffMode.NORMAL,
        base_shift=0, skip_lines=0, compress=None,
        show_rodata_refs=True, show_branches=True, show_line_numbers=False,
        show_source=False, stop_at_ret=False, ignore_large_imms=False,
        ignore_addr_diffs=False, algorithm="levenshtein",
        reg_categories={}, diff_function_symbols=False,
    )


def strip_objdump_headers(objdump_output: str) -> str:
    lines, in_disassembly = [], False
    for line in objdump_output.splitlines():
        if line.startswith("Disassembly of section "):
            in_disassembly = True
            continue
        if in_disassembly:
            lines.append(line)
    return "\n".join(lines)


def disassemble_elf(path: Path, config: Any) -> str:
    command = [OBJDUMP, *config.arch.arch_flags, "-dr", str(path)]
    completed = run_command(command, OBJDUMP_TIMEOUT_SECONDS)
    if completed.returncode != 0:
        raise RuntimeError(completed.stderr or completed.stdout)
    if hasattr(asm_differ, "preprocess_objdump_out"):
        return asm_differ.preprocess_objdump_out(None, path.read_bytes(), completed.stdout, config)
    return strip_objdump_headers(completed.stdout)


def score_disassembly(expected_disasm: str, actual_disasm: str, config: Any) -> int:
    expected_lines = asm_differ.process(expected_disasm, config)
    actual_lines = asm_differ.process(actual_disasm, config)
    return int(asm_differ.do_diff(expected_lines, actual_lines, config).score)


def compile_and_score(generated_c: str, object_bytes: bytes, compiler: str,
                      flags: str, context: str, function_name: str) -> dict:
    if not ASM_DIFFER_OK:
        return {"status": "scoring_unavailable"}
    if compiler not in COMPILER_PATHS:
        return {"status": "unsupported_compiler", "compiler": compiler}

    config = make_diff_config()
    predicted = strip_line_comments(generated_c)

    with tempfile.TemporaryDirectory(prefix="n64_space_") as tmp_name:
        tmp = Path(tmp_name)
        expected_path = tmp / "expected.o"
        expected_path.write_bytes(object_bytes)

        actual_path, error = compile_source(compiler, flags, context, tmp, predicted, "actual")
        if error is not None:
            return {"status": error["type"], "error": error, "compiler": compiler}

        baseline_path, baseline_error = compile_source(
            compiler, flags, context, tmp, make_empty_stub_source(function_name), "baseline"
        )

        try:
            expected_disasm = disassemble_elf(expected_path, config)
            actual_disasm = disassemble_elf(actual_path, config)
            score = score_disassembly(expected_disasm, actual_disasm, config)
        except Exception as exc:  # noqa: BLE001
            return {"status": "diff_failed", "error": str(exc), "compiler": compiler}

        result = {"status": "evaluated", "score": score, "compiler": compiler}
        if baseline_error is None:
            try:
                baseline_disasm = disassemble_elf(baseline_path, config)
                result["baseline_score"] = score_disassembly(expected_disasm, baseline_disasm, config)
            except Exception:  # noqa: BLE001
                pass
        return result


# --------------------------------------------------------------------------- #
# Model  --  loads at startup; if MODEL_ID is wrong the Space still boots and
#            generate_c reports the error instead of crashing.
# --------------------------------------------------------------------------- #

_MODEL = None
_TOKENIZER = None
_MODEL_LOAD_ERROR = None
try:
    from transformers import AutoModelForCausalLM, AutoTokenizer
    _TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID)
    _MODEL = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
except Exception as exc:  # noqa: BLE001
    _MODEL_LOAD_ERROR = str(exc)


def l4d_prompt(pseudo_c: str) -> str:
    return f"# This is the assembly code:\n{pseudo_c}\n# What is the source code?\n"


@GPU
def generate_c(pseudo_c: str) -> str:
    prompt = l4d_prompt(pseudo_c)
    if _MODEL is None or _TOKENIZER is None:
        msg = _MODEL_LOAD_ERROR or "model not loaded — check MODEL_ID"
        return f"/* model unavailable: {msg} */\n/* prompt that would be sent:\n\n{prompt}\n*/\n"

    _MODEL.to("cuda")
    inputs = _TOKENIZER(prompt, return_tensors="pt").to("cuda")
    out = _MODEL.generate(
        **inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False,
        repetition_penalty=1.1, pad_token_id=_TOKENIZER.eos_token_id,
    )
    completion = out[0][inputs["input_ids"].shape[1]:]
    return _TOKENIZER.decode(completion, skip_special_tokens=True)


# --------------------------------------------------------------------------- #
# Rating render
# --------------------------------------------------------------------------- #

def render_rating(result: dict | None, scored: bool) -> str:
    if not scored:
        return ("**No match rating.** Upload the target object alongside your paste "
                "(or use object mode) to score the compiled output against it.")
    if result is None:
        return ""
    status = result.get("status")
    if status == "scoring_unavailable":
        return ("**Scoring unavailable** — `diff.py` (asm-differ) isn't in the repo yet. "
                "The C is below; add diff.py to enable match ratings.")
    if status == "unsupported_compiler":
        return f"**Unsupported compiler** `{result.get('compiler')}`. The C is below."
    if status == "compiler_not_found":
        return ("**Compiler binary not found** — check that `compilers/` is in the Space "
                "and executable. The C is below.")
    if status == "compile_timeout":
        return "**Compile timed out.** The C is below — you may be able to simplify it."
    if status == "compile_failed":
        err = (result.get("error") or {}).get("stderr") or "compile error"
        return (f"**Did not compile** with `{result.get('compiler')}`. "
                f"The C is below anyway — you may be able to fix it.\n\n```\n{err[:600]}\n```")
    if status == "diff_failed":
        return f"**Diff failed:** {str(result.get('error'))[:300]}. The C is below."

    score = result["score"]
    line = f"**{score_label(score)}** — asm-differ score `{score}` (0 = byte-perfect, lower is better)."
    base = result.get("baseline_score")
    if isinstance(base, int):
        verb = "beats" if score < base else "does not beat"
        line += f"\n\nThis {verb} the empty-stub baseline (`{base}`)."
    return line


# --------------------------------------------------------------------------- #
# Orchestrator
# --------------------------------------------------------------------------- #

def _read_bytes(file_path) -> bytes | None:
    if file_path is None:
        return None
    with open(file_path, "rb") as handle:
        return handle.read()


def run(pasted_pseudo_c, paste_object_file, compiler, flags, context):
    context = context or ""

    pseudo_c = (pasted_pseudo_c or "").strip()
    if not pseudo_c:
        return "", "**Paste some Ghidra pseudo-C to start.**"
    generated_c = generate_c(pseudo_c)

    object_bytes = _read_bytes(paste_object_file)
    if object_bytes is not None:
        result = compile_and_score(generated_c, object_bytes, compiler, flags, context, "")
        return generated_c, render_rating(result, scored=True)
    return generated_c, render_rating(None, scored=False)


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #

CSS = """
.gradio-container { max-width: 1100px !important; }
#out-c textarea, #pseudo textarea { font-family: ui-monospace, SFMono-Regular, Menlo, monospace; }
"""

with gr.Blocks(title="N64LLMDecompile", theme=gr.themes.Base(), css=CSS) as demo:
    gr.Markdown(
        "# N64LLMDecompile\n"
        "Turn N64 MIPS into C with LLM4Decompile-v2 fine-tuned on decomp.me. "
        "Paste Ghidra pseudo-C to get C back — optionally upload the target object "
        "to score the result with a match rating."
    )

    with gr.Row():
        with gr.Column(scale=1):
            pasted_pseudo_c = gr.Textbox(
                label="Ghidra pseudo-C", placeholder="Paste the decompiler output here…",
                lines=12, elem_id="pseudo",
            )
            paste_object_file = gr.File(
                label="Target object for scoring (optional) — upload to rate the output",
                file_types=[".o", ".elf", ".bin"], type="filepath",
            )

            compiler = gr.Dropdown(choices=COMPILERS, value="ido7.1",
                                   label="Compiler (for scoring)")
            flags = gr.Textbox(label="Compiler flags (for scoring)",
                               placeholder=FLAG_HINTS["ido7.1"])
            context = gr.Textbox(
                label="Context / headers (optional, for scoring)",
                placeholder="typedefs, structs, externs the function needs to compile", lines=4,
            )
            run_btn = gr.Button("Decompile to C", variant="primary")

        with gr.Column(scale=1):
            out_c = gr.Code(label="Generated C", language="c", elem_id="out-c")
            rating = gr.Markdown()

    compiler.change(lambda c: gr.update(placeholder=FLAG_HINTS.get(c, "")),
                    inputs=compiler, outputs=flags)
    run_btn.click(
        run,
        inputs=[pasted_pseudo_c, paste_object_file, compiler, flags, context],
        outputs=[out_c, rating],
    )


if __name__ == "__main__":
    demo.launch()