top-sort-asm/topsort_asm/toposort.py

"""Topological sorting, three ways.

All functions take a graph as ``n`` plus ``succ``: a mapping from each node
``0..n-1`` to the set of its successors. They know nothing about assembly.

* :func:`kahn` -- repeatedly emit a node with no remaining predecessors.
  The set of such nodes (the *ready set*) is exactly the freedom a scheduler
  has; the ``choose`` hook decides which ready node goes next.
* :func:`dfs_order` -- reverse DFS post-order; finds a cycle if there is one.
* :func:`all_orders` / :func:`count_orders` -- every valid order, for small
  graphs, to show how much freedom the dependences actually leave.
"""

from functools import lru_cache
from typing import Callable, Dict, Iterator, List, Mapping, Optional, Sequence, Set


class CycleError(ValueError):
    def __init__(self, cycle: Sequence[int]):
        self.cycle = list(cycle)
        super().__init__("cycle: " + " -> ".join(map(str, self.cycle)))


Succ = Mapping[int, Set[int]]


def indegrees(n: int, succ: Succ) -> List[int]:
    deg = [0] * n
    for u in range(n):
        for v in succ[u]:
            deg[v] += 1
    return deg


def kahn(n: int, succ: Succ,
         choose: Optional[Callable[[List[int]], int]] = None,
         trace: Optional[List[dict]] = None) -> List[int]:
    """Kahn's algorithm.

    ``choose`` picks one node from the sorted ready list (default: the
    smallest, which for a dependence DAG reproduces program order). If
    ``trace`` is given, one dict per step is appended to it describing the
    ready set, the choice, and which nodes it released.
    """
    choose = choose or (lambda ready: ready[0])
    deg = indegrees(n, succ)
    ready = sorted(v for v in range(n) if deg[v] == 0)
    order: List[int] = []
    while ready:
        u = choose(ready)
        ready.remove(u)
        order.append(u)
        released = []
        for v in sorted(succ[u]):
            deg[v] -= 1
            if deg[v] == 0:
                released.append(v)
        if trace is not None:
            trace.append({"ready": sorted(ready + [u]), "chose": u,
                          "released": released})
        ready = sorted(ready + released)
    if len(order) != n:
        # Every leftover node still has a predecessor among the leftovers,
        # so walking predecessors must eventually repeat a node.
        left = {v for v in range(n) if deg[v] > 0}
        raise CycleError(_find_cycle(left, succ))
    return order


def _find_cycle(nodes: Set[int], succ: Succ) -> List[int]:
    pred: Dict[int, int] = {}
    for u in nodes:
        for v in succ[u]:
            if v in nodes:
                pred.setdefault(v, u)
    path: List[int] = []
    seen: Dict[int, int] = {}
    v = min(nodes)
    while v not in seen:
        seen[v] = len(path)
        path.append(v)
        v = pred[v]
    cyc = path[seen[v]:][::-1]
    return cyc + [cyc[0]]


def dfs_order(n: int, succ: Succ) -> List[int]:
    """Reverse post-order of a DFS. Raises CycleError on a back edge."""
    WHITE, GREY, BLACK = 0, 1, 2
    color = [WHITE] * n
    post: List[int] = []
    for root in range(n):
        if color[root] != WHITE:
            continue
        # Iterative DFS so long blocks don't hit the recursion limit.
        stack = [(root, iter(sorted(succ[root])))]
        path = [root]
        color[root] = GREY
        while stack:
            u, it = stack[-1]
            v = next(it, None)
            if v is None:
                stack.pop()
                path.pop()
                color[u] = BLACK
                post.append(u)
            elif color[v] == GREY:
                cyc = path[path.index(v):]
                raise CycleError(cyc + [v])
            elif color[v] == WHITE:
                color[v] = GREY
                path.append(v)
                stack.append((v, iter(sorted(succ[v]))))
    return post[::-1]


def is_topological(order: Sequence[int], n: int, succ: Succ) -> bool:
    if sorted(order) != list(range(n)):
        return False
    pos = {v: i for i, v in enumerate(order)}
    return all(pos[u] < pos[v] for u in range(n) for v in succ[u])


def all_orders(n: int, succ: Succ) -> Iterator[List[int]]:
    """Yield every topological order (exponential; small graphs only)."""
    deg = indegrees(n, succ)
    order: List[int] = []

    def rec() -> Iterator[List[int]]:
        if len(order) == n:
            yield list(order)
            return
        for u in range(n):
            if deg[u] == 0 and u not in placed:
                placed.add(u)
                order.append(u)
                for v in succ[u]:
                    deg[v] -= 1
                yield from rec()
                for v in succ[u]:
                    deg[v] += 1
                order.pop()
                placed.remove(u)

    placed: Set[int] = set()
    yield from rec()


def count_orders(n: int, succ: Succ) -> int:
    """Number of topological orders, by DP over the set of placed nodes.

    O(2^n * n): fine up to ~20 nodes, where enumeration would be hopeless.
    """
    if n > 24:
        raise ValueError("count_orders is exponential; n too large")
    predmask = [0] * n
    for u in range(n):
        for v in succ[u]:
            predmask[v] |= 1 << u
    full = (1 << n) - 1

    @lru_cache(maxsize=None)
    def ways(placed: int) -> int:
        if placed == full:
            return 1
        total = 0
        for v in range(n):
            bit = 1 << v
            if not placed & bit and predmask[v] & placed == predmask[v]:
                total += ways(placed | bit)
        return total

    return ways(0)