Struct DenseMask 

pub struct DenseMask {
    pub bits: Vec<u64>,
    pub states: Vec<u32>,
    pub depth: u32,
    pub vocab_size: u32,
}

Bit-packed dense mask for the first depth trie layers.

For depth = 2 and vocab_size = V this is a V × V bit matrix stored as packed u64 words (row-major). A set bit at linear index i * vocab_size + j means the 2-token prefix [i, j] exists in C.

states[i * vocab_size + j] is the trie node reached after emitting [i, j]; 0 is used as a sentinel for invalid entries (the root is never a valid post-prefix destination in practice).

Fields

bits: Vec<u64>

Packed u64 bits. Length = ceil(vocab_size^depth / 64).

states: Vec<u32>

Flat node-ID lookup: length = vocab_size ^ depth. Entry is 0 (sentinel) when the corresponding prefix is absent.

depth: u32

Number of dense layers d.

vocab_size: u32

Vocabulary size |V|.

Implementations

impl DenseMask

pub fn from_constraints( constraints: &[Vec<u32>], vocab_size: u32, depth: u32, node_ids: &[u32], ) -> Self

Build a DenseMask directly from a full constraint set.

This is the canonical constructor used by build_static_index. It is equivalent to calling DenseMask::new followed by repeated insert calls, but avoids recomputing flat_index twice per entry.

Arguments

• constraints — every sequence must have length ≥ depth
• vocab_size — |V|
• depth — number of dense layers d (typically 2)
• node_ids — parallel slice: node_ids[i] is the trie node reached after the first depth tokens of constraints[i]. Pass an all-zeros slice when node IDs are not yet known and will be back-filled by transition.rs.

pub fn first_token_valid(&self, first_token: u32) -> bool

Returns true if any full-depth prefix starts with first_token.

Operates on packed u64 words without deserialising individual bits. Used by decoder.rs at step 0 to expose the valid first-token set.

Complexity

O(|V|^(depth-1) / 64) ≈ O(1) for small depth and typical |V|.

pub fn partial_prefix_has_extension(&self, partial: &[u32]) -> bool

Returns true if partial (length < depth) can be extended to a valid full-depth prefix in the constraint set.

Panics (debug)

Panics if partial.len() >= depth.

pub fn intersect(&self, other: &DenseMask) -> DenseMask

Returns a new DenseMask that is the intersection of self and other.

Two masks can be intersected to find the set of prefixes that appear in both constraint sets — useful for multi-constraint filtering.

Panics

Panics if self and other have different vocab_size or depth.

pub fn union(&self, other: &DenseMask) -> DenseMask

Returns a new DenseMask that is the union of self and other.

Used when merging two separately-built index shards. Where both masks have a valid entry, self’s node ID takes precedence.

Panics

Panics if vocab_size or depth differ.

pub fn first_token_packed_mask(&self) -> Vec<u64>

Returns the first-token marginal as a packed Vec<u64> of length ceil(vocab_size / 64).

Bit t is set iff token t is a valid first token in the constraint set. This vec can be ANDed directly with the model’s top-k bitmask.

pub fn second_token_packed_mask(&self, first_token: u32) -> Vec<u64>

Returns the second-token marginal given that first_token was chosen, packed as a Vec<u64> of length ceil(vocab_size / 64).

Only defined for depth >= 2.

Panics (debug)

Panics if depth < 2.

pub fn count_valid(&self) -> u64

Returns the total number of valid prefixes stored in the mask.

pub fn count_valid_first_tokens(&self) -> u32

Returns the number of distinct valid first tokens.

pub fn to_bytes(&self) -> Vec<u8>

Serialises the mask into a flat byte buffer.

Layout (little-endian):

[u32 vocab_size][u32 depth]
[u32 bits_len][u64 * bits_len]
[u32 states_len][u32 * states_len]

pub fn from_bytes(buf: &[u8]) -> Option<Self>

Deserialises a DenseMask from the byte layout produced by to_bytes.

Returns None if the buffer is malformed.

impl DenseMask

pub fn new(vocab_size: u32, depth: u32) -> Self

Allocate a zeroed mask for vocab_size^depth entries.

pub fn flat_index(&self, tokens: &[u32]) -> usize

Converts a token sequence of length depth to a flat index.

Panics

Panics in debug builds if tokens.len() != depth.

pub fn insert(&mut self, tokens: &[u32], node_id: u32)

Sets the bit and stores the destination node_id for prefix tokens.

pub fn contains(&self, tokens: &[u32]) -> bool

Returns true if the prefix encoded by tokens is marked valid.

pub fn get(&self, v1: u32, v2: u32) -> bool

Two-token shorthand used pervasively in tests and the decoder. Equivalent to contains(&[v1, v2]) when depth == 2.

pub fn state_for(&self, tokens: &[u32]) -> Option<u32>

Returns the trie node reached after the valid prefix tokens, or None if the prefix is absent.

pub fn iter_valid(&self) -> impl Iterator<Item = (Vec<u32>, u32)> + '_

Iterates over all valid prefixes as (tokens, node_id).

Trait Implementations

impl Clone for DenseMask

fn clone(&self) -> DenseMask

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for DenseMask

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Display for DenseMask

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl PartialEq for DenseMask

fn eq(&self, other: &DenseMask) -> bool

Tests for self and other values to be equal, and is used by ==.

fn ne(&self, other: &Rhs) -> bool

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl Eq for DenseMask

impl StructuralPartialEq for DenseMask