beam: fix toposort _key partial-order bug (sets treated as unorderable)
This commit is contained in:
beam committed
1 parent
e9d3283242
commit
b646cbb497
1 file changed
+180
@@ -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)
|
||||
Reference in new issue
Block a user