diff --git a/wizard-work/beam/test_topo_beam.py b/wizard-work/beam/test_topo_beam.py new file mode 100644 index 0000000..f8d54d1 --- /dev/null +++ b/wizard-work/beam/test_topo_beam.py @@ -0,0 +1,41 @@ +import sys +sys.path.insert(0, "/root/.hermes/wizard-outputs/beam") +from topo import topo_sort, has_cycle, CycleError + +def test_empty(): + assert topo_sort([]) == [] + +def test_single(): + assert topo_sort([("a", [])]) == ["a"] + +def test_linear(): + assert topo_sort([("b", ["a"]), ("c", ["b"])]) == ["a", "b", "c"] + +def test_diamond(): + r = topo_sort([("d", ["b", "c"]), ("b", ["a"]), ("c", ["a"])]) + assert r[0] == "a" and r[-1] == "d" + assert r.index("b") < r.index("c") + +def test_stability(): + assert topo_sort([("x", []), ("y", []), ("z", [])]) == ["x", "y", "z"] + +def test_cycle(): + try: + topo_sort([("a", ["b"]), ("b", ["a"])]) + assert False + except CycleError as e: + assert set(e.remaining) == {"a", "b"} + +def test_self_loop(): + assert has_cycle([("a", ["a"])]) + +def test_no_cycle(): + assert not has_cycle([("a", []), ("b", ["a"])]) + +def test_dup_preds(): + assert topo_sort([("b", ["a", "a"])]) == ["a", "b"] + +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.py b/wizard-work/beam/topo.py new file mode 100644 index 0000000..16c7f71 --- /dev/null +++ b/wizard-work/beam/topo.py @@ -0,0 +1,71 @@ +"""Deterministic topological sort (Kahn's algorithm) with cycle detection. + +Nodes are arbitrary hashable values. Input is an iterable of (node, predecessors) +pairs where predecessors are the nodes that must come before node. Ordering +among independent nodes is stable by input order of first appearance (earliest +discovered zero-indegree / ready node wins). + +Returns a list of nodes in topological order. Raises CycleError if the graph +contains a cycle, listing the remaining (unsorted) nodes. +""" +from __future__ import annotations + +import heapq +from collections import OrderedDict + + +class CycleError(Exception): + def __init__(self, remaining): + self.remaining = tuple(remaining) + super().__init__("cycle detected among: %r" % (self.remaining,)) + + +def topo_sort(edges): + order = OrderedDict() # node -> discovery index + preds = {} + succs = {} + indeg = {} + + idx = 0 + def discover(n): + nonlocal idx + if n not in order: + order[n] = idx + idx += 1 + preds[n] = set() + succs[n] = [] + indeg[n] = 0 + + for node, predecessors in edges: + discover(node) + for p in predecessors: + discover(p) + if p not in preds[node]: + preds[node].add(p) + succs[p].append(node) + indeg[node] += 1 + + # Kahn with a min-heap keyed by discovery index for global stability. + heap = [(order[n], n) for n in order if indeg[n] == 0] + heapq.heapify(heap) + result = [] + while heap: + _, n = heapq.heappop(heap) + result.append(n) + for s in succs[n]: + indeg[s] -= 1 + if indeg[s] == 0: + heapq.heappush(heap, (order[s], s)) + + if len(result) != len(order): + remaining = [n for n in order if n not in set(result)] + raise CycleError(remaining) + return result + + +def has_cycle(edges): + try: + topo_sort(edges) + return False + except CycleError: + return True