Skip to main content

qualia_core_db/inference/
neuro_symbolic_sieve.rs

1//! Grammar-constrained FSM sieve for neuro-symbolic LLM output (zero-heap hot path).
2
3use crate::q_hash;
4use crate::NQuin;
5
6/// Max allowed token IDs per FSM state (stack-only mask).
7pub const MAX_SIEVE_ALLOW: usize = 16;
8
9/// One lexicon-bound token slot in a state mask.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct SieveSlot {
12    pub token_id: u32,
13    pub lexicon_hash: u64,
14}
15
16/// Stack mask: linear scan during chunked argmax (no `HashMap` / `Vec`).
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct SieveStateMask {
19    pub slots: [SieveSlot; MAX_SIEVE_ALLOW],
20    pub len: u8,
21}
22
23impl SieveStateMask {
24    pub const EMPTY: Self = Self {
25        slots: [SieveSlot {
26            token_id: 0,
27            lexicon_hash: 0,
28        }; MAX_SIEVE_ALLOW],
29        len: 0,
30    };
31
32    #[inline]
33    pub fn allows(&self, token_id: u32) -> bool {
34        if self.len == 0 {
35            return true;
36        }
37        for i in 0..self.len as usize {
38            if self.slots[i].token_id == token_id {
39                return true;
40            }
41        }
42        false
43    }
44
45    #[inline]
46    pub fn lexicon_hash_for(&self, token_id: u32) -> Option<u64> {
47        for i in 0..self.len as usize {
48            if self.slots[i].token_id == token_id {
49                return Some(self.slots[i].lexicon_hash);
50            }
51        }
52        None
53    }
54
55    pub(crate) fn push(&mut self, token_id: u32, lexicon_hash: u64) {
56        let n = self.len as usize;
57        if n >= MAX_SIEVE_ALLOW {
58            return;
59        }
60        for i in 0..n {
61            if self.slots[i].token_id == token_id {
62                return;
63            }
64        }
65        self.slots[n] = SieveSlot {
66            token_id,
67            lexicon_hash,
68        };
69        self.len += 1;
70    }
71}
72
73/// FSM states for Subject → Predicate → Object graph emission.
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
75pub enum SieveState {
76    ExpectSubject = 0,
77    ExpectPredicate = 1,
78    ExpectObject = 2,
79    Complete = 3,
80}
81
82/// Sieve exhausted: no allowed token had finite logit mass.
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub enum SieveError {
85    Misaligned,
86    AlreadyComplete,
87}
88
89/// Stack-only IRI hash lists for dynamic `.q42.lex` mask population (cold path).
90#[derive(Debug, Clone, Copy)]
91pub struct SieveLexSpec {
92    pub subjects: [u64; MAX_SIEVE_ALLOW],
93    pub subjects_len: u8,
94    pub predicates: [u64; MAX_SIEVE_ALLOW],
95    pub predicates_len: u8,
96    pub objects: [u64; MAX_SIEVE_ALLOW],
97    pub objects_len: u8,
98}
99
100impl SieveLexSpec {
101    pub const EMPTY: Self = Self {
102        subjects: [0; MAX_SIEVE_ALLOW],
103        subjects_len: 0,
104        predicates: [0; MAX_SIEVE_ALLOW],
105        predicates_len: 0,
106        objects: [0; MAX_SIEVE_ALLOW],
107        objects_len: 0,
108    };
109
110    pub fn push_subject(&mut self, hash: u64) {
111        let n = self.subjects_len as usize;
112        if n < MAX_SIEVE_ALLOW {
113            self.subjects[n] = hash;
114            self.subjects_len += 1;
115        }
116    }
117
118    pub fn push_predicate(&mut self, hash: u64) {
119        let n = self.predicates_len as usize;
120        if n < MAX_SIEVE_ALLOW {
121            self.predicates[n] = hash;
122            self.predicates_len += 1;
123        }
124    }
125
126    pub fn push_object(&mut self, hash: u64) {
127        let n = self.objects_len as usize;
128        if n < MAX_SIEVE_ALLOW {
129            self.objects[n] = hash;
130            self.objects_len += 1;
131        }
132    }
133
134    /// Default graph-mutation triple for clinical / conduct intents.
135    pub fn graph_mutation_default() -> Self {
136        let mut s = Self::EMPTY;
137        s.push_subject(q_hash("schema:Patient"));
138        s.push_subject(q_hash("q42:subject"));
139        s.push_predicate(q_hash("snomed:hasFever"));
140        s.push_predicate(q_hash("q42:conductViolation"));
141        s.push_predicate(q_hash("q42:hasGuardian"));
142        s.push_object(q_hash("xsd:true"));
143        s.push_object(q_hash("q42:entity"));
144        s
145    }
146
147    /// Fever observation triple used by the e2e WAL pipeline test.
148    pub fn fever_observation() -> Self {
149        let mut s = Self::EMPTY;
150        s.push_subject(q_hash("Patient"));
151        s.push_predicate(q_hash("fever"));
152        s.push_object(q_hash("True"));
153        s
154    }
155}
156
157/// Neuro-symbolic grammar sieve (masks built once from tokenizer + lex — cold path may alloc).
158#[derive(Debug, Clone)]
159pub struct NeuroSymbolicSieve {
160    masks: [SieveStateMask; 3],
161    state: SieveState,
162    subject_hash: u64,
163    predicate_hash: u64,
164    object_hash: u64,
165    emitted_tokens: [u32; 3],
166    emitted_len: u8,
167}
168
169impl NeuroSymbolicSieve {
170    /// Build masks from a memory-mapped `.q42.lex` view and GGUF tokenizer (load-time only).
171    #[cfg(not(target_arch = "wasm32"))]
172    pub fn from_lex_and_tokenizer(
173        lex: &crate::q42_lex::Q42LexMmap<'_>,
174        tok: &crate::gguf_sharder::GgufTokenizer,
175        spec: &SieveLexSpec,
176    ) -> Self {
177        let mut sieve = Self::empty_fsm();
178        fill_mask_from_lex(
179            &mut sieve.masks[0],
180            lex,
181            tok,
182            &spec.subjects[..spec.subjects_len as usize],
183        );
184        fill_mask_from_lex(
185            &mut sieve.masks[1],
186            lex,
187            tok,
188            &spec.predicates[..spec.predicates_len as usize],
189        );
190        fill_mask_from_lex(
191            &mut sieve.masks[2],
192            lex,
193            tok,
194            &spec.objects[..spec.objects_len as usize],
195        );
196        sieve
197    }
198
199    /// Tokenizer-only fallback when no `.q42.lex` sidecar is loaded (dev / tests).
200    pub fn from_gguf_tokenizer(tok: &crate::gguf_sharder::GgufTokenizer) -> Self {
201        let mut sieve = Self::empty_fsm();
202        const SUBJECTS: &[(&str, u64)] = &[
203            ("Webizen", q_hash("q42:webizenAgent")),
204            ("Agent", q_hash("q42:agent")),
205            ("Subject", q_hash("q42:subject")),
206        ];
207        const PREDICATES: &[(&str, u64)] = &[
208            ("conductViolation", q_hash("q42:conductViolation")),
209            ("hasGuardian", q_hash("q42:hasGuardian")),
210            ("violation", q_hash("q42:conductViolation")),
211            ("Guardian", q_hash("q42:hasGuardian")),
212        ];
213        const OBJECTS: &[(&str, u64)] = &[
214            ("guardian", q_hash("q42:guardianEntity")),
215            ("Entity", q_hash("q42:entity")),
216            ("Object", q_hash("q42:object")),
217        ];
218        fill_mask_literal(&mut sieve.masks[0], tok, SUBJECTS);
219        fill_mask_literal(&mut sieve.masks[1], tok, PREDICATES);
220        fill_mask_literal(&mut sieve.masks[2], tok, OBJECTS);
221        sieve
222    }
223
224    pub(crate) fn empty_fsm() -> Self {
225        Self {
226            masks: [SieveStateMask::EMPTY; 3],
227            state: SieveState::ExpectSubject,
228            subject_hash: 0,
229            predicate_hash: 0,
230            object_hash: 0,
231            emitted_tokens: [0; 3],
232            emitted_len: 0,
233        }
234    }
235
236    #[inline]
237    pub fn state(&self) -> SieveState {
238        self.state
239    }
240
241    #[inline]
242    pub fn is_complete(&self) -> bool {
243        self.state == SieveState::Complete
244    }
245
246    #[inline]
247    pub fn emitted_len(&self) -> u8 {
248        self.emitted_len
249    }
250
251    #[inline]
252    pub fn current_mask(&self) -> &SieveStateMask {
253        match self.state {
254            SieveState::ExpectSubject => &self.masks[0],
255            SieveState::ExpectPredicate => &self.masks[1],
256            SieveState::ExpectObject => &self.masks[2],
257            SieveState::Complete => &SieveStateMask::EMPTY,
258        }
259    }
260
261    /// Apply a sieve-selected token and advance the FSM.
262    pub fn apply_token(&mut self, token_id: u32) -> Result<(), SieveError> {
263        if self.state == SieveState::Complete {
264            return Err(SieveError::AlreadyComplete);
265        }
266        let mask = self.current_mask();
267        if mask.len == 0 {
268            return Err(SieveError::Misaligned);
269        }
270        let hash = mask
271            .lexicon_hash_for(token_id)
272            .ok_or(SieveError::Misaligned)?;
273        match self.state {
274            SieveState::ExpectSubject => {
275                self.subject_hash = hash;
276                self.state = SieveState::ExpectPredicate;
277            }
278            SieveState::ExpectPredicate => {
279                self.predicate_hash = hash;
280                self.state = SieveState::ExpectObject;
281            }
282            SieveState::ExpectObject => {
283                self.object_hash = hash;
284                self.state = SieveState::Complete;
285            }
286            SieveState::Complete => return Err(SieveError::AlreadyComplete),
287        }
288        let n = self.emitted_len as usize;
289        if n < 3 {
290            self.emitted_tokens[n] = token_id;
291            self.emitted_len += 1;
292        }
293        Ok(())
294    }
295
296    /// Assemble the 48-byte `NQuin` from constrained emissions (stack only).
297    pub fn assemble_quin(&self, context_hash: u64) -> NQuin {
298        let mut quin = NQuin {
299            subject: self.subject_hash,
300            predicate: self.predicate_hash,
301            object: self.object_hash,
302            context: context_hash,
303            metadata: 0,
304            parity: 0,
305        };
306        quin.parity = quin.subject ^ quin.predicate ^ quin.object ^ quin.context;
307        quin
308    }
309
310    pub fn masks_ready(&self) -> bool {
311        self.masks[0].len > 0 && self.masks[1].len > 0 && self.masks[2].len > 0
312    }
313
314    /// First resolved token ID per FSM slot (cold-path / test helper).
315    pub fn resolved_token_triple(&self) -> Option<(u32, u32, u32)> {
316        if !self.masks_ready() {
317            return None;
318        }
319        Some((
320            self.masks[0].slots[0].token_id,
321            self.masks[1].slots[0].token_id,
322            self.masks[2].slots[0].token_id,
323        ))
324    }
325}
326
327fn fill_mask_literal(
328    mask: &mut SieveStateMask,
329    tok: &crate::gguf_sharder::GgufTokenizer,
330    entries: &[(&str, u64)],
331) {
332    for &(text, hash) in entries {
333        let ids = tok.encode(text);
334        if let Some(&id) = ids.first() {
335            mask.push(id, hash);
336        }
337    }
338}
339
340#[cfg(not(target_arch = "wasm32"))]
341fn fill_mask_from_lex(
342    mask: &mut SieveStateMask,
343    lex: &crate::q42_lex::Q42LexMmap<'_>,
344    tok: &crate::gguf_sharder::GgufTokenizer,
345    hashes: &[u64],
346) {
347    for &hash in hashes {
348        if let Some(text) = lex.lookup_hash(hash) {
349            let ids = tok.encode(text);
350            if let Some(&id) = ids.first() {
351                mask.push(id, hash);
352            }
353        }
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360
361    fn write_lex_bytes(entries: &[(u64, &str)]) -> Vec<u8> {
362        let mut sorted: Vec<(u64, &str)> = entries.to_vec();
363        sorted.sort_unstable_by_key(|(h, _)| *h);
364        let entry_count = sorted.len() as u64;
365        let strings_offset = 32 + entry_count * 16;
366        let mut blob = Vec::new();
367        let mut index = Vec::new();
368        for (hash, text) in &sorted {
369            let str_off = blob.len() as u64;
370            let b = text.as_bytes();
371            let len = b.len().min(65535) as u16;
372            // Q42LEX string payload format: [LEX_TAG_STRING=0x01][len_lo][len_hi][utf8...]
373            // Must match Q42LexMmap::read_string_at which checks the tag byte first.
374            blob.push(0x01u8); // LEX_TAG_STRING
375            blob.extend_from_slice(&len.to_le_bytes());
376            blob.extend_from_slice(&b[..len as usize]);
377            index.extend_from_slice(&hash.to_le_bytes());
378            index.extend_from_slice(&str_off.to_le_bytes());
379        }
380        let mut out = Vec::new();
381        out.extend_from_slice(b"Q42LEX\0\0");
382        out.extend_from_slice(&entry_count.to_le_bytes());
383        out.extend_from_slice(&strings_offset.to_le_bytes());
384        out.extend_from_slice(&1u64.to_le_bytes());
385        out.extend_from_slice(&index);
386        out.extend_from_slice(&blob);
387        out
388    }
389
390    #[test]
391    fn sieve_mask_allows_linear_scan() {
392        let mut m = SieveStateMask::EMPTY;
393        m.push(42, q_hash("a"));
394        m.push(99, q_hash("b"));
395        assert!(m.allows(42));
396        assert!(!m.allows(1));
397        assert_eq!(m.lexicon_hash_for(99), Some(q_hash("b")));
398    }
399
400    #[test]
401    fn sieve_fsm_transitions_to_complete() {
402        let mut s = NeuroSymbolicSieve::empty_fsm();
403        s.masks[0].push(10, q_hash("sub"));
404        s.masks[1].push(20, q_hash("pred"));
405        s.masks[2].push(30, q_hash("obj"));
406        assert!(s.apply_token(10).is_ok());
407        assert!(s.apply_token(20).is_ok());
408        assert!(s.apply_token(30).is_ok());
409        assert!(s.is_complete());
410        let q = s.assemble_quin(q_hash("ctx"));
411        assert_eq!(q.subject, q_hash("sub"));
412        assert_eq!(q.predicate, q_hash("pred"));
413        assert_eq!(q.object, q_hash("obj"));
414    }
415
416    #[test]
417    fn sieve_builds_masks_from_mmap_lex() {
418        let h_sub = q_hash("Patient");
419        let h_pred = q_hash("fever");
420        let h_obj = q_hash("True");
421        let bytes = write_lex_bytes(&[(h_sub, "Patient"), (h_pred, "fever"), (h_obj, "True")]);
422        let lex = crate::q42_lex::Q42LexMmap::from_bytes(&bytes).unwrap();
423        let tok = crate::gguf_sharder::GgufTokenizer::default();
424        let spec = SieveLexSpec::fever_observation();
425        let sieve = NeuroSymbolicSieve::from_lex_and_tokenizer(&lex, &tok, &spec);
426        assert!(sieve.masks_ready());
427    }
428}