diff --git a/wizard-work/beam/test_topo_levels.py b/wizard-work/beam/test_topo_levels.py new file mode 100644 index 0000000..f91b110 --- /dev/null +++ b/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") diff --git a/wizard-work/beam/topo_levels.py b/wizard-work/beam/topo_levels.py new file mode 100644 index 0000000..6a7e073 --- /dev/null +++ b/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))