diff --git a/wizard-work/beam/toposort.py b/wizard-work/beam/toposort.py new file mode 100644 index 0000000..d377185 --- /dev/null +++ b/wizard-work/beam/toposort.py @@ -0,0 +1,180 @@ +"""Topological sort with deterministic ordering and cycle-path reporting. + +Public API: + topo_sort(graph) -> TopoResult +where graph is a mapping node -> iterable of dependencies (edges node -> dep). + +A node's dependencies must appear before it in the output. Ties broken by +node's natural ordering (so results are deterministic). If a cycle exists, +result.sorted is None and result.cycle is a list of nodes forming the cycle +(in dependency order, last -> first closing the loop). +""" +from __future__ import annotations + +from collections import defaultdict, deque +from dataclasses import dataclass +from typing import Any, Iterable, Mapping + + +@dataclass(frozen=True) +class TopoResult: + sorted: list[Any] | None + cycle: list[Any] | None + + +def topo_sort(graph: Mapping[Any, Iterable[Any]]) -> TopoResult: + # Normalize: ensure every referenced node exists as a key. + nodes: set[Any] = set(graph.keys()) + for deps in graph.values(): + for d in deps: + nodes.add(d) + + # adjacency: node -> set of dependencies + deps_of: dict[Any, set[Any]] = {n: set() for n in nodes} + for n, ds in graph.items(): + for d in ds: + if d == n: + # self-loop is a cycle of length 1 + deps_of[n].add(d) + else: + deps_of[n].add(d) + + # Kahn's algorithm with a sorted ready queue for determinism. + # in-degree = number of unprocessed dependencies + in_deg: dict[Any, int] = {n: len(deps_of[n]) for n in nodes} + # reverse adjacency: dep -> dependents + dependents: dict[Any, list[Any]] = defaultdict(list) + for n in nodes: + for d in deps_of[n]: + dependents[d].append(n) + + # ready set as a min-heap via sorted list (nodes sortable; if not, fall back) + try: + ready = sorted((n for n in nodes if in_deg[n] == 0), key=_key) + except TypeError: + ready = [n for n in nodes if in_deg[n] == 0] + out: list[Any] = [] + # We'll use a list as a priority queue by re-sorting on pop; fine for tests. + while ready: + # pop smallest + ready.sort(key=_key) + cur = ready.pop(0) + out.append(cur) + for dep in sorted(dependents[cur], key=_key): + in_deg[dep] -= 1 + if in_deg[dep] == 0: + ready.append(dep) + + if len(out) == len(nodes): + return TopoResult(sorted=out, cycle=None) + + # Cycle exists. Reconstruct a cycle path. + # Nodes not emitted are in cycles. Pick the smallest unemitted node. + unemitted = [n for n in nodes if n not in set(out)] + try: + start = min(unemitted) + except TypeError: + start = unemitted[0] + + # Walk dependency edges among unemitted nodes to find a cycle. + # Use DFS with color marking. + color: dict[Any, int] = {n: 0 for n in nodes} # 0 unvisited,1 visiting,2 done + parent: dict[Any, Any] = {} + + def dfs(u): + color[u] = 1 + for v in sorted(deps_of[u], key=_key): + if color.get(v, 0) == 0: + parent[v] = u + r = dfs(v) + if r is not None: + return r + elif color.get(v, 0) == 1: + # found cycle: v is ancestor of u + # reconstruct v -> ... -> u -> v + cycle_nodes = [v] + cur = u + while cur != v and cur is not None: + cycle_nodes.append(cur) + cur = parent.get(cur) + if cur is None: + break + cycle_nodes.reverse() + # ensure v at front already; make path v...u + # cycle_nodes currently [v, ..., u]? Let's rebuild cleanly. + path = [v] + cur = u + chain = [u] + while cur != v and cur is not None: + cur = parent.get(cur) + if cur is not None and cur != v: + chain.append(cur) + chain.reverse() + # chain now from v's child ... u? Reconstruct differently below. + return ("FOUND", v, u) + color[u] = 2 + return None + + res = None + for s in sorted(unemitted, key=_key): + if color[s] == 0: + parent[s] = None # type: ignore + r = dfs(s) + if r is not None: + res = r + break + + # Reconstruct cycle path robustly using parent pointers from the found edge. + # Find the back edge (u -> v) where v is an ancestor. + back_u = back_v = None + # Re-run a focused search to capture back edge. + color2: dict[Any, int] = {n: 0 for n in nodes} + parent2: dict[Any, Any] = {} + + def dfs2(u): + nonlocal back_u, back_v + color2[u] = 1 + for v in sorted(deps_of[u], key=_key): + if color2.get(v, 0) == 0: + parent2[v] = u + r = dfs2(v) + if r: + return r + elif color2.get(v, 0) == 1: + back_u, back_v = u, v + return True + color2[u] = 2 + return False + + for s in sorted(unemitted, key=_key): + if color2[s] == 0: + parent2[s] = None # type: ignore + if dfs2(s): + break + + # Build path from back_v to back_u via parent2, then close. + path = [] + cur = back_u + while cur is not None and cur != back_v: + path.append(cur) + cur = parent2.get(cur) + if cur == back_v: + path.append(back_v) + path.reverse() + # path: back_v ... back_u + cycle = path # back_v -> ... -> back_u, and edge back_u -> back_v closes + else: + cycle = [back_v, back_u] if back_v is not None else unemitted[:1] + + return TopoResult(sorted=None, cycle=cycle) + + +def _key(x): + tn = type(x).__name__ + if isinstance(x, (set, frozenset)): + return (1, tn, repr(x), id(x)) + try: + _ = x < x + except Exception: + return (1, tn, repr(x), id(x)) + return (0, tn, x) \ No newline at end of file