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

"""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).

The declared contract is `Iterable`, not "re-iterable": a dependency set may
arrive as a generator, a map object, a filter, or any other one-shot iterator.
Every dependency sequence is therefore materialised exactly once, up front, and
all later passes read that single materialised copy. Iterating `graph.values()`
twice would silently exhaust a generator, leaving the node with an empty
dependency set, in-degree 0 and a wrong-but-successful ordering.
"""
from __future__ import annotations

from collections import defaultdict
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:
    # Materialise each dependency iterable once. `None` means "no dependencies".
    # This is the only place any dependency sequence is consumed.
    deps_src: dict[Any, tuple[Any, ...]] = {
        n: () if ds is None else tuple(ds) for n, ds in graph.items()
    }

    # Normalise: every referenced node exists as a key, even if it has no entry.
    nodes: set[Any] = set(graph.keys())
    for ds in deps_src.values():
        nodes.update(ds)

    # adjacency: node -> set of dependencies (a self-dependency is a 1-cycle)
    deps_of: dict[Any, set[Any]] = {n: set() for n in nodes}
    for n, ds in deps_src.items():
        deps_of[n].update(ds)

    # 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)

    # A cycle exists. Everything Kahn could not emit is in a cycle or reaches
    # one, so a DFS restricted to those nodes is guaranteed to find a back edge.
    emitted = set(out)
    unemitted = [n for n in nodes if n not in emitted]

    color: dict[Any, int] = {n: 0 for n in nodes}  # 0 unvisited, 1 visiting, 2 done
    parent: dict[Any, Any] = {}

    def find_back_edge(u: Any) -> tuple[Any, Any] | None:
        """DFS from u; return the (node, ancestor) back edge that closes a cycle."""
        color[u] = 1
        for v in sorted(deps_of[u], key=_key):
            state = color.get(v, 0)
            if state == 0:
                parent[v] = u
                found = find_back_edge(v)
                if found is not None:
                    return found
            elif state == 1:
                return (u, v)
        color[u] = 2
        return None

    back: tuple[Any, Any] | None = None
    for s in sorted(unemitted, key=_key):
        if color.get(s, 0) == 0:
            parent[s] = None
            back = find_back_edge(s)
            if back is not None:
                break

    if back is None:
        # Unreachable for a real cycle; never report a silent success instead.
        return TopoResult(sorted=None, cycle=list(unemitted[:1]))

    back_u, back_v = back
    if back_u == back_v:
        # self-loop: a cycle of length 1
        return TopoResult(sorted=None, cycle=[back_v])

    # back_u depends on back_v, and parent[x] = y means "y depends on x", so
    # walking parents from back_u reaches back_v. Reversed, that walk is the
    # cycle in dependency order: every node is followed by one of its own deps,
    # and the last node depends on the first, closing the loop.
    path: list[Any] = [back_u]
    cur = parent.get(back_u)
    while cur is not None and cur != back_v:
        path.append(cur)
        cur = parent.get(cur)
    if cur != back_v:
        # Parent chain broken (should not happen); degrade to a live node
        # rather than inventing an edge that does not exist.
        return TopoResult(sorted=None, cycle=list(unemitted[:1]))
    path.append(back_v)
    path.reverse()
    return TopoResult(sorted=None, cycle=path)


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)