1 files changed, 149 insertions, 0 deletions
diff --git a/experiments/graph_spike.py b/experiments/graph_spike.py
new file mode 100644
index 0000000..547c829
--- /dev/null
+++ b/experiments/graph_spike.py
@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+"""Render a CFG with idatui.graph, offline, to stdout.
+
+The layout engine now lives in ``idatui/graph.py`` (this started life as the
+spike that proved it). What survives here is the offline harness: feed it a
+corpus from ``cfg_dump.py`` and look at a function, or lay out the whole corpus
+and check the cost. No IDA, no Textual, no worker -- so it iterates in
+milliseconds when you're changing layout heuristics.
+
+    /usr/bin/python3 experiments/cfg_dump.py targets/echo -o /tmp/cfg-echo.json
+    python3 experiments/graph_spike.py /tmp/cfg-echo.json --func sub_61D0
+    python3 experiments/graph_spike.py /tmp/cfg-echo.json --stats
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from idatui import graph as G  # noqa: E402
+
+COLOR = {
+    G.E_UNCOND: "\033[38;5;39m", G.E_TRUE: "\033[38;5;40m",
+    G.E_FALSE: "\033[38;5;203m", G.E_SWITCH: "\033[38;5;178m",
+    G.E_BACK: "\033[38;5;135m",
+}
+DIM, RESET = "\033[38;5;244m", "\033[0m"
+PAD = 1
+
+
+def build(rec: dict, max_lines: int):
+    """(blocks, sizer, texts) for one cfg_dump record."""
+    texts: dict[int, list[str]] = {}
+    for b in rec["blocks"]:
+        lines = list(b["lines"])
+        if max_lines and len(lines) > max_lines:
+            lines = lines[:max_lines - 1] + [f"... {len(b['lines']) - max_lines + 1} more"]
+        texts[b["id"]] = lines
+    blocks = [G.Block(id=b["id"], start=b["start"], end=b["end"],
+                      succs=[(d, k) for d, k in b["succs"]])
+              for b in rec["blocks"]]
+
+    def sizer(b: G.Block) -> tuple[int, int]:
+        lines = texts[b.id]
+        label = f"loc_{b.start:X}"
+        widest = max([len(label) + 4] + [len(l) for l in lines] or [4])
+        return (widest + 2 * PAD + 2, max(len(lines), 1) + 2)
+
+    return blocks, sizer, texts
+
+
+def render(lay: G.Layout, texts: dict[int, list[str]], color: bool) -> str:
+    rows = []
+    for row in range(lay.height):
+        cells: dict[int, tuple[str, str]] = {}
+        for col, (ch, kind, _eid) in lay.painting.cells_at_row(
+                row, 0, lay.width).items():
+            cells[col] = (ch, COLOR.get(kind, ""))
+        for n in lay.nodes_at_row(row):
+            label = f"loc_{n.block.start:X}" if n.block else ""
+            if row == n.y:
+                cells[n.x] = (G.BOX["tl"], DIM)
+                for i in range(1, n.w - 1):
+                    cells[n.x + i] = (G.BOX["h"], DIM)
+                cells[n.x + n.w - 1] = (G.BOX["tr"], DIM)
+                tag = f" {label} "
+                if len(tag) <= n.w - 4:
+                    for k, c in enumerate(tag):
+                        cells[n.x + 2 + k] = (c, DIM)
+                if n.block is not None and n.block.selfloop:
+                    cells[n.x + n.w - 2] = ("\u21ba", COLOR[G.E_BACK])
+            elif row == n.y + n.h - 1:
+                cells[n.x] = (G.BOX["bl"], DIM)
+                for i in range(1, n.w - 1):
+                    cells[n.x + i] = (G.BOX["h"], DIM)
+                cells[n.x + n.w - 1] = (G.BOX["br"], DIM)
+            else:
+                cells[n.x] = (G.BOX["v"], DIM)
+                cells[n.x + n.w - 1] = (G.BOX["v"], DIM)
+                for i in range(1, n.w - 1):
+                    cells[n.x + i] = (" ", "")
+                lines = texts.get(n.id, [])
+                i = row - n.y - 1
+                if 0 <= i < len(lines):
+                    for k, c in enumerate(lines[i]):
+                        cells[n.x + 1 + PAD + k] = (c, "")
+        line, cur, last = [], "", -1
+        for col in sorted(cells):
+            ch, st = cells[col]
+            line.append(" " * (col - last - 1))
+            if color and st != cur:
+                line.append(st or RESET)
+                cur = st
+            line.append(ch)
+            last = col
+        if color and cur:
+            line.append(RESET)
+        rows.append("".join(line))
+    return "\n".join(rows)
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("corpus")
+    ap.add_argument("--func", help="function name (default: the smallest)")
+    ap.add_argument("--stats", action="store_true", help="lay out the whole corpus")
+    ap.add_argument("--max-lines", type=int, default=8,
+                    help="collapse blocks longer than this (0 = never)")
+    ap.add_argument("--no-color", action="store_true")
+    args = ap.parse_args()
+
+    recs = json.load(open(args.corpus))
+    if args.stats:
+        print(f"{'blocks':>7} {'nodes':>6} {'dummy':>6} {'layer':>6} "
+              f"{'canvas':>12} {'ms':>8}  name")
+        tot = 0.0
+        for r in sorted(recs, key=lambda r: len(r["blocks"])):
+            blocks, sizer, _ = build(r, args.max_lines)
+            lay = G.layout(blocks, sizer)
+            s = lay.stats
+            tot += s["ms"]
+            print(f"{s['blocks']:>7} {s['nodes']:>6} {s['dummies']:>6} "
+                  f"{s['layers']:>6} {lay.width:>5}x{lay.height:<6} "
+                  f"{s['ms']:>8.1f}  {r['name']}")
+        print(f"total {tot:.0f} ms over {len(recs)} functions")
+        return 0
+
+    if args.func:
+        rec = next((r for r in recs if r["name"] == args.func), None)
+        if rec is None:
+            print("no such function; have: "
+                  f"{', '.join(r['name'] for r in recs[:20])}", file=sys.stderr)
+            return 1
+    else:
+        rec = min(recs, key=lambda r: len(r["blocks"]))
+
+    blocks, sizer, texts = build(rec, args.max_lines)
+    lay = G.layout(blocks, sizer)
+    print(render(lay, texts, color=not args.no_color))
+    print(f"\n{rec['name']}: {lay.stats}  canvas {lay.width}x{lay.height}",
+          file=sys.stderr)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())