194 lines
6.3 KiB
Python
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) |