1#![deny(unsafe_op_in_unsafe_fn)]
29
30pub(crate) struct BigramBloom {
37 bits: Box<[u64; 1024]>,
38 short_anchors: Option<aho_corasick::AhoCorasick>,
40 width_mask: u8,
42 minimum_anchor_bytes: u8,
43 state: BigramPrefilterState,
45}
46
47#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub enum BigramPrefilterState {
53 Healthy,
54 Saturated,
55 Invalid,
56}
57
58#[derive(Clone, Copy, Debug, Eq, PartialEq)]
60pub struct BigramPrefilterStatus {
61 pub populated_slots: u32,
62 pub total_slots: u32,
63 pub saturation_threshold_slots: u32,
64 pub density_basis_points: u16,
65 pub state: BigramPrefilterState,
66}
67
68#[derive(Clone, Copy, Debug, Eq, PartialEq)]
70pub struct BigramPrefilterCorpusStatus<'a> {
71 pub corpus_name: &'a str,
72 pub input_count: u64,
73 pub eligible_inputs: u64,
75 pub rejected_inputs: u64,
76 pub rejection_basis_points: u16,
78}
79
80const SATURATION_NUMERATOR: u32 = 3;
86const SATURATION_DENOMINATOR: u32 = 5;
87const TABLE_SLOTS: u32 = 65_536;
88const SATURATION_THRESHOLD_SLOTS: u32 =
89 (TABLE_SLOTS * SATURATION_NUMERATOR + SATURATION_DENOMINATOR - 1) / SATURATION_DENOMINATOR;
90
91const MAX_ANCHOR_BYTES: usize = 8;
92
93#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
94struct AnchorKey {
95 bytes: [u8; MAX_ANCHOR_BYTES],
96 len: u8,
97}
98
99impl AnchorKey {
100 fn from_slice(bytes: &[u8]) -> Self {
101 debug_assert!(!bytes.is_empty() && bytes.len() <= MAX_ANCHOR_BYTES);
102 let mut key = Self {
103 bytes: [0; MAX_ANCHOR_BYTES],
104 len: bytes.len() as u8,
105 };
106 key.bytes[..bytes.len()].copy_from_slice(bytes);
107 for byte in &mut key.bytes[..bytes.len()] {
108 *byte = byte.to_ascii_lowercase();
109 }
110 key
111 }
112
113 fn as_slice(&self) -> &[u8] {
114 &self.bytes[..usize::from(self.len)]
115 }
116}
117
118impl Clone for BigramBloom {
119 fn clone(&self) -> Self {
120 Self {
121 bits: Box::new(*self.bits),
122 short_anchors: self.short_anchors.clone(),
123 width_mask: self.width_mask,
124 minimum_anchor_bytes: self.minimum_anchor_bytes,
125 state: self.state,
126 }
127 }
128}
129
130impl BigramBloom {
131 pub(crate) fn empty() -> Self {
132 Self {
133 bits: Box::new([0; 1024]),
134 short_anchors: None,
135 width_mask: width_bit(2),
136 minimum_anchor_bytes: 2,
137 state: BigramPrefilterState::Healthy,
138 }
139 }
140
141 fn blank() -> Self {
142 Self {
143 bits: Box::new([0; 1024]),
144 short_anchors: None,
145 width_mask: 0,
146 minimum_anchor_bytes: 0,
147 state: BigramPrefilterState::Healthy,
148 }
149 }
150
151 #[inline]
152 fn insert_anchor(&mut self, anchor: &[u8]) {
153 for slot in ngram_slots(anchor) {
154 self.bits[slot >> 6] |= 1u64 << (slot & 63);
155 }
156 }
157
158 fn insert_folded_anchor(&mut self, anchor: AnchorKey) {
159 self.insert_anchor(anchor.as_slice());
160 }
161
162 pub(crate) fn from_literal_prefixes(literals: &[String]) -> Self {
170 if literals.is_empty() || literals.iter().any(String::is_empty) {
171 return Self::invalid_for_test();
172 }
173
174 let mut frequencies = std::collections::HashMap::<AnchorKey, u32>::new();
175 let mut short_literals = Vec::<&[u8]>::new();
176 for literal in literals {
177 let bytes = literal.as_bytes();
178 if bytes.len() < MAX_ANCHOR_BYTES {
179 short_literals.push(bytes);
180 continue;
181 }
182 for window in bytes.windows(MAX_ANCHOR_BYTES) {
183 let key = AnchorKey::from_slice(window);
184 frequencies
185 .entry(key)
186 .and_modify(|count| *count = count.saturating_add(1))
187 .or_insert(1);
188 }
189 }
190
191 let mut bloom = Self::blank();
192 let Some(minimum_literal_bytes) = literals.iter().map(|literal| literal.len()).min() else {
193 return Self::invalid_for_test();
194 };
195 bloom.minimum_anchor_bytes = minimum_literal_bytes.min(MAX_ANCHOR_BYTES) as u8;
196 if !short_literals.is_empty() {
197 bloom.short_anchors = match aho_corasick::AhoCorasick::builder()
198 .ascii_case_insensitive(true)
199 .build(short_literals)
200 {
201 Ok(anchors) => Some(anchors),
202 Err(error) => {
203 tracing::error!(%error, "selective short-anchor automaton build failed; filter is invalid and fail-open");
204 return Self::invalid_for_test();
205 }
206 };
207 }
208 for literal in literals {
209 let bytes = literal.as_bytes();
210 if bytes.len() < MAX_ANCHOR_BYTES {
211 continue;
212 }
213 let mut selected = None;
214 for (position, window) in bytes.windows(MAX_ANCHOR_BYTES).enumerate() {
215 let key = AnchorKey::from_slice(window);
216 let Some(frequency) = frequencies.get(&key).copied() else {
217 return Self::invalid_for_test();
218 };
219 let candidate = (frequency, key, position);
220 if selected.is_none_or(|current| candidate < current) {
221 selected = Some(candidate);
222 }
223 }
224 let Some((_, selected, _)) = selected else {
225 return Self::invalid_for_test();
226 };
227 bloom.width_mask |= width_bit(MAX_ANCHOR_BYTES);
228 bloom.insert_folded_anchor(selected);
229 }
230 bloom.recompute_saturation();
231 bloom
232 }
233
234 fn recompute_saturation(&mut self) {
235 self.state = classify_population(self.popcount(), TABLE_SLOTS);
236 }
237
238 pub(crate) fn maybe_overlaps(&self, chunk: &[u8]) -> bool {
245 if self.state != BigramPrefilterState::Healthy {
246 return true;
247 }
248 if chunk.len() < usize::from(self.minimum_anchor_bytes) {
249 return true;
250 }
251 if self
252 .short_anchors
253 .as_ref()
254 .is_some_and(|anchors| anchors.is_match(chunk))
255 {
256 return true;
257 }
258 if self.width_mask == 0 {
259 return false;
260 }
261 chunk
262 .windows(MAX_ANCHOR_BYTES)
263 .any(|window| self.contains_anchor(window))
264 }
265
266 #[inline]
267 fn contains_anchor(&self, anchor: &[u8]) -> bool {
268 let mut folded = [0u8; MAX_ANCHOR_BYTES];
269 for (target, byte) in folded.iter_mut().zip(anchor.iter().copied()) {
270 *target = byte.to_ascii_lowercase();
271 }
272 ngram_slots(&folded[..anchor.len()])
273 .into_iter()
274 .all(|slot| self.bits[slot >> 6] & (1u64 << (slot & 63)) != 0)
275 }
276
277 pub(crate) fn popcount(&self) -> u32 {
278 self.bits.iter().map(|word| word.count_ones()).sum()
279 }
280
281 pub(crate) fn status(&self) -> BigramPrefilterStatus {
282 let populated_slots = self.popcount();
283 let derived_state = classify_population(populated_slots, TABLE_SLOTS);
284 let has_anchor_owner = self.width_mask != 0 || self.short_anchors.is_some();
285 let state = if self.state == BigramPrefilterState::Invalid
286 || self.state != derived_state
287 || !has_anchor_owner
288 {
289 BigramPrefilterState::Invalid
290 } else {
291 derived_state
292 };
293 BigramPrefilterStatus {
294 populated_slots,
295 total_slots: TABLE_SLOTS,
296 saturation_threshold_slots: SATURATION_THRESHOLD_SLOTS,
297 density_basis_points: share_basis_points(populated_slots as u64, TABLE_SLOTS as u64),
298 state,
299 }
300 }
301
302 pub(crate) fn corpus_status<'a, I>(
303 &self,
304 corpus_name: &'a str,
305 inputs: I,
306 minimum_input_bytes: usize,
307 ) -> BigramPrefilterCorpusStatus<'a>
308 where
309 I: IntoIterator<Item = &'a [u8]>,
310 {
311 let mut input_count = 0u64;
312 let mut eligible_inputs = 0u64;
313 let mut rejected_inputs = 0u64;
314 for input in inputs {
315 input_count += 1;
316 if input.len() >= minimum_input_bytes {
317 eligible_inputs += 1;
318 if !self.maybe_overlaps(input) {
319 rejected_inputs += 1;
320 }
321 }
322 }
323 BigramPrefilterCorpusStatus {
324 corpus_name,
325 input_count,
326 eligible_inputs,
327 rejected_inputs,
328 rejection_basis_points: share_basis_points(rejected_inputs, input_count),
329 }
330 }
331
332 pub(crate) fn is_saturated(&self) -> bool {
333 self.status().state == BigramPrefilterState::Saturated
334 }
335
336 #[cfg(test)]
337 pub(crate) fn scalar_overlaps_reference(&self, chunk: &[u8]) -> bool {
338 if self.state != BigramPrefilterState::Healthy {
339 return true;
340 }
341 if chunk.len() < usize::from(self.minimum_anchor_bytes) {
342 return true;
343 }
344 if self
345 .short_anchors
346 .as_ref()
347 .is_some_and(|anchors| anchors.is_match(chunk))
348 {
349 return true;
350 }
351 self.width_mask != 0
352 && chunk
353 .windows(MAX_ANCHOR_BYTES)
354 .any(|window| self.contains_anchor(window))
355 }
356
357 #[cfg(test)]
358 pub(crate) fn saturated_for_test() -> Self {
359 Self::with_population_for_test(SATURATION_THRESHOLD_SLOTS)
360 }
361
362 #[doc(hidden)]
363 pub(crate) fn with_population_for_test(populated_slots: u32) -> Self {
364 let mut bloom = Self::empty();
365 let bounded = populated_slots.min(TABLE_SLOTS) as usize;
366 for slot in 0..bounded {
367 bloom.bits[slot >> 6] |= 1u64 << (slot & 63);
368 }
369 bloom.recompute_saturation();
370 bloom
371 }
372
373 #[doc(hidden)]
374 pub(crate) fn invalid_for_test() -> Self {
375 Self {
376 bits: Box::new([0; 1024]),
377 short_anchors: None,
378 width_mask: 0,
379 minimum_anchor_bytes: 0,
380 state: BigramPrefilterState::Invalid,
381 }
382 }
383}
384
385const fn classify_population(populated_slots: u32, total_slots: u32) -> BigramPrefilterState {
386 if total_slots == 0 || populated_slots > total_slots {
387 return BigramPrefilterState::Invalid;
388 }
389 if populated_slots >= SATURATION_THRESHOLD_SLOTS {
390 BigramPrefilterState::Saturated
391 } else {
392 BigramPrefilterState::Healthy
393 }
394}
395
396fn share_basis_points(numerator: u64, denominator: u64) -> u16 {
397 if denominator == 0 {
398 return 0;
399 }
400 let basis_points = (u128::from(numerator) * 10_000) / u128::from(denominator);
401 basis_points.min(10_000) as u16
402}
403
404#[inline(always)]
405fn width_bit(width: usize) -> u8 {
406 1 << (width - 1)
407}
408
409#[inline(always)]
413fn ngram_slots(bytes: &[u8]) -> [usize; 2] {
414 debug_assert!(!bytes.is_empty() && bytes.len() <= MAX_ANCHOR_BYTES);
415 let mut first = 0x811c_9dc5u32 ^ bytes.len() as u32;
416 let mut second = 0x9e37_79b9u32 ^ (bytes.len() as u32).rotate_left(16);
417 for byte in bytes {
418 first ^= u32::from(*byte);
419 first = first.wrapping_mul(0x0100_0193);
420 second ^= u32::from(*byte);
421 second = second.rotate_left(5).wrapping_mul(0x85eb_ca6b);
422 }
423 [
424 usize::from(((first ^ (first >> 16)) & 0xffff) as u16),
425 usize::from(((second ^ (second >> 16)) & 0xffff) as u16),
426 ]
427}