Blumer et al. suffix automaton algorithm in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/blumer-algorithm/. The authors are listed in prov/EMPS9K.jsonl.


class SuffixAutomaton:
    """The suffix automaton (DAWG) of a string, built online by the algorithm of Blumer et al.

    State 0 is the initial state. For each state v, length[v] is the length of the
    longest string that reaches v, link[v] is its suffix link (-1 for the initial state)
    and nxt[v] maps a symbol to the target state.
    """

    def __init__(self, s: str = "") -> None:
        self.length: list[int] = [0]
        self.link: list[int] = [-1]
        self.nxt: list[dict[str, int]] = [{}]
        self.last = 0
        self.text: list[str] = []
        for c in s:
            self.extend(c)

    def _new_state(self, length: int, link: int, nxt: dict[str, int]) -> int:
        self.length.append(length)
        self.link.append(link)
        self.nxt.append(nxt)
        return len(self.length) - 1

    def extend(self, c: str) -> None:
        """Append the symbol c to the text and update the automaton."""
        self.text.append(c)
        cur = self._new_state(self.length[self.last] + 1, -1, {})  # @line L1
        p = self.last
        while p != -1 and c not in self.nxt[p]:  # @line L2
            self.nxt[p][c] = cur
            p = self.link[p]
        if p == -1:  # @line L3
            self.link[cur] = 0
        else:
            q = self.nxt[p][c]
            if self.length[q] == self.length[p] + 1:  # @line L4
                self.link[cur] = q
            else:
                clone = self._new_state(self.length[p] + 1, self.link[q], dict(self.nxt[q]))  # @line L5
                while p != -1 and self.nxt[p].get(c) == q:  # @line L6
                    self.nxt[p][c] = clone
                    p = self.link[p]
                self.link[q] = clone  # @line L7
                self.link[cur] = clone
        self.last = cur  # @line L8

    def num_states(self) -> int:
        return len(self.length)

    def num_transitions(self) -> int:
        return sum(len(d) for d in self.nxt)

    def contains(self, w: str) -> bool:
        """Is w a substring of the text?"""
        v = 0
        for c in w:
            if c not in self.nxt[v]:
                return False
            v = self.nxt[v][c]
        return True

    def is_suffix(self, w: str) -> bool:
        """Is w a suffix of the text? The terminal states are those on the link path from last."""
        v = 0
        for c in w:
            if c not in self.nxt[v]:
                return False
            v = self.nxt[v][c]
        u = self.last
        while u != -1:
            if u == v:
                return True
            u = self.link[u]
        return False

    def count_distinct_substrings(self) -> int:
        """The number of distinct nonempty substrings: the sum of len(v) - len(link(v))."""
        return sum(self.length[v] - self.length[self.link[v]] for v in range(1, len(self.length)))

    def check_invariant(self) -> None:
        n = len(self.text)
        k = len(self.length)
        assert self.length[0] == 0 and self.link[0] == -1
        assert self.length[self.last] == n
        for v in range(1, k):
            u = self.link[v]
            assert 0 <= u < k and self.length[u] < self.length[v], "suffix link must shorten"
            for c, w in self.nxt[v].items():
                assert self.length[w] >= self.length[v] + 1
        for c, w in self.nxt[0].items():
            assert self.length[w] >= 1
        if n >= 2:
            assert k <= 2 * n - 1
        if n >= 3:
            assert self.num_transitions() <= 3 * n - 4
        # every state is reached by a string of each length in (len(link), len], here checked via
        # the longest string: following the transitions of a longest string reaches the state.
        # Compute one longest string per state by a walk in topological (length) order.
        word: list[str | None] = [None] * k
        word[0] = ""
        for v in sorted(range(k), key=lambda x: self.length[x]):
            wv = word[v]
            assert wv is not None, "state not reachable"
            for c, w in self.nxt[v].items():
                if self.length[w] == self.length[v] + 1 and word[w] is None:
                    word[w] = wv + c
        text = "".join(self.text)
        for v in range(k):
            wv = word[v]
            assert wv is not None and len(wv) == self.length[v] and wv in text