Suffix automaton in Python

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

from collections.abc import Hashable, Iterable, Sequence
from typing import Any


class SuffixAutomaton:
    """The suffix automaton (DAWG) of a text, built online.

    Symbols are any hashable values; transitions are dicts, so only equality
    and hashing of symbols are used. Positions are 0-based; no sentinel.
    State 0 is the initial state. For each state v: nxt[v] maps a symbol to a
    state, link[v] is the suffix link, length[v] is the length of the longest
    string of the class, and first[v] is the end position of its first
    occurrence. At most 2n - 1 states for n >= 2. No recursion.
    """

    def __init__(self, text: Iterable[Hashable] = ()) -> None:
        self.nxt: list[dict[Any, int]] = [{}]
        self.link: list[int] = [-1]
        self.length: list[int] = [0]
        self.first: list[int] = [-1]
        self.is_clone: list[bool] = [False]
        self.last = 0
        self.n = 0
        self._occ: list[int] | None = None
        for c in text:
            self.extend(c)

    def _state(self, length: int, first: int, nxt: dict[Any, int], link: int, clone: bool) -> int:
        self.nxt.append(nxt)
        self.link.append(link)
        self.length.append(length)
        self.first.append(first)
        self.is_clone.append(clone)
        return len(self.nxt) - 1

    def extend(self, c: Hashable) -> None:
        """Append one symbol: the classes of the new suffixes are created or split."""
        self._occ = None
        cur = self._state(self.length[self.last] + 1, self.n, {}, 0, False)  # @line L1
        p = self.last  # @line L2
        while p != -1 and c not in self.nxt[p]:  # @line L3
            self.nxt[p][c] = cur
            p = self.link[p]
        if p != -1:  # @line L4
            q = self.nxt[p][c]
            if self.length[p] + 1 == self.length[q]:  # @line L5: q is solid
                self.link[cur] = q
            else:  # @line L6: split q by cloning
                cl = self._state(self.length[p] + 1, self.first[q], dict(self.nxt[q]), self.link[q], True)
                while p != -1 and self.nxt[p].get(c) == q:  # @line L7
                    self.nxt[p][c] = cl
                    p = self.link[p]
                self.link[q] = self.link[cur] = cl  # @line L8
        self.last = cur  # @line L9
        self.n += 1

    def _walk(self, p: Sequence[Hashable]) -> int:
        v = 0
        for c in p:  # @line L10
            v = self.nxt[v].get(c, -1)
            if v < 0:
                return -1
        return v  # @line L11

    def contains(self, p: Sequence[Hashable]) -> bool:
        return self._walk(p) >= 0

    def is_suffix(self, p: Sequence[Hashable]) -> bool:
        """Whether p is a suffix of the text: its state lies on the link path of last."""
        v = self._walk(p)
        if v < 0:
            return False
        u = self.last
        while u != -1 and self.length[u] >= len(p):
            if u == v:
                return True
            u = self.link[u]
        return False

    def _counts(self) -> list[int]:
        if self._occ is None:
            occ = [0 if self.is_clone[v] or v == 0 else 1 for v in range(len(self.nxt))]
            for v in sorted(range(1, len(self.nxt)), key=lambda v: -self.length[v]):
                occ[self.link[v]] += occ[v]
            self._occ = occ
        return self._occ

    def count(self, p: Sequence[Hashable]) -> int:
        """The number of occurrences of p (n for the empty pattern)."""
        if len(p) == 0:
            return self.n
        v = self._walk(p)
        return 0 if v < 0 else self._counts()[v]

    def first_occurrence(self, p: Sequence[Hashable]) -> int:
        """The start of the leftmost occurrence of p, or -1."""
        v = self._walk(p)
        if v < 0:
            return -1
        return 0 if len(p) == 0 else self.first[v] - len(p) + 1

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

    def state_count(self) -> int:
        return len(self.nxt)

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

    def check_invariant(self) -> None:
        """Links shorten lengths, transitions go forward in length, sizes are bounded."""
        m = len(self.nxt)
        n = self.n
        assert self.link[0] == -1 and self.length[0] == 0
        for v in range(1, m):
            u = self.link[v]
            assert 0 <= u < m and self.length[u] < self.length[v]
            for w in self.nxt[v].values():
                assert self.length[w] > self.length[v]
        for w in self.nxt[0].values():
            assert self.length[w] >= 1
        if n >= 2:
            assert m <= 2 * n - 1
        if n >= 3:
            assert self.transition_count() <= 3 * n - 4
        # the state of the longest string of each class is reached by it
        for v in range(1, m):
            assert self.length[v] <= n