HAT-trie in Python

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

from typing import Generic, TypeVar

V = TypeVar("V")


def _hash(s: str) -> int:
    """Shift-add-xor string hash, deterministic across runs."""
    h = 220373
    for c in s:
        h ^= ((h << 5) + (h >> 2) + ord(c)) & 0xFFFFFFFF
    return h


class Container(Generic[V]):
    """An array hash table: each slot is one array of (suffix, value) entries.

    A pure container hangs from one pointer and stores the suffixes without the
    symbol of that pointer. A hybrid container hangs from the pointers lo .. hi - 1
    of its parent and keeps the first symbol of each suffix.
    """

    def __init__(self, slots: int, pure: bool, lo: int, hi: int) -> None:
        self.slots: list[list[tuple[str, V]]] = [[] for _ in range(slots)]
        self.pure = pure
        self.lo, self.hi = lo, hi
        self.size = 0

    def find(self, t: str) -> V | None:
        for k, v in self.slots[_hash(t) % len(self.slots)]:  # @line L5
            if k == t:
                return v
        return None

    def put(self, t: str, v: V) -> bool:
        """Insert or update t; return True if t is new."""
        slot = self.slots[_hash(t) % len(self.slots)]
        for i, (k, _) in enumerate(slot):
            if k == t:
                slot[i] = (t, v)
                return False
        slot.append((t, v))  # @line L7
        self.size += 1
        return True

    def items(self) -> list[tuple[str, V]]:
        return [e for slot in self.slots for e in slot]


class Node(Generic[V]):
    __slots__ = ("child", "value")

    def __init__(self, r: int) -> None:
        self.child: list[Node[V] | Container[V]] = []
        self.value: V | None = None  # the value of the key that ends at this node


class HatTrie(Generic[V]):
    """A HAT-trie over the symbols with codes 0 .. r - 1.

    Trie nodes have an array of r pointers; leaves are array hash containers of at
    most `limit` entries. Values must not be None.
    """

    def __init__(self, r: int = 128, limit: int = 1024, slots: int = 64) -> None:
        if limit < 1 or slots < 1:
            raise ValueError("the limit and the number of slots must be positive")
        self.r, self.limit, self.nslots = r, limit, slots
        self.root: Node[V] = self._node()
        self.n = 0
        self.bursts = 0

    def _node(self) -> Node[V]:
        x: Node[V] = Node(self.r)
        h: Container[V] = Container(self.nslots, False, 0, self.r)
        x.child = [h] * self.r
        return x

    def _code(self, c: str) -> int:
        k = ord(c)
        if k >= self.r:
            raise ValueError(f"symbol {c!r} is outside the alphabet")
        return k

    def _descend(self, s: str) -> tuple[Node[V], int]:
        """The deepest trie node on the path of s and the number of symbols it consumed."""
        x, d = self.root, 0  # @line L1
        while d < len(s):  # @line L2
            y = x.child[self._code(s[d])]  # @line L3
            if not isinstance(y, Node):
                break
            x, d = y, d + 1
        return x, d

    def get(self, s: str) -> V | None:
        x, d = self._descend(s)
        if d == len(s):
            return x.value
        y = x.child[self._code(s[d])]
        assert isinstance(y, Container)
        return y.find(s[d + 1:] if y.pure else s[d:])  # @line L4

    def put(self, s: str, v: V) -> None:
        if v is None:
            raise ValueError("a value must not be None")
        x, d = self._descend(s)  # @line L6
        if d == len(s):
            if x.value is None:
                self.n += 1
            x.value = v
            return
        y = x.child[self._code(s[d])]
        assert isinstance(y, Container)
        if y.put(s[d + 1:] if y.pure else s[d:], v):
            self.n += 1
        if y.size > self.limit:  # @line L8
            self._burst(x, y)

    def _burst(self, x: Node[V], y: Container[V]) -> None:
        work = [(x, y)]
        while work:  # @line L12
            x, y = work.pop()
            if y.size <= self.limit:
                continue
            self.bursts += 1
            if y.pure:  # @line L9
                z = self._node()
                x.child[y.lo] = z
                h = z.child[0]
                assert isinstance(h, Container)
                for t, v in y.items():
                    if t == "":
                        z.value = v
                    else:
                        h.put(t, v)
                work.append((z, h))
                continue
            counts = [0] * (y.hi - y.lo)
            for t, _ in y.items():
                counts[ord(t[0]) - y.lo] += 1
            if y.hi - y.lo == 1:  # @line L11
                p: Container[V] = Container(self.nslots, True, y.lo, y.hi)
                for t, v in y.items():
                    p.put(t[1:], v)
                x.child[y.lo] = p
                work.append((x, p))
                continue
            used = [y.lo + i for i, k in enumerate(counts) if k]
            if len(used) == 1:  # @line L10
                c = used[0]  # one symbol only: cut out its range
                cuts = [y.lo, c, c + 1, y.hi]
            else:
                best, split, left = None, used[1], 0
                for s in range(used[0] + 1, used[-1] + 1):
                    left += counts[s - 1 - y.lo]
                    gap = abs(2 * left - y.size)
                    if best is None or gap < best:
                        best, split = gap, s
                cuts = [y.lo, split, y.hi]
            for lo, hi in zip(cuts, cuts[1:]):
                if lo == hi:
                    continue
                part: Container[V] = Container(self.nslots, False, lo, hi)
                for t, v in y.items():
                    if lo <= ord(t[0]) < hi:
                        part.put(t, v)
                for c in range(lo, hi):
                    x.child[c] = part
                work.append((x, part))

    def size(self) -> int:
        return self.n

    def keys(self) -> list[str]:
        """All keys in sorted order; each container is sorted when it is reached."""
        out: list[str] = []
        stack: list[tuple[Node[V] | None, str, Container[V] | None]] = [(self.root, "", None)]
        while stack:
            x, prefix, c = stack.pop()
            if c is not None:
                out.extend(sorted(prefix + t for t, _ in c.items()))
                continue
            assert x is not None
            if x.value is not None:
                out.append(prefix)
            entries: list[tuple[Node[V] | None, str, Container[V] | None]] = []
            k = 0
            while k < self.r:
                y = x.child[k]
                if isinstance(y, Node):
                    entries.append((y, prefix + chr(k), None))
                    k += 1
                elif y.pure:
                    entries.append((None, prefix + chr(k), y))
                    k += 1
                else:
                    entries.append((None, prefix, y))
                    k = y.hi
            stack.extend(reversed(entries))
        return out

    def check_invariant(self) -> None:
        count = 0
        stack = [self.root]
        while stack:
            x = stack.pop()
            assert len(x.child) == self.r
            count += x.value is not None
            k = 0
            while k < self.r:
                y = x.child[k]
                if isinstance(y, Node):
                    stack.append(y)
                    k += 1
                    continue
                assert y.lo == k and y.size <= self.limit
                assert all(x.child[c] is y for c in range(y.lo, y.hi))
                assert y.size == len(y.items())
                if y.pure:
                    assert y.hi == k + 1
                else:
                    assert all(t != "" and y.lo <= ord(t[0]) < y.hi for t, _ in y.items())
                for t, _ in y.items():
                    assert y.slots[_hash(t) % len(y.slots)].count((t, y.find(t))) == 1
                count += y.size
                k = y.hi
        assert count == self.n