A weighted graph-structured stack

A graph-structured stack compresses a set of stacks into a graph where each path represents one stack. When stacks share a run of items, their paths use the same nodes; the graph branches only where they differ.

WeightedGSS extends this representation by assigning a weight to each stack.

For example, suppose we construct one from three stacks:

stacks = [
    ([0, 5, 8, 12, 18, 31], p),
    ([0, 5, 14, 27, 31], p),
    ([0, 4, 13, 15, 27, 31], q),
]

WeightedGSS.from_stacks(stacks)

The graph is split by an interface of weight-bearing nodes. Every non-empty path crosses this interface exactly once, so each represented stack has exactly one weight. Nodes on the root side of the interface are upper nodes; those on the leaf side are lower nodes.

To limit how much of the graph must be traversed to reach each weight, WeightedGSS keeps the interface close to the root. This matters because mask generation reads the weights frequently.

Each operation returns a new graph, creating only the nodes in the part that changed and reusing the rest. The unchanged subgraphs are immutable, so they can be shared safely between the old and new graphs.

Most of the time, however, the graph is largely linear. Ambiguity is usually local and regular, and often absent altogether. WeightedGSS therefore packs deterministic chains into segments and exposes a mutable LinearPrefix for callers that want a fast path for the linear case.²

Why does a weight need join?

A pop can make two stacks identical, requiring their weights to be joined:

WeightedGSS.from_stacks([    ([0, 4, 9, 18], p),
    ([0, 4, 9, 27], q),
    ([0, 4, 10, 35], r),
])

WeightedGSS.from_stacks([    ([0, 4, 9], p),
    ([0, 4, 9], q),
    ([0, 4, 10], r),
])

WeightedGSS.from_stacks([    ([0, 4, 9], p ⊔ q),
    ([0, 4, 10], r),
])

Motivation

weighted-gss grew out of some awkward state I needed to track in glrmask. In general, WeightedGSS does not care what the weight is, as long as it implements a join operation. In glrmask, for reasons I will not get into deeply here (see GLRMask: Constrained Decoding with Weighted Automata), the weights are disallowed-terminals maps used to preserve greedy lexer semantics.

That left the data structure with a few requirements:

• Rather obviously, stacks have no stable identity apart from their paths through the GSS, so their weights must be encoded in those paths somehow.
• Every non-empty path through the GSS should contain exactly one weight-bearing node. The alternative would be distributing a weight across several pieces and somehow reconstructing it, which is a bit silly. So: one weight per path. If you think about what kind of shape this induces in the graph, the weight nodes form a kind of interface layer that cuts the graph in half, with all other nodes falling either above or below it. So our Upper/Interface/Lower node types arise naturally as a consequence of requiring one weight per path.
• Mask generation reads the weights frequently, so this interface should remain as close to the root as possible.
• The graph must support frequent incremental changes, including merges.

Much of the difficulty lies in lifting weights towards the root while keeping updates cheap. The graph should also share as much internal structure as possible without making pushes, pops, and merges expensive.

Validation

A weighted GSS represents a set of stacks with a weight attached to each one. This makes it easy to test. We can write a slow implementation that stores every stack directly in a dict[tuple[T, ...], W]. The property tests and fuzzer apply the same random operations to both implementations, then expand the GSS and compare the resulting stack-to-weight mappings after every step.

from collections.abc import Callable
from dataclasses import dataclass, field


type Stack[T] = tuple[T, ...]


@dataclass
class ExplicitStacks[T, W]:
    join: Callable[[W, W], W]
    stacks: dict[Stack[T], W] = field(default_factory=dict)

    def pop(self):
        result = ExplicitStacks(self.join)

        for stack, weight in self.stacks.items():
            if not stack:
                continue

            new_stack = stack[:-1]
            if new_stack in result.stacks:
                weight = self.join(result.stacks[new_stack], weight)
            result.stacks[new_stack] = weight

        return result