beam: topo_levels.py Kahn parallel-level topological sort + tests

wizard-work/beam/test_topo_levels.py

@@ -0,0 +1,63 @@
+import sys
+sys.path.insert(0, "/root/.hermes/wizard-outputs/beam")
+from topo_levels import topo_levels, flatten, CycleError
+
+def test_basic():
+    nodes = ["a","b","c","d"]
+    edges = [("a","c"),("b","c"),("c","d")]
+    levels = topo_levels(nodes, edges)
+    assert levels == [["a","b"],["c"],["d"]], levels
+
+def test_disconnected():
+    nodes = [1,2,3]
+    assert topo_levels(nodes, []) == [[1,2,3]]
+
+def test_cycle():
+    nodes = ["a","b","c"]
+    edges = [("a","b"),("b","c"),("c","a")]
+    try:
+        topo_levels(nodes, edges)
+    except CycleError as e:
+        assert set(e.remaining) == {"a","b","c"}
+    else:
+        raise AssertionError("expected CycleError")
+
+def test_unknown_node():
+    try:
+        topo_levels(["a"], [("a","z")])
+    except KeyError:
+        pass
+    else:
+        raise AssertionError("expected KeyError")
+
+def test_single_node_no_edges():
+    assert topo_levels(["x"], []) == [["x"]]
+
+def test_diamond():
+    nodes = ["a","b","c","d"]
+    edges = [("a","b"),("a","c"),("b","d"),("c","d")]
+    levels = topo_levels(nodes, edges)
+    assert levels == [["a"],["b","c"],["d"]], levels
+
+def test_determinism():
+    nodes = list("zyxwv")
+    edges = [("z","w"),("y","w")]
+    r1 = topo_levels(nodes, edges)
+    r2 = topo_levels(list(reversed(nodes)), edges)
+    assert r1 == r2, (r1, r2)
+
+def test_flatten_roundtrip():
+    nodes = ["a","b","c","d","e"]
+    edges = [("a","c"),("b","c"),("c","d"),("c","e")]
+    levels = topo_levels(nodes, edges)
+    flat = flatten(levels)
+    assert sorted(flat) == sorted(nodes)
+    # every edge prerequisite before dependent
+    pos = {n:i for i,n in enumerate(flat)}
+    for u,v in edges:
+        assert pos[u] < pos[v]
+
+if __name__ == "__main__":
+    for name in [n for n in globals() if n.startswith("test_")]:
+        globals()[name]()
+    print("ALL TESTS PASSED")

wizard-work/beam/topo_levels.py

@@ -0,0 +1,97 @@
+"""Topological ordering with parallel execution levels.
+
+Kahn's algorithm, but instead of emitting one node at a time we emit
+*levels*: each level contains every node whose in-degree has just hit zero,
+so all nodes in a level can run in parallel. Deterministic tie-break by node
+key (string) so output is reproducible. Raises CycleError listing the
+remaining (unresolved) nodes if the graph has a cycle.
+"""
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from typing import Any, Iterable, List, Tuple
+
+
+class CycleError(ValueError):
+    def __init__(self, remaining: List[Any]):
+        super().__init__(f"graph contains a cycle; unresolved nodes: {remaining!r}")
+        self.remaining = remaining
+
+
+def topo_levels(
+    nodes: Iterable[Any],
+    edges: Iterable[Tuple[Any, Any]],
+) -> List[List[Any]]:
+    """Return levels of nodes in dependency order.
+
+    `edges` are (prerequisite, dependent) pairs: prerequisite must come before
+    dependent. Nodes with no edges still appear (in their own level 0 group).
+    """
+    node_set = list(dict.fromkeys(nodes))  # preserve insertion order, dedupe
+    index = {n: i for i, n in enumerate(node_set)}
+
+    # validate edges reference known nodes
+    for a, b in edges:
+        if a not in index:
+            raise KeyError(f"edge references unknown node: {a!r}")
+        if b not in index:
+            raise KeyError(f"edge references unknown node: {b!r}")
+
+    indeg = {n: 0 for n in node_set}
+    adj = defaultdict(list)
+    for a, b in edges:
+        adj[a].append(b)
+        indeg[b] += 1
+
+    # deterministic: process ready queue sorted by key
+    ready = sorted([n for n in node_set if indeg[n] == 0], key=_key)
+    levels: List[List[Any]] = []
+    processed = 0
+    # Kahn in waves
+    current = ready
+    while current:
+        level = sorted(current, key=_key)
+        levels.append(level)
+        next_ready: List[Any] = []
+        for n in level:
+            processed += 1
+            for m in adj[n]:
+                indeg[m] -= 1
+                if indeg[m] == 0:
+                    next_ready.append(m)
+        # dedupe next_ready (a node can be pushed twice if two prereqs resolve same wave? no, indeg hits 0 once)
+        current = next_ready
+
+    if processed != len(node_set):
+        remaining = [n for n in node_set if indeg[n] > 0]
+        raise CycleError(sorted(remaining, key=_key))
+
+    # flatten check: every node appears exactly once
+    seen = []
+    for lvl in levels:
+        seen.extend(lvl)
+    assert len(seen) == len(node_set) == len(set(seen)), "internal invariant violated"
+    return levels
+
+
+def _key(n: Any):
+    return (type(n).__name__, n) if _comparable(n) else (type(n).__name__, repr(n))
+
+
+def _comparable(n: Any) -> bool:
+    # Always comparable within the same concrete type for primitives we care about.
+    return isinstance(n, (int, float, str, bytes))
+
+
+def flatten(levels: List[List[Any]]) -> List[Any]:
+    out = []
+    for lvl in levels:
+        out.extend(lvl)
+    return out
+
+
+if __name__ == "__main__":
+    # smoke
+    nodes = ["a", "b", "c", "d"]
+    edges = [("a", "c"), ("b", "c"), ("c", "d")]
+    print(topo_levels(nodes, edges))