Files
troupe-standup/wizard-work/beam/toposort.py
T

194 lines
6.3 KiB
Python

"""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)
class _Sortable:
__slots__ = ("x", "tn")
def __init__(self, x):
self.x = x
self.tn = type(x).__name__
def __lt__(self, other):
if not isinstance(other, _Sortable):
return NotImplemented
if self.tn != other.tn:
return self.tn < other.tn
try:
return self.x < other.x
except Exception:
rx, ry = repr(self.x), repr(other.x)
if rx != ry:
return rx < ry
return id(self.x) < id(other.x)
def __eq__(self, other):
return isinstance(other, _Sortable) and self.x == other.x
def _key(x):
return _Sortable(x)