Skip to main content

keyhog_scanner/
bigram_bloom.rs

1//! Selective literal-anchor prefilter (the public diagnostic name remains
2//! "bigram Bloom" for compatibility).
3//!
4//! The original Layer-0.5 table inserted every bigram from every detector
5//! literal and widened each terminal byte to a full row. The production table
6//! was only 28.13% populated, but ordinary source lines still overlapped it
7//! almost universally: density was truthful while rejection was zero.
8//!
9//! Construction now chooses one mandatory anchor for every direct-matcher
10//! literal alternative. Alternatives shorter than eight bytes use one exact
11//! ASCII-case-insensitive automaton; longer alternatives select the rarest
12//! eight-byte window measured across the compiled literal corpus, with
13//! deterministic byte and position tie-breaking. Every complete alternative
14//! therefore contains its selected anchor. Empty/unextractable alternatives
15//! invalidate this gate and fail open.
16//!
17//! The 65,536-bit table stores two independent stable hashes of each selected
18//! eight-byte anchor. A query requires both bits. Hash collisions can only
19//! admit extra chunks; they cannot reject a real candidate.
20//! Prefixless and dynamic regexes are not trained into this table. They remain
21//! in the scanner's explicit phase-2 always-admit/no-hit lane, which is evaluated
22//! even when this direct-literal gate rejects a chunk.
23//!
24//! Keeping the existing 8 KB table and diagnostic surface preserves operator
25//! compatibility while replacing an overbroad two-byte representation with a
26//! selective, proof-carrying one.
27
28#![deny(unsafe_op_in_unsafe_fn)]
29
30/// Scanner-owned selective anchor gate: one exact short-literal automaton plus
31/// a 65,536-bit (8 KB) double-hash table for selected eight-byte anchors.
32///
33/// `Box<[u64; 1024]>` (not inline) keeps the `CompiledScanner` struct compact:
34/// the scanner is moved during compile, and 8 KB inline would force stack
35/// spill on every move.
36pub(crate) struct BigramBloom {
37    bits: Box<[u64; 1024]>,
38    /// Exact owner for mandatory literal alternatives shorter than eight bytes.
39    short_anchors: Option<aho_corasick::AhoCorasick>,
40    /// Bit `n - 1` is set when hashed anchors of byte width `n` exist.
41    width_mask: u8,
42    minimum_anchor_bytes: u8,
43    /// Cached build-time state. Saturated and invalid states fail open.
44    state: BigramPrefilterState,
45}
46
47/// Operator-visible health state for the Layer-0.5 selective anchor prefilter.
48///
49/// `Saturated` and `Invalid` are deliberately fail-open states: production
50/// scanning bypasses the prefilter and retains full downstream matcher recall.
51#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub enum BigramPrefilterState {
53    Healthy,
54    Saturated,
55    Invalid,
56}
57
58/// Build-time hash-table density diagnostics for the Layer-0.5 prefilter.
59#[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/// Rejection effectiveness measured over one explicitly named input corpus.
69#[derive(Clone, Copy, Debug, Eq, PartialEq)]
70pub struct BigramPrefilterCorpusStatus<'a> {
71    pub corpus_name: &'a str,
72    pub input_count: u64,
73    /// Inputs large enough for the production bloom gate to run.
74    pub eligible_inputs: u64,
75    pub rejected_inputs: u64,
76    /// Rejected share in basis points (10_000 = 100.00%).
77    pub rejection_basis_points: u16,
78}
79
80/// When the set-bit fraction of the 65,536-slot hash table reaches this share,
81/// collision admissions make the hashed owner ineffective. At that point the
82/// downstream AC/HS automaton should run unconditionally rather than paying for
83/// a dead prefilter pass. The exact short-anchor owner does not alter the table
84/// population. 60% (39,322 slots) retains conservative saturation headroom.
85const 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    /// Build one mandatory anchor for every literal alternative.
163    ///
164    /// Long-anchor candidate frequency is measured over the complete compiled
165    /// literal corpus. A hash collision or common selected anchor only causes
166    /// extra admissions. Short alternatives use exact matching. An empty
167    /// alternative cannot provide a rejection proof, so construction marks the
168    /// filter invalid and every query fails open.
169    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    /// Return whether a selected mandatory anchor may occur in `chunk`.
239    ///
240    /// Saturated/invalid states and chunks shorter than the shortest compiled
241    /// anchor fail open. A healthy miss proves that none of the direct matcher
242    /// literal alternatives can occur. Prefixless/dynamic phase-2 alternatives
243    /// are owned by the scanner's separate always-admit lane.
244    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/// Two stable 16-bit slots for an anchor of one to eight bytes. Requiring both
410/// bits keeps long real-corpus lines selective without enlarging the public
411/// 65,536-slot diagnostic table. Collisions remain fail-open admissions.
412#[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}