maskforge_core/mask/
bitmask.rs

1//! Stores allowed-token masks as packed `u32` words.
2
3use crate::error::{CompileError, ErrorCode, Stage};
4use crate::primitives::TokenId;
5
6const WORD_BITS: usize = 32;
7
8/// A dense allowed-token bitmask over `vocab_size` tokens, packed into little-endian `u32` words.
9#[derive(Clone, PartialEq, Eq, Debug)]
10pub struct Bitmask {
11    words: Vec<u32>,
12    vocab_size: usize,
13}
14
15impl Bitmask {
16    /// Allocates an all-zero mask sized for `vocab_size` tokens (`ceil(vocab_size/32)` words).
17    #[must_use]
18    pub fn zeros(vocab_size: usize) -> Self {
19        Self {
20            words: vec![0; vocab_size.div_ceil(WORD_BITS)],
21            vocab_size,
22        }
23    }
24
25    /// Builds a mask from packed words, clearing unused high bits.
26    #[must_use]
27    pub fn from_words(words: &[u32], vocab_size: usize) -> Self {
28        let need = vocab_size.div_ceil(WORD_BITS);
29        if need == 0 {
30            return Self {
31                words: Vec::new(),
32                vocab_size,
33            };
34        }
35        let mut buf = vec![0u32; need];
36        let take = words.len().min(need);
37        buf[..take].copy_from_slice(&words[..take]);
38        let remainder = vocab_size % WORD_BITS;
39        if remainder != 0 {
40            buf[need - 1] &= (1u32 << remainder) - 1;
41        }
42        Self {
43            words: buf,
44            vocab_size,
45        }
46    }
47
48    /// Sets the bit for `id`. Fails (never truncates) when `id` is outside `0..vocab_size`.
49    pub fn set(&mut self, id: TokenId) -> Result<(), CompileError> {
50        let idx = id.get() as usize;
51        if idx >= self.vocab_size {
52            return Err(CompileError::new(
53                ErrorCode::ArtifactOutOfBounds,
54                Stage::ArtifactValidate,
55                "token id is outside the mask vocab size",
56            )
57            .with_observed(idx.to_string()));
58        }
59        self.words[idx / WORD_BITS] |= 1u32 << (idx % WORD_BITS);
60        Ok(())
61    }
62
63    /// Returns whether the bit for `id` is set (`false` for out-of-range ids).
64    #[must_use]
65    pub fn get(&self, id: TokenId) -> bool {
66        let idx = id.get() as usize;
67        if idx >= self.vocab_size {
68            return false;
69        }
70        self.words[idx / WORD_BITS] & (1u32 << (idx % WORD_BITS)) != 0
71    }
72
73    /// The packed words (the wire format): `ceil(vocab_size/32)` little-endian `u32`s.
74    #[must_use]
75    pub fn as_words(&self) -> &[u32] {
76        &self.words
77    }
78
79    /// The number of tokens this mask covers.
80    #[must_use]
81    pub fn vocab_size(&self) -> usize {
82        self.vocab_size
83    }
84
85    /// Returns the number of set token bits.
86    #[must_use]
87    pub fn count_ones(&self) -> usize {
88        self.words
89            .iter()
90            .map(|w| usize::try_from(w.count_ones()).expect("popcount of a u32 fits usize"))
91            .sum()
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn word_count_is_ceil_div_32() {
101        assert_eq!(Bitmask::zeros(0).as_words().len(), 0);
102        assert_eq!(Bitmask::zeros(1).as_words().len(), 1);
103        assert_eq!(Bitmask::zeros(32).as_words().len(), 1);
104        assert_eq!(Bitmask::zeros(33).as_words().len(), 2);
105    }
106
107    #[test]
108    fn set_get_and_unused_high_bits_zero() {
109        let mut m = Bitmask::zeros(40);
110        m.set(TokenId(0)).unwrap();
111        m.set(TokenId(39)).unwrap();
112        assert!(m.get(TokenId(0)));
113        assert!(m.get(TokenId(39)));
114        assert!(!m.get(TokenId(1)));
115        assert_eq!(m.count_ones(), 2);
116        assert_eq!(m.as_words()[1] & !0x0000_00FF, 0);
117    }
118
119    #[test]
120    fn oversized_id_fails_not_truncates() {
121        let mut m = Bitmask::zeros(10);
122        assert_eq!(
123            m.set(TokenId(10)).unwrap_err().code,
124            ErrorCode::ArtifactOutOfBounds
125        );
126        assert!(!m.get(TokenId(10)));
127    }
128
129    #[test]
130    fn word_boundary_ids_31_32_63() {
131        let mut m = Bitmask::zeros(64);
132        for id in [31u32, 32, 63] {
133            m.set(TokenId(id)).unwrap();
134        }
135        assert!(m.get(TokenId(31)) && m.get(TokenId(32)) && m.get(TokenId(63)));
136        assert!(!m.get(TokenId(30)) && !m.get(TokenId(33)));
137        assert_eq!(m.as_words().len(), 2);
138        assert_eq!(m.count_ones(), 3);
139    }
140
141    #[test]
142    fn empty_vocab_has_zero_words() {
143        let m = Bitmask::zeros(0);
144        assert_eq!(m.as_words().len(), 0);
145        assert_eq!(m.count_ones(), 0);
146        assert!(!m.get(TokenId(0)));
147    }
148
149    #[test]
150    fn non_multiple_of_32_vocab_sizes_63_and_65() {
151        assert_eq!(Bitmask::zeros(63).as_words().len(), 2);
152        assert_eq!(Bitmask::zeros(65).as_words().len(), 3);
153        let mut m = Bitmask::zeros(65);
154        m.set(TokenId(64)).unwrap();
155        assert!(m.get(TokenId(64)));
156        assert!(!m.get(TokenId(63)));
157        assert_eq!(
158            m.as_words()[2] & !0b1,
159            0,
160            "id 64 is bit 0 of the third word"
161        );
162    }
163
164    #[test]
165    fn token_id_65_is_out_of_range_for_a_65_vocab() {
166        let mut m = Bitmask::zeros(65);
167        assert_eq!(
168            m.set(TokenId(65)).unwrap_err().code,
169            ErrorCode::ArtifactOutOfBounds
170        );
171    }
172
173    #[test]
174    fn repeated_writes_to_the_same_bit_are_idempotent() {
175        let mut m = Bitmask::zeros(40);
176        m.set(TokenId(5)).unwrap();
177        m.set(TokenId(5)).unwrap();
178        m.set(TokenId(5)).unwrap();
179        assert!(m.get(TokenId(5)));
180        assert_eq!(m.count_ones(), 1);
181    }
182
183    #[test]
184    fn far_oversized_token_id_fails_not_truncates() {
185        let mut m = Bitmask::zeros(10);
186        assert_eq!(
187            m.set(TokenId(u32::MAX)).unwrap_err().code,
188            ErrorCode::ArtifactOutOfBounds
189        );
190        assert!(!m.get(TokenId(u32::MAX)));
191    }
192
193    #[test]
194    fn from_words_masks_dirty_high_bits_beyond_vocab_size() {
195        assert_eq!(Bitmask::from_words(&[u32::MAX], 1).as_words(), &[1]);
196        let m31 = Bitmask::from_words(&[u32::MAX], 31);
197        assert_eq!(m31.as_words(), &[0x7FFF_FFFF]);
198        assert_eq!(m31.count_ones(), 31);
199        let m32 = Bitmask::from_words(&[u32::MAX], 32);
200        assert_eq!(m32.as_words(), &[u32::MAX]);
201        assert_eq!(m32.count_ones(), 32);
202        let m33 = Bitmask::from_words(&[u32::MAX, u32::MAX], 33);
203        assert_eq!(m33.as_words(), &[u32::MAX, 1]);
204        assert_eq!(m33.count_ones(), 33);
205    }
206
207    #[test]
208    fn from_words_at_vocab_size_zero_does_not_panic_and_has_no_words() {
209        let m = Bitmask::from_words(&[u32::MAX], 0);
210        assert_eq!(m.as_words(), &[] as &[u32]);
211        assert_eq!(m.count_ones(), 0);
212    }
213
214    #[test]
215    fn from_words_count_ones_never_counts_out_of_range_bits() {
216        for size in [1usize, 31, 32, 33, 65] {
217            let words = vec![u32::MAX; size.div_ceil(32)];
218            let m = Bitmask::from_words(&words, size);
219            assert_eq!(m.count_ones(), size, "size {size}");
220        }
221    }
222}
maskforge_core/mask/bitmask.rs

maskforge_core/mask/
bitmask.rs
