#!/usr/bin/env python3
"""Redundancy ranker for the SF3 match worklist.

Phase 11 finding 79/88/99. The cost of a matching row is set by TIE-BREAK DENSITY, not
by size. Two independent measurements:

  * worker A: 548 B / 460 B / 356 B / 204 B all matched on the FIRST spelling, while its
    one nine-attempt failure was the SMALLEST row it attacked (176 B);
  * worker D: 1232 B matched on the 2nd spelling and 700 B on the 3rd, while its 248 B row
    took 4.

So rank rows by how much of the body is REPEATED. A body made of a few repeated blocks is
cheap at any size; a tie-break-dense body (a long straight-line call sequence, a pointer
walk, a table walk) is expensive at any size.

The metric counts, for n = 2, 3 and 4, the fraction of n-instruction subsequences that have
already been seen earlier in the body, and takes the **MAXIMUM** of the three. Nops are excluded
because `lw`/`nop` pairs otherwise inflate ordinary arithmetic rows (worker A).

**REDUNDANCY IS CORRELATED WITH SIZE, WHICH BIASES THE TOP OF THE LIST (worker D).**
A larger body simply has more chances to repeat a 3-gram, so the highest scores are dominated by
very large rows — worker D measured 3288 B at 0.91 and 8508 B at 0.87 sitting above the 248 B rows
it was actually succeeding on. **Since the Phase 11 milestone counts BODIES and not bytes, a large
row is worth exactly one body at many times the context cost.** So use `--max-size` to match the
tool to how the work is actually picked:

    ./tools/sf3_rank <partition> --max-size 300 --top 25

Worker D's effective filter was **redundancy per unit size**, which is what a size cap
approximates. The tool is not wrong to score what it was asked to score; the *view* needed the cap.

**The statistic is a MAX, not a mean — this matters for comparability.** A body that repeats
2-grams heavily but 4-grams not at all scores high, so a high score means "some length scale is
very repetitive", not "repetitive at every scale". Worker B raised this: if any implementation
uses a mean or a different normalisation the numbers are not comparable across partitions.
**This tool is the reference implementation; compare scores only against its output.**

Companion rule (finding 99): once you have matched a row, the row ADJACENT to it is an even
better bet than the top of this list — the binary is laid out by translation unit, so
neighbours share the author's habits. That rule was 4-for-4 across two workers.

Usage:
    sf3_rank [worklist.tsv ...] [--top N] [--min-size B] [--max-size B]

With no worklist arguments it ranks config/match_worklist.tsv. Rows are printed
highest-score first, so the top of the output is the cheapest expected work.
"""
from __future__ import annotations

import struct
import sys
from collections import Counter
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
EXE = REPO / "extracted/SCUS_946.40;1"
WORKLIST = REPO / "config/match_worklist.tsv"
LOAD_ADDRESS = 0x80010000
HEADER_SIZE = 0x800


def body_words(payload: bytes, start: int, end: int) -> list[int]:
    """The instruction words of a body, in order."""
    offset = HEADER_SIZE + (start - LOAD_ADDRESS)
    count = (end - start) // 4
    return list(struct.unpack_from(f"<{count}I", payload, offset))


def is_nop(word: int) -> bool:
    return word == 0


def redundancy(words: list[int]) -> float:
    """Best repeated-subsequence fraction over n = 2, 3, 4, ignoring nops."""
    insns = [word for word in words if not is_nop(word)]
    scores: list[float] = []
    for n in (2, 3, 4):
        if len(insns) < n + 1:
            continue
        grams = [tuple(word >> 26 for word in insns[i:i + n])
                 for i in range(len(insns) - n + 1)]
        seen: Counter[tuple[int, ...]] = Counter()
        repeats = 0
        for gram in grams:
            if seen[gram]:
                repeats += 1
            seen[gram] += 1
        scores.append(repeats / len(grams))
    return max(scores) if scores else 0.0


CALLEE_SAVED = frozenset(range(16, 24))  # s0..s7


