bootloader/tools/sizes.py

#!/usr/bin/env python3
"""What the loader images actually measure, and whether the README still agrees.

The size matrix asserts every image fits its slot; it says nothing about the
numbers the README prints, and those drift. Every row of that table was eight
bytes stale once `startup::caller_page()` landed — common code, so every build
moved at once and no test noticed, because none of them was over budget.

Two questions, both answered from built trees:

  sizes.py max            the largest image per chip, and anything over budget
  sizes.py check-readme   the README's per-chip table against what is built

Nothing here knows a chip's geometry. The (image, budget) pairs come from each
build's own `CTestTestfile.cmake` — the same values the gate checks — so the
slot rules stay where they belong, in `pureboot/CMakeLists.txt`, and a chip
added or a budget changed needs no edit here. Only trees a configure preset
still owns are read: a stale directory keeps its last build, and a loader built
before a slot changed will happily report a size that was true once
(`tools/prune-build-trees.sh` in libavr removes them).

Sizes come from `avr-size`, and a target is only as current as its last build —
run the gate first if you want the table checked against today's source.
"""

from __future__ import annotations

import argparse
import pathlib
import re
import shutil
import subprocess
import sys

ROOT = pathlib.Path(__file__).resolve().parents[1]
# add_test(<name>.size ... -DELF=<path> ... -DLIMIT=<n> ...) — the gate's own
# pairing of an image with the budget it must fit.
# ctest writes the name as a bracket argument ([=[name.size]=]) and quotes the
# rest, so the name starts after the bracket and the path ends at the quote.
SIZE_TEST = re.compile(r'add_test\(\s*\[=\[(?P<name>[^\]]+?)\.size\]=\][^\n]*?'
                       r'-DELF=(?P<elf>[^"\s]+)[^\n]*?-DLIMIT=(?P<limit>\d+)')


def avr_size() -> str:
    for env in (ROOT / "../../toolchain").resolve().glob("avr-gcc-*/bin/avr-size"):
        if env.is_file():
            return str(env)
    found = shutil.which("avr-size")
    if not found:
        sys.exit("no avr-size found (build the toolchain, or put it on PATH)")
    return found


def preset_dirs() -> list[pathlib.Path]:
    """Build trees a configure preset still owns, newest-listed first."""
    listing = subprocess.run(["cmake", "--list-presets"], cwd=ROOT, capture_output=True, text=True)
    names = re.findall(r'^\s*"(.+)"$', listing.stdout, re.MULTILINE)
    if not names:
        sys.exit("cmake --list-presets returned nothing — run from a configured checkout")
    return [d for d in (ROOT / "build" / n for n in names) if (d / "CTestTestfile.cmake").is_file()]


def measure(paths: list[str], tool: str) -> dict[str, int]:
    """.text per ELF, in one avr-size call per batch."""
    sizes: dict[str, int] = {}
    for start in range(0, len(paths), 400):
        batch = [p for p in paths[start:start + 400] if pathlib.Path(p).is_file()]
        if not batch:
            continue
        out = subprocess.run([tool, *batch], capture_output=True, text=True).stdout
        for line in out.splitlines()[1:]:
            fields = line.split()
            if len(fields) >= 6 and fields[0].isdigit():
                sizes[fields[5]] = int(fields[0])
    return sizes


def collect() -> dict[str, list[tuple[str, int, int]]]:
    """chip -> [(target, text, limit)], from every owned build tree."""
    tool = avr_size()
    found: dict[str, list[tuple[str, str, int]]] = {}
    for tree in preset_dirs():
        chip = tree.name.split("-")[0]
        for match in SIZE_TEST.finditer((tree / "CTestTestfile.cmake").read_text()):
            found.setdefault(chip, []).append((match["name"], match["elf"], int(match["limit"])))
    sizes = measure([elf for rows in found.values() for _, elf, _ in rows], tool)
    measured = {
        chip: sorted(((name, sizes[elf], limit) for name, elf, limit in rows if elf in sizes),
                     key=lambda row: -row[1])
        for chip, rows in sorted(found.items())
    }
    # A configured-but-unbuilt preset registers its tests with no images behind
    # them; it is not a chip with nothing to say, it is a chip not built yet.
    return {chip: rows for chip, rows in measured.items() if rows}


def cmd_max(args) -> int:
    measured = collect()
    if not measured:
        sys.exit("nothing built — configure and build a preset first")
    over = []
    print(f"{'chip':<13} {'largest image':<34} {'.text':>6} {'budget':>7}  headroom")
    for chip, rows in measured.items():
        name, text, limit = rows[0]
        flag = "OVER" if text > limit else f"{limit - text:>5} B"
        print(f"{chip:<13} {name:<34} {text:>6} {limit:>7}  {flag}")
        over += [(chip, n, t, l) for n, t, l in rows if t > l]
    total = sum(len(rows) for rows in measured.values())
    print(f"\n{total} images across {len(measured)} chips")
    if over:
        print("\nOVER BUDGET:")
        for chip, name, text, limit in over:
            print(f"  {chip} {name}: {text} > {limit}")
        return 1
    tightest = min(((chip, n, t, l) for chip, rows in measured.items() for n, t, l in rows),
                   key=lambda row: row[3] - row[2])
    chip, name, text, limit = tightest
    print(f"tightest fit: {chip} {name} — {text} of {limit}, {limit - text} B spare")
    return 0


def cmd_check_readme(args) -> int:
    """The README's per-chip table, against the stock and autobaud builds."""
    readme = (ROOT / "pureboot" / "README.md").read_text()
    measured = collect()
    rows = re.findall(r"^\|\s*(AT\w+[^|]*?)\s*\|[^|]*\|[^|]*\|[^|]*\|\s*(\d+) B\s*\|\s*(\d+) B\s*\|$",
                      readme, re.MULTILINE)
    if not rows:
        sys.exit("no size table found in pureboot/README.md")
    bad = skipped = 0
    for chips, stock_doc, auto_doc in rows:
        # "ATmega48, 48A, 48P, 48PA †" — the first name is the family's base.
        chip = re.sub(r"[^a-z0-9]", "", chips.split(",")[0].strip().lower())
        built = {name: text for name, text, _ in measured.get(chip, [])}
        for target, documented in (("pureboot", stock_doc), ("pureboot_autobaud", auto_doc)):
            if target not in built:
                skipped += 1
                continue
            if built[target] != int(documented):
                print(f"  {chip:<12} {target:<18} README says {documented} B, built is {built[target]} B")
                bad += 1
    if bad:
        print(f"\n{bad} row(s) stale — update pureboot/README.md")
        return 1
    print(f"README size table matches every built image ({len(rows)} rows"
          + (f", {skipped} not built" if skipped else "") + ")")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    subs = parser.add_subparsers(dest="cmd", required=True)
    subs.add_parser("max", help="largest image per chip, and anything over budget")
    subs.add_parser("check-readme", help="the README's size table against what is built")
    args = parser.parse_args()
    return {"max": cmd_max, "check-readme": cmd_check_readme}[args.cmd](args)


if __name__ == "__main__":
    raise SystemExit(main())