From dbb2dae51f4bedb2dfd09ced5de97e701c97fb90 Mon Sep 17 00:00:00 2001 From: Jonas Volman Date: Sat, 18 Jul 2026 10:02:46 -0700 Subject: [PATCH] Add exact donor shape atlas --- tools/shape_atlas.py | 168 ++++++++++++++++++++++++++++++++++++++ tools/test_shape_atlas.py | 63 ++++++++++++++ 2 files changed, 231 insertions(+) create mode 100644 tools/shape_atlas.py create mode 100644 tools/test_shape_atlas.py diff --git a/tools/shape_atlas.py b/tools/shape_atlas.py new file mode 100644 index 00000000..20bd009e --- /dev/null +++ b/tools/shape_atlas.py @@ -0,0 +1,168 @@ +"""Find reusable target-code shapes without granting exact-match credit. + +The atlas is a scheduling aid. It groups target functions only when their +relocation-masked bytes, padded size, relocation count, and relocation-type +sequence agree. An exact donor means that a proven source template exists; +the recipient must still pass the ordinary strict comparator after it is +implemented in its own translation unit. +""" + +import argparse +import json +from collections import defaultdict +from pathlib import Path + +from .coff_compare import CoffError, load, section_info + + +def _shape_key(measurement): + return ( + measurement["size"], + measurement["relocation_count"], + measurement["normalized_sha256"], + tuple(relocation["type"] for relocation in measurement["relocations"]), + ) + + +def build_atlas(catalog, minimum_bulk_members=3): + """Return transfer and unsolved-bulk groups from catalog entries.""" + groups = defaultdict(list) + for entry in catalog: + groups[tuple(entry["shape_key"])].append(entry) + + transfers = [] + bulk = [] + for shape_key, members in groups.items(): + exact = [member for member in members if member["status"] == "exact"] + nonexact = [member for member in members if member["status"] != "exact"] + if exact and nonexact: + transfers.append({ + "shape": _shape_record(shape_key), + "donors": sorted(exact, key=_member_sort_key), + "recipients": sorted(nonexact, key=_member_sort_key), + "recipient_bytes": sum(member["size"] for member in nonexact), + }) + elif not exact and len(nonexact) >= minimum_bulk_members: + bulk.append({ + "shape": _shape_record(shape_key), + "members": sorted(nonexact, key=_member_sort_key), + "total_bytes": sum(member["size"] for member in nonexact), + }) + + transfers.sort(key=lambda group: (-group["recipient_bytes"], group["shape"]["sha256"])) + bulk.sort(key=lambda group: (-group["total_bytes"], group["shape"]["sha256"])) + return transfers, bulk + + +def _shape_record(shape_key): + size, relocation_count, sha256, relocation_types = shape_key + return { + "size": size, + "relocation_count": relocation_count, + "sha256": sha256, + "relocation_types": list(relocation_types), + } + + +def _member_sort_key(member): + return member["unit"], member["function"] + + +def analyze(project_root, report_path, semantic_path, config_path): + project_root = Path(project_root) + report = json.loads(Path(report_path).read_text(encoding="utf-8")) + semantic = json.loads(Path(semantic_path).read_text(encoding="utf-8")) + config = json.loads(Path(config_path).read_text(encoding="utf-8")) + config_units = {unit["name"]: unit for unit in config.get("units", [])} + exact = { + (entry["unit"], entry["function"]) + for entry in semantic.get("accepted_ledger", []) + } + catalog = [] + errors = [] + + for report_unit in report.get("units", []): + unit_name = report_unit["name"] + categories = report_unit.get("metadata", {}).get("progress_categories") or [] + if "halobetacache" not in categories: + continue + config_unit = config_units.get(unit_name) + if not config_unit: + continue + try: + target = load(project_root / config_unit["target_path"]) + base_path = config_unit.get("base_path") + base = load(project_root / base_path) if base_path else None + except (CoffError, OSError, KeyError) as error: + errors.append({"unit": unit_name, "error": str(error)}) + continue + + for function in report_unit.get("functions", []): + function_name = function["name"] + if function_name.startswith("$"): + continue + try: + measurement = section_info(target, function_name) + except CoffError: + continue + if base is None: + base_present = False + else: + try: + section_info(base, function_name) + base_present = True + except CoffError: + base_present = False + status = "exact" if (unit_name, function_name) in exact else ( + "implemented" if base_present else "absent" + ) + catalog.append({ + "unit": unit_name, + "function": function_name, + "status": status, + "size": measurement["size"], + "relocation_count": measurement["relocation_count"], + "normalized_sha256": measurement["normalized_sha256"], + "shape_key": list(_shape_key(measurement)), + }) + + transfers, bulk = build_atlas(catalog) + return { + "summary": { + "functions": len(catalog), + "distinct_shapes": len({tuple(entry["shape_key"]) for entry in catalog}), + "transfer_groups": len(transfers), + "transfer_recipients": sum(len(group["recipients"]) for group in transfers), + "transfer_bytes": sum(group["recipient_bytes"] for group in transfers), + "bulk_groups": len(bulk), + "bulk_bytes": sum(group["total_bytes"] for group in bulk), + "errors": len(errors), + }, + "transfer": transfers, + "bulk": bulk, + "errors": errors, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--project", type=Path, default=Path.cwd()) + parser.add_argument("--report", type=Path, default=Path("build/report.json")) + parser.add_argument("--semantic", type=Path, default=Path("build/semantic_report.json")) + parser.add_argument("--config", type=Path, default=Path("objdiff.json")) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + root = args.project.resolve() + result = analyze(root, root / args.report, root / args.semantic, root / args.config) + serialized = json.dumps(result, indent=2) + "\n" + if args.output: + output = root / args.output + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(serialized, encoding="utf-8") + print(serialized, end="") + if result["errors"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/test_shape_atlas.py b/tools/test_shape_atlas.py new file mode 100644 index 00000000..a152d5f2 --- /dev/null +++ b/tools/test_shape_atlas.py @@ -0,0 +1,63 @@ +import unittest + +from tools.shape_atlas import build_atlas + + +def entry(unit, function, status, size=32, relocation_types=(20,), sha256="shape"): + return { + "unit": unit, + "function": function, + "status": status, + "size": size, + "shape_key": [size, len(relocation_types), sha256, tuple(relocation_types)], + } + + +class ShapeAtlasTests(unittest.TestCase): + def test_exact_donor_creates_transfer_group(self): + transfers, bulk = build_atlas([ + entry("source/a", "_donor", "exact"), + entry("source/b", "_recipient", "absent"), + ]) + self.assertEqual(len(transfers), 1) + self.assertEqual(transfers[0]["donors"][0]["function"], "_donor") + self.assertEqual(transfers[0]["recipients"][0]["function"], "_recipient") + self.assertEqual(bulk, []) + + def test_three_unsolved_members_create_bulk_group(self): + transfers, bulk = build_atlas([ + entry("source/a", "_a", "absent"), + entry("source/b", "_b", "implemented"), + entry("source/c", "_c", "absent"), + ]) + self.assertEqual(transfers, []) + self.assertEqual(len(bulk), 1) + self.assertEqual(bulk[0]["total_bytes"], 96) + + def test_padding_difference_does_not_transfer(self): + transfers, bulk = build_atlas([ + entry("source/a", "_donor", "exact", size=544), + entry("source/b", "_near", "absent", size=543), + ]) + self.assertEqual(transfers, []) + self.assertEqual(bulk, []) + + def test_relocation_type_difference_does_not_transfer(self): + transfers, bulk = build_atlas([ + entry("source/a", "_donor", "exact", relocation_types=(20,)), + entry("source/b", "_near", "absent", relocation_types=(6,)), + ]) + self.assertEqual(transfers, []) + self.assertEqual(bulk, []) + + def test_two_unsolved_members_are_not_bulk(self): + transfers, bulk = build_atlas([ + entry("source/a", "_a", "absent"), + entry("source/b", "_b", "implemented"), + ]) + self.assertEqual(transfers, []) + self.assertEqual(bulk, []) + + +if __name__ == "__main__": + unittest.main()