def fragment_reason(insns: list[int]) -> str:
    """Worker A's fragment check: does the body use a frame it never set up?

    A function may legally start with `beq a0,a1,...` or `sh a0,0(a0)`. It may NOT start by
    touching a slot of a frame that does not exist yet. Worker A found `0x800C3490`, whose
    first instruction is `sw v0,32(sp)` -- a stack write before any frame setup, using
    registers set by the code BEFORE the row -- and whose identical tail also appears at
    0x800C3470. It is a shared/jump-target block Ghidra promoted to a function, so it is
    inside no region and nobody can match it standalone.

    QUALIFIED BY WORKER B, WHICH MADE IT DISJOINT. The first version fired on any nonzero
    `sp` offset and flagged 2 of the 555 registered regions (`0x800923E8`, `0x80099DC4`).
    Worker B adjudicated its own hit `0x800B704C` and showed it is a **legal frameless leaf
    with eight arguments**: in o32 the callee's `sp` is unchanged at entry, so `sp+16..sp+28`
    IS the caller's outgoing area -- arguments 4-7. Reading it before any `addiu sp,sp,-N` is
    exactly what a frameless >4-argument leaf looks like.

    So the incoming argument area (`sp+0..sp+31`) is excluded, and only a **negative** offset
    (a frame slot below an unmoved `sp`) or an offset **beyond** the 8-argument area is a real
    fragment. With that qualification the check is **disjoint from the corpus: 0 of 555
    registered regions**, and the suspects across all four partitions drop from 5 to 1.
    """
    if not insns:
        return "empty"
    word = insns[0]
    opcode = word >> 26
    rs = (word >> 21) & 0x1F
    rt = (word >> 16) & 0x1F
    immediate = word & 0xFFFF
    if opcode in (0x23, 0x2B, 0x0F, 0x20, 0x24, 0x25, 0x28, 0x29, 0x2C, 0x2D, 0x2E, 0x3F):
        if (rs == 29 or rt == 29) and immediate != 0:
            offset = immediate - 0x10000 if immediate >= 0x8000 else immediate
            # In o32 the callee's `sp` is unchanged at entry, so the incoming argument
            # area sits at sp+0..sp+28 (args 0-3 live in a0-a3 but their home slots are
            # there; args 4-7 are passed at sp+16..sp+28). Worker B adjudicated
            # 0x800B704C as a legal frameless leaf with EIGHT arguments whose first
            # instruction is `lw t0,16(sp)` -- so reading that area at entry is legal
            # and must not fire. Only a NEGATIVE offset (a frame slot below an unmoved
            # sp) or an offset beyond the 8-argument incoming area is a real fragment.
            if offset < 0 or offset >= 32:
                return f"first instruction touches sp{offset:+d} (outside the incoming arg area)"
    used = {rs, rt}
    if opcode == 0:
        used.add((word >> 11) & 0x1F)
    hit = used & CALLEE_SAVED
    if hit and opcode != 0x09 and len(insns) > 8:
        saves = any((x >> 26) == 0x2B and ((x >> 21) & 0x1F) == 29 for x in insns[:12])
        if not saves:
            return f"uses {sorted(hit)} with no save in the first 12 instructions"
    return ""


def load_rows(paths: list[Path], payload: bytes, min_size: int, max_size: int) -> list[tuple[float, int, str]]:
    rows: list[tuple[float, int, str]] = []
    for path in paths:
        for line in path.read_text().splitlines():
            if line.startswith("#") or not line.strip():
                continue
            fields = line.split("\t")
            start, end = int(fields[1], 16), int(fields[2], 16)
            size = end - start
            if size < min_size or size > max_size:
                continue
            rows.append((redundancy(body_words(payload, start, end)), size, fields[1]))
    return rows


def main(argv: list[str]) -> int:
    paths, top, min_size, max_size, fragments = [], 40, 0, 1 << 30, False
    index = 0
    while index < len(argv):
        arg = argv[index]
        if arg == "--fragments":
            fragments = True
        elif arg == "--top":
            index += 1
            top = int(argv[index])
        elif arg == "--min-size":
            index += 1
            min_size = int(argv[index])
        elif arg == "--max-size":
            index += 1
            max_size = int(argv[index])
        elif arg.startswith("-"):
            print(f"unknown option: {arg}", file=sys.stderr)
            return 2
        else:
            paths.append(Path(arg))
        index += 1
    if not paths:
        paths = [WORKLIST]
    for path in paths:
        if not path.exists():
            print(f"missing worklist: {path}", file=sys.stderr)
            return 2

    payload = EXE.read_bytes()

    if fragments:
        print("ADVISORY ONLY -- measured at 2 false positives across 555 registered regions.")
        print(f"{'size':>5}  address  reason")
        count = 0
        for path in paths:
            for line in path.read_text().splitlines():
                if line.startswith("#") or not line.strip():
                    continue
                fields = line.split("\t")
                start, end = int(fields[1], 16), int(fields[2], 16)
                if end - start < min_size or end - start > max_size:
                    continue
                reason = fragment_reason(body_words(payload, start, end))
                if reason:
                    count += 1
                    print(f"{end - start:5d}  {fields[1]}  {reason}")
        print(f"\nsuspects={count}")
        return 0

    rows = load_rows(paths, payload, min_size, max_size)
    rows.sort(key=lambda row: (-row[0], row[1], row[2]))

    print(f"{'score':>6} {'size':>5}  address")
    for score, size, start in rows[:top]:
        print(f"{score:6.2f} {size:5d}  {start}")
    print(f"\nrows={len(rows)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
