Building a mini llguidance

I had read the llguidance source several times before I felt I understood it. I knew there was a vocabulary trie, an Earley parser, a regex lexer, and something called a slicer. I could identify each piece in the repository. I could not have written down the complete sequence of events in a mask call without quietly reopening six source files.

So I built a small version.

It takes a grammar, a model vocabulary, and the bytes generated so far. Its only useful operation is:

get_mask(state) -> set[token_id]

The result contains every model token that can be appended without making the grammar impossible to complete.

My version omits most of llguidance. It has no JSON Schema compiler, captures, stop nodes, token forcing, lazy lexemes, substring constraints, compact trie representation, or serious concern for performance. The pieces it does have are enough to explain the ordinary mask path: a vocabulary trie speculatively feeds bytes into a lexer, and the lexer occasionally wakes an Earley parser.

A very small constraint

We will use a tiny grammar for an object containing one Boolean field:

start  ::= LBRACE WS KEY COLON WS BOOL WS RBRACE

LBRACE  = {
RBRACE  = }
KEY     = "ok"
COLON   = :
BOOL    = true | false
WS      = [ ]*

And an extremely bad model vocabulary:

0   {
1   }
2   "
3   "ok"
4   '": '
5   true
6   false
7   null
8   trout

After the model has generated {, tokens ", "ok", and perhaps a token beginning with spaces may be useful. null and trout are not.

There are two kinds of token here that make constrained decoding slightly awkward. Token 3 contains a complete KEY lexeme. Token 4 contains the end of that key, a colon, and trailing whitespace. Model-token boundaries and grammar-lexeme boundaries have no obligation to line up.

The model may also emit only part of a lexeme. A token containing "o should remain valid even though the closing quote has not arrived. The constraint only needs the prefix to be capable of becoming valid later.

Try every token

The first mask generator is wonderfully stupid:

def get_mask_slow(recognizer, vocabulary):
    allowed = set()

    for token_id, token_bytes in vocabulary.items():
        candidate = recognizer.clone()

        if candidate.push_bytes(token_bytes):
            allowed.add(token_id)

    return allowed

For each model token, copy the recognizer and feed it the token’s bytes. Keep the token when every byte succeeds and the resulting state can still reach a complete parse.

This implementation is useful even after we replace it. It gives us a reference mask for testing the trie walk. If the clever version disagrees with it, the clever version has a bug.

It also repeats a ridiculous amount of work. Tokens ", "ok", "other", and "occasionally" all begin with the same byte. A real vocabulary contains tens or hundreds of thousands of tokens with shared prefixes. The slow loop feeds those shared bytes into fresh recognizer copies over and over again.

Put the vocabulary in a trie

A trie stores one byte on each edge:

root
 ├─ {
 ├─ }
 ├─ "
 │   └─ o
 │       └─ k
 │           └─ "
 ├─ t
 │   └─ r
 │       ├─ u → e
 │       └─ o → u → t
 ├─ f → a → l → s → e
 └─ n → u → l → l

The recognizer reads the first t once. If t is impossible, both true and trout disappear. If tr remains possible but tro fails, the walk skips the rest of the trout branch while continuing through true.

The trie walk needs a recognizer with three operations:

checkpoint = recognizer.save()
ok = recognizer.push_byte(byte)
recognizer.restore(checkpoint)

Then the mask is a depth-first traversal:

def walk(node, recognizer, allowed):
    if node.token_id is not None:
        if recognizer.can_continue():
            allowed.add(node.token_id)

    saved = recognizer.save()

    for byte, child in node.children.items():
        recognizer.restore(saved)

        if recognizer.push_byte(byte):
            walk(child, recognizer, allowed)

The real TokTrie::add_bias_inner stores the trie in a flat array. Each node records its subtree size, so a rejected edge becomes an array jump rather than a recursive visit. Its recognizer has matching push/pop operations for returning from a speculative branch.

Our tree of dictionaries is slower and much easier to look at.

The recognizer has two floors

Earley parsers consume grammar symbols such as KEY, COLON, and BOOL. The trie supplies bytes. A lexer lives between them.

The recognizer therefore carries:

class Recognizer:
    earley_row: EarleyRow
    lexer_state: LexerState

The current Earley row tells us which terminals may occur next. At the start of the object, only LBRACE is useful. After scanning LBRACE, Earley advances to a row expecting WS, and eventually KEY.

The lexer activates the regular expressions belonging to those terminals. Feeding it a byte advances the active regex states. llguidance uses derivre, whose states are based on regular-expression derivatives and constructed lazily as they are encountered.

For the miniature version, we can pretend we already have a DFA for every terminal:

def advance_lexer(active, byte):
    next_states = {
        terminal: dfa[terminal].step(state, byte)
        for terminal, state in active.items()
    }

    return {
        terminal: state
        for terminal, state in next_states.items()
        if not state.is_dead
    }

If every active state dies, the byte is impossible in the current lexeme. If some survive, the trie walk continues.

An accepting lexer state does not necessarily mean the lexeme should be emitted immediately. After reading t, a literal such as true is still incomplete. After reading true, it is complete. An identifier regex might be complete after one letter while remaining capable of consuming another hundred.

llguidance supports several lexeme modes. Our miniature recognizer will use a simple greedy rule: continue the current lexeme until the next byte kills it. If the previous lexer state was accepting, scan that lexeme into Earley and retry the killing byte using the terminals allowed by the new row.

def push_byte(self, byte):
    next_lexer = self.lexer.advance(byte)

    if not next_lexer.is_dead():
        self.lexer = next_lexer
        return True

    matches = self.lexer.accepting_terminals()
    if not matches:
        return False

    next_row = self.earley.scan(matches)
    if next_row is None:
        return False

    self.earley = next_row
    self.lexer = Lexer.start(next_row.allowed_terminals())
    return self.push_byte(byte)

Several details are hiding in scan. More than one terminal may match the same byte span, and the Earley row may accept more than one of them. The parser keeps the alternatives that lead somewhere and drops the rest.

The recursive push_byte(byte) is also important. The byte that killed the old lexeme has not been consumed. It may be the first byte of the next lexeme. A token containing "ok": crosses the KEY, COLON, and WS boundaries during one trie descent.

What Earley contributes

An Earley row is a set of dotted grammar items. For a production:

start ::= LBRACE WS KEY COLON WS BOOL WS RBRACE

an item might look like:

start ::= LBRACE WS • KEY COLON WS BOOL WS RBRACE

The dot records how far the production has advanced. Prediction adds productions for a nonterminal after the dot. Scanning moves the dot across a terminal supplied by the lexer. Completion advances items that were waiting for a finished nonterminal.

We do not need to teach the trie any of this. The trie asks whether one byte survived. Earley wakes only after the lexer has assembled enough bytes to report a lexeme.

This is an excellent division of labour for JSON. A hundred bytes inside a string may produce a hundred lexer transitions and no Earley work. Quotes, colons, commas, and brackets create the parser-visible events.

The Earley rows also survive trie backtracking. Suppose several vocabulary branches finish the same lexeme sequence. They can arrive at the same row instead of reconstructing it from the beginning.

Tokens may end between lexemes

At a trie node carrying a token ID, can_continue() has to answer a slightly different question from “is the complete grammar accepting?”

A model token ending halfway through "ok" is usable. Its lexer state is non-accepting but live: later bytes can still finish the key. A token ending immediately after the key is also usable, even though the parser still expects a colon. The generated prefix has not completed the grammar, but it has not made completion impossible.

Our miniature recognizer accepts a token node when:

the current lexer state can still reach a match accepted by the current Earley row; or
a lexeme has just completed and the resulting Earley row can still reach the grammar’s end.

End of sequence needs separate treatment. A completed greedy lexeme may be waiting for a killing byte that never comes. Before declaring the whole generation complete, the recognizer flushes an accepting lexer state into Earley and checks whether the grammar can finish.

Model-token endings do not perform that flush. The next model token may continue the same lexeme.

Skipping parts of the trie

The trie helps most when many candidates fail near their first few bytes. It helps less inside a permissive string, where a large fraction of the vocabulary remains legal.

llguidance has a slicer for this case. During tokenizer construction, it groups useful sets of token byte strings and stores their token masks. One such set can cover bounded runs of ordinary JSON-string bytes:

[^"\\\x00-\x1F\x7F]{1,30}

When the current lexer state accepts every byte string in a slice as a valid prefix, the mask generator ORs the precomputed slice mask into the result. It does not visit those trie nodes individually.

A toy slice can be added to our walk:

def walk(node, recognizer, allowed):
    for slice in node.slices:
        if recognizer.subsumes(slice.regex):
            allowed |= slice.token_mask
            node = node.without(slice)

    # Continue through whatever the slices did not cover.
    ...

The real test is more careful than this sketch. It relates the slice regex to the active lexer’s prefix language, and the vocabulary partition is designed to make the operation worthwhile. Still, the mechanical effect is simple: admit a prepared group of tokens at once.

Pruning makes restrictive masks cheap. Slicing helps with permissive masks.

Finding the same machine in llguidance

The miniature components have fairly direct source counterparts in llguidance v1.6.1:

TokenParser::compute_mask begins the mask operation.
The parser-side recognizer in parser/src/earley/parser.rs connects byte advancement to lexer and Earley state.
TokTrie::add_bias_inner performs the flat trie walk.
parser/src/earley/lexer.rs contains the lexer handoff.
parser/src/earley/slicer.rs implements the vocabulary-slice optimization.

The production code is compact, cached, flattened, and full of cases our version ignores. The basic journey remains recognisable: trie byte, lexer transition, occasional Earley scan, then back up the trie.

There is one awkward question hiding inside our greedy push_byte: what happens after a lexeme matches, continues into a live non-matching state, and dies later? Our miniature implementation loses the older match. The llguidance path I tested does too. Five carefully selected bytes are enough to expose it.

That deserved a separate article.