Skip to main content

vyre_runtime/megakernel/
rule_catalog.rs

1//! DFA rule catalog packing for batched megakernel dispatch.
2
3use super::staging_reserve::{
4    reserve_hash_map_capacity as reserve_catalog_map, reserve_vec_capacity as reserve_catalog_vec,
5};
6use crate::PipelineError;
7use rustc_hash::FxHashMap;
8
9/// Dense byte alphabet used by the DFA transition table as the INPUT
10/// (`BatchRuleProgram`) representation: every rule still arrives as a dense
11/// `state * 256 + byte` table. The on-device packed table is byte-class
12/// compressed (see [`pack_rule_catalog_into`]); this constant is the source
13/// alphabet width the compressor folds DOWN from.
14pub const ALPHABET_SIZE: u32 = 256;
15const ALPHABET_SIZE_USIZE: usize = 256;
16
17/// Number of `u32` words per rule metadata entry. The kernel reads these in
18/// order: `transition_base`, `accept_base`, `state_count`, `class_map_base`,
19/// `num_classes`. Bump in lockstep with [`RuleMeta`] and the dispatcher's
20/// `dfa_byte_scanner` if the per-rule metadata grows.
21pub const RULE_META_WORDS: usize = 5;
22
23/// One compiled DFA-backed rule program consumed by the batch dispatcher.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct BatchRuleProgram {
26    /// Stable rule-table index.
27    pub rule_idx: u32,
28    /// Dense DFA transition table (`state * 256 + byte -> next_state`).
29    pub transitions: Vec<u32>,
30    /// Dense DFA accept table (`state -> non-zero match marker`).
31    pub accept: Vec<u32>,
32    /// DFA state count.
33    pub state_count: u32,
34}
35
36impl BatchRuleProgram {
37    /// Build one DFA-backed rule program.
38    ///
39    /// # Errors
40    ///
41    /// Returns [`PipelineError::Backend`] when the DFA buffers do not match
42    /// `state_count`.
43    pub fn new(
44        rule_idx: u32,
45        transitions: Vec<u32>,
46        accept: Vec<u32>,
47        state_count: u32,
48    ) -> Result<Self, PipelineError> {
49        validate_rule_shape(rule_idx, &transitions, &accept, state_count)?;
50        Ok(Self {
51            rule_idx,
52            transitions,
53            accept,
54            state_count,
55        })
56    }
57}
58
59/// Packed metadata for one byte-class-compressed DFA rule entry.
60///
61/// The on-device transition table is `transitions[transition_base + state *
62/// num_classes + class_maps[class_map_base + byte]]`: each rule carries a
63/// 256-entry byte→class map (into the shared `class_maps` buffer) and a
64/// compressed `state_count * num_classes` transition block, instead of a dense
65/// `state_count * 256` block. The compression is LOSSLESS, bytes share a class
66/// only when their transition column is identical across every state, so GPU
67/// firings are byte-for-byte identical to the dense table.
68#[repr(C)]
69#[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)]
70pub struct RuleMeta {
71    /// Word offset into the flattened (compressed) transition table.
72    pub transition_base: u32,
73    /// Word offset into the flattened accept table.
74    pub accept_base: u32,
75    /// DFA state count for this rule.
76    pub state_count: u32,
77    /// Word offset into the shared 256-entry-per-rule byte→class map table.
78    pub class_map_base: u32,
79    /// Number of distinct byte classes for this rule (the compressed row width).
80    pub num_classes: u32,
81}
82
83/// One rule rejected from a megakernel batch while other rules still ran.
84#[derive(Debug, Clone, PartialEq, Eq)]
85pub struct BatchRuleRejection {
86    /// Caller-supplied rule index when present.
87    pub rule_idx: Option<u32>,
88    /// Human-readable rejection reason.
89    pub reason: String,
90}
91
92/// Packed rule catalog uploaded to device storage buffers.
93pub struct PackedRuleCatalog {
94    /// Dense per-rule metadata table.
95    pub rule_meta: Vec<RuleMeta>,
96    /// Deduplicated flattened byte-class-COMPRESSED DFA transition storage.
97    /// Indexed `transitions[rule.transition_base + state * rule.num_classes +
98    /// class]` where `class = class_maps[rule.class_map_base + byte]`.
99    pub transitions: Vec<u32>,
100    /// Deduplicated flattened DFA accept storage.
101    pub accept: Vec<u32>,
102    /// Deduplicated flattened 256-entry-per-rule byte→class maps. Indexed
103    /// `class_maps[rule.class_map_base + byte]`.
104    pub class_maps: Vec<u32>,
105    /// Rules rejected during validation or dense-slot assignment.
106    pub rejected_rules: Vec<BatchRuleRejection>,
107}
108
109/// Caller-owned storage for packing rule catalogs without rebuilding host
110/// allocations on every refresh.
111#[derive(Default)]
112pub struct RuleCatalogPackingScratch {
113    /// Dense per-rule metadata table.
114    pub rule_meta: Vec<RuleMeta>,
115    /// Deduplicated flattened byte-class-COMPRESSED DFA transition storage.
116    pub transitions: Vec<u32>,
117    /// Deduplicated flattened DFA accept storage.
118    pub accept: Vec<u32>,
119    /// Deduplicated flattened 256-entry-per-rule byte→class maps.
120    pub class_maps: Vec<u32>,
121    /// Rules rejected during validation or dense-slot assignment.
122    pub rejected_rules: Vec<BatchRuleRejection>,
123    /// fingerprint -> (transition_base, accept_base, state_count,
124    /// class_map_base, num_classes) for storage dedup across identical DFAs.
125    unique_storage: FxHashMap<[u8; 32], UniqueStorageLayout>,
126    occupied: Vec<bool>,
127    addressed: Vec<bool>,
128    /// Reusable 256-entry byte→class scratch built per unique DFA.
129    class_scratch: Vec<u32>,
130}
131
132/// Resident-buffer layout for one deduplicated unique DFA storage block.
133#[derive(Clone, Copy)]
134struct UniqueStorageLayout {
135    transition_base: u32,
136    accept_base: u32,
137    state_count: u32,
138    class_map_base: u32,
139    num_classes: u32,
140}
141
142/// Fingerprints for the valid dense catalog entries.
143#[must_use]
144pub fn accepted_rule_fingerprints(
145    rules: &[BatchRuleProgram],
146) -> (Vec<[u8; 32]>, Vec<BatchRuleRejection>) {
147    let mut fingerprints = Vec::new();
148    let mut occupied = Vec::new();
149    let mut addressed = Vec::new();
150    let rejections =
151        accepted_rule_fingerprints_into(rules, &mut fingerprints, &mut occupied, &mut addressed);
152    (fingerprints, rejections)
153}
154
155/// Fill caller-owned storage with fingerprints for valid dense catalog entries.
156///
157/// The output fingerprint order matches dense rule-table order, not input
158/// order. `fingerprints`, `occupied`, and `addressed` are cleared and reused so
159/// dispatchers can check resident catalog identity without allocating on every
160/// cache-hit dispatch.
161pub fn accepted_rule_fingerprints_into(
162    rules: &[BatchRuleProgram],
163    fingerprints: &mut Vec<[u8; 32]>,
164    occupied: &mut Vec<bool>,
165    addressed: &mut Vec<bool>,
166) -> Vec<BatchRuleRejection> {
167    let mut rejections = Vec::new();
168    accepted_rule_fingerprints_and_rejections_into(
169        rules,
170        fingerprints,
171        occupied,
172        addressed,
173        &mut rejections,
174    );
175    rejections
176}
177
178/// Fill caller-owned storage with fingerprints and rejection details for valid
179/// dense catalog entries.
180///
181/// This is the allocation-stable form used by hot dispatchers. All scratch
182/// vectors are cleared and reused; valid unchanged catalogs perform no host
183/// allocations while checking resident rule-buffer identity.
184pub fn accepted_rule_fingerprints_and_rejections_into(
185    rules: &[BatchRuleProgram],
186    fingerprints: &mut Vec<[u8; 32]>,
187    occupied: &mut Vec<bool>,
188    addressed: &mut Vec<bool>,
189    rejections: &mut Vec<BatchRuleRejection>,
190) {
191    fingerprints.clear();
192    fingerprints.resize(rules.len(), [0; 32]);
193    occupied.clear();
194    occupied.resize(rules.len(), false);
195    addressed.clear();
196    addressed.resize(rules.len(), false);
197    rejections.clear();
198
199    for rule in rules {
200        mark_addressed(addressed, rule.rule_idx);
201        match validate_rule_shape(
202            rule.rule_idx,
203            &rule.transitions,
204            &rule.accept,
205            rule.state_count,
206        ) {
207            Ok(()) => match claim_dense_index(occupied, rule.rule_idx, rules.len()) {
208                Ok(index) => fingerprints[index] = rule_fingerprint(rule),
209                Err(rejection) => rejections.push(rejection),
210            },
211            Err(error) => rejections.push(BatchRuleRejection {
212                rule_idx: Some(rule.rule_idx),
213                reason: error.to_string(),
214            }),
215        }
216    }
217
218    extend_missing_rejections(occupied, addressed, rejections);
219    let mut write = 0;
220    for read in 0..occupied.len() {
221        if occupied[read] {
222            fingerprints[write] = fingerprints[read];
223            write += 1;
224        }
225    }
226    fingerprints.truncate(write);
227}
228
229/// Pack valid DFA rules into compact shared device tables.
230///
231/// Rules with identical `(transitions, accept, state_count)` share backing
232/// transition and accept storage while retaining distinct dense metadata slots.
233pub fn pack_rule_catalog(rules: &[BatchRuleProgram]) -> Result<PackedRuleCatalog, PipelineError> {
234    let mut scratch = RuleCatalogPackingScratch::default();
235    pack_rule_catalog_into(rules, &mut scratch)?;
236    Ok(PackedRuleCatalog {
237        rule_meta: scratch.rule_meta,
238        transitions: scratch.transitions,
239        accept: scratch.accept,
240        class_maps: scratch.class_maps,
241        rejected_rules: scratch.rejected_rules,
242    })
243}
244
245/// Pack valid DFA rules into caller-owned storage.
246///
247/// Existing vector and hash-map allocations in `scratch` are reused across
248/// calls. This is the hot-path form for resident megakernel dispatchers that
249/// refresh device rule buffers repeatedly.
250pub fn pack_rule_catalog_into(
251    rules: &[BatchRuleProgram],
252    scratch: &mut RuleCatalogPackingScratch,
253) -> Result<(), PipelineError> {
254    scratch.unique_storage.clear();
255    reserve_catalog_map(
256        &mut scratch.unique_storage,
257        rules.len(),
258        "unique DFA storage",
259    )?;
260    // Inert rule slot 0: a 1-state DFA that self-loops on every byte and never
261    // accepts. Compressed it is num_classes=1, a 1-word transition row `[0]`,
262    // and a 256-entry all-zero byte→class map. Rejected / missing rules point
263    // their metadata here so the kernel reads a well-formed (no-match) DFA
264    // instead of out-of-bounds storage.
265    scratch.transitions.clear();
266    reserve_catalog_vec(&mut scratch.transitions, 1, "inert transition row")?;
267    scratch.transitions.push(0);
268    scratch.accept.clear();
269    reserve_catalog_vec(&mut scratch.accept, 1, "inert accept row")?;
270    scratch.accept.push(0);
271    scratch.class_maps.clear();
272    reserve_catalog_vec(
273        &mut scratch.class_maps,
274        ALPHABET_SIZE_USIZE,
275        "inert byte-class map",
276    )?;
277    scratch.class_maps.resize(ALPHABET_SIZE_USIZE, 0);
278    scratch.rule_meta.clear();
279    reserve_catalog_vec(&mut scratch.rule_meta, rules.len(), "rule metadata")?;
280    scratch.rule_meta.resize(
281        rules.len(),
282        RuleMeta {
283            transition_base: 0,
284            accept_base: 0,
285            state_count: 1,
286            class_map_base: 0,
287            num_classes: 1,
288        },
289    );
290    scratch.rejected_rules.clear();
291    reserve_catalog_vec(
292        &mut scratch.rejected_rules,
293        rules.len(),
294        "rule rejection rows",
295    )?;
296    scratch.occupied.clear();
297    reserve_catalog_vec(&mut scratch.occupied, rules.len(), "dense occupancy bitmap")?;
298    scratch.occupied.resize(rules.len(), false);
299    scratch.addressed.clear();
300    reserve_catalog_vec(
301        &mut scratch.addressed,
302        rules.len(),
303        "dense addressed bitmap",
304    )?;
305    scratch.addressed.resize(rules.len(), false);
306
307    for rule in rules {
308        mark_addressed(&mut scratch.addressed, rule.rule_idx);
309        if let Err(error) = validate_rule_shape(
310            rule.rule_idx,
311            &rule.transitions,
312            &rule.accept,
313            rule.state_count,
314        ) {
315            scratch.rejected_rules.push(BatchRuleRejection {
316                rule_idx: Some(rule.rule_idx),
317                reason: error.to_string(),
318            });
319            continue;
320        }
321
322        let meta_index = match claim_dense_index(
323            &mut scratch.occupied,
324            rule.rule_idx,
325            scratch.rule_meta.len(),
326        ) {
327            Ok(index) => index,
328            Err(rejection) => {
329                scratch.rejected_rules.push(rejection);
330                continue;
331            }
332        };
333
334        let storage_fingerprint = dfa_storage_fingerprint(rule);
335        let layout = if let Some(layout) = scratch.unique_storage.get(&storage_fingerprint) {
336            *layout
337        } else {
338            // Build the LOSSLESS byte→class map for this DFA into reusable
339            // scratch, then emit the compressed `state * num_classes + class`
340            // transition block. `num_classes <= 256`, with equality only when
341            // every byte transitions differently in some state, the common
342            // secret-detector DFAs collapse to a handful of classes.
343            let num_classes = build_byte_class_map_for_table(
344                &rule.transitions,
345                rule.state_count as usize,
346                &mut scratch.class_scratch,
347            );
348
349            let class_map_base =
350                u32::try_from(scratch.class_maps.len()).map_err(|_| PipelineError::QueueFull {
351                    queue: "submission",
352                    fix: "flattened byte-class map table exceeds u32::MAX words; split the rule catalog into smaller groups",
353                })?;
354            let class_map_target = scratch
355                .class_maps
356                .len()
357                .checked_add(ALPHABET_SIZE_USIZE)
358                .ok_or(PipelineError::QueueFull {
359                    queue: "submission",
360                    fix: "flattened byte-class map length overflows usize; split the rule catalog into smaller groups",
361                })?;
362            reserve_catalog_vec(
363                &mut scratch.class_maps,
364                class_map_target,
365                "flattened byte-class map storage",
366            )?;
367            scratch.class_maps.extend_from_slice(&scratch.class_scratch);
368
369            let transition_base =
370                u32::try_from(scratch.transitions.len()).map_err(|_| PipelineError::QueueFull {
371                    queue: "submission",
372                    fix: "flattened transition table exceeds u32::MAX words; split the rule catalog into smaller groups",
373                })?;
374            let accept_base = u32::try_from(scratch.accept.len()).map_err(|_| PipelineError::QueueFull {
375                queue: "submission",
376                fix: "flattened accept table exceeds u32::MAX words; split the rule catalog into smaller groups",
377            })?;
378            // Compressed block size = state_count * num_classes. Both are
379            // bounded (state_count is validated, num_classes <= 256), so the
380            // product cannot exceed the dense state_count * 256 size that
381            // already validated.
382            let compressed_words = (rule.state_count as usize)
383                .checked_mul(num_classes as usize)
384                .ok_or(PipelineError::QueueFull {
385                    queue: "submission",
386                    fix: "compressed transition block size overflows usize; split the rule catalog into smaller groups",
387                })?;
388            let transition_target = scratch
389                .transitions
390                .len()
391                .checked_add(compressed_words)
392                .ok_or(PipelineError::QueueFull {
393                    queue: "submission",
394                    fix: "flattened transition table length overflows usize; split the rule catalog into smaller groups",
395                })?;
396            reserve_catalog_vec(
397                &mut scratch.transitions,
398                transition_target,
399                "flattened transition storage",
400            )?;
401            let accept_target = scratch
402                .accept
403                .len()
404                .checked_add(rule.accept.len())
405                .ok_or(PipelineError::QueueFull {
406                    queue: "submission",
407                    fix: "flattened accept table length overflows usize; split the rule catalog into smaller groups",
408                })?;
409            reserve_catalog_vec(
410                &mut scratch.accept,
411                accept_target,
412                "flattened accept storage",
413            )?;
414            // Emit the compressed `state * num_classes + class` transition block
415            // via the shared primitive (lossless: every byte in a class shares
416            // its dense column).
417            compress_dense_transitions_into(
418                &rule.transitions,
419                rule.state_count as usize,
420                &scratch.class_scratch,
421                num_classes,
422                &mut scratch.transitions,
423            );
424            scratch.accept.extend_from_slice(&rule.accept);
425
426            let layout = UniqueStorageLayout {
427                transition_base,
428                accept_base,
429                state_count: rule.state_count,
430                class_map_base,
431                num_classes,
432            };
433            scratch.unique_storage.insert(storage_fingerprint, layout);
434            layout
435        };
436        scratch.rule_meta[meta_index] = RuleMeta {
437            transition_base: layout.transition_base,
438            accept_base: layout.accept_base,
439            state_count: layout.state_count,
440            class_map_base: layout.class_map_base,
441            num_classes: layout.num_classes,
442        };
443    }
444
445    extend_missing_rejections(
446        &scratch.occupied,
447        &scratch.addressed,
448        &mut scratch.rejected_rules,
449    );
450    Ok(())
451}
452
453/// Build the LOSSLESS byte→class map for one dense DFA into `out` (resized to
454/// 256) and return the class count.
455///
456/// Two bytes share a class iff their transition COLUMN is identical across every
457/// state: `transitions[s*256 + a] == transitions[s*256 + b]` for all `s`. Class
458/// ids are assigned in order of first byte appearance, so the map is `0` for
459/// byte 0's class and deterministic. The returned width `num_classes` is `<=
460/// 256`; for the secret-detector DFAs (long fixed prefixes + a few char
461/// classes) it collapses to a handful, shrinking the per-state transition row
462/// from 256 words to `num_classes` words, a lossless ~16x reduction on the
463/// ~987 MB catalog without changing a single firing.
464///
465/// `out` is cleared/reused so the hot resident-refresh path allocates nothing.
466/// Build the LOSSLESS byte→class map for a dense `state_count * 256` DFA
467/// transition table into `out` (resized to 256) and return the class count.
468///
469/// Two bytes share a class iff their transition COLUMN is identical across every
470/// state: `transitions[s*256 + a] == transitions[s*256 + b]` for all `s`. Class
471/// ids are assigned in order of first byte appearance, so the map is
472/// deterministic. The returned `num_classes` is `<= 256`.
473///
474/// This is the shared compression primitive: the per-rule catalog packer and
475/// the combined-AC megakernel both call it so a single definition owns the
476/// "identical column ⇒ same class" contract. `out` is cleared/reused so hot
477/// paths allocate nothing.
478#[must_use]
479pub fn build_byte_class_map_for_table(
480    transitions: &[u32],
481    state_count: usize,
482    out: &mut Vec<u32>,
483) -> u32 {
484    out.clear();
485    out.resize(ALPHABET_SIZE_USIZE, 0);
486    // Column signature per byte = its next-state across every state. Group by
487    // signature. `FxHashMap` keyed on the column bytes; the first byte to
488    // produce a signature owns a fresh class id.
489    let mut signature_to_class: FxHashMap<Vec<u32>, u32> = FxHashMap::default();
490    let mut next_class: u32 = 0;
491    let mut signature = Vec::with_capacity(state_count);
492    for byte in 0..ALPHABET_SIZE_USIZE {
493        signature.clear();
494        for state in 0..state_count {
495            signature.push(transitions[state * ALPHABET_SIZE_USIZE + byte]);
496        }
497        let class = match signature_to_class.get(&signature) {
498            Some(&class) => class,
499            None => {
500                let class = next_class;
501                next_class += 1;
502                signature_to_class.insert(signature.clone(), class);
503                class
504            }
505        };
506        out[byte] = class;
507    }
508    next_class
509}
510
511/// Append the compressed `state_count * num_classes` transition block for a
512/// dense `state_count * 256` table to `out`, given the byte→class `class_map`
513/// from [`build_byte_class_map_for_table`].
514///
515/// For class `c` it copies the dense column of ANY byte mapping to `c` (the
516/// first; every byte in a class has an identical column by construction, so the
517/// value is well-defined and LOSSLESS). Shared by the per-rule packer and the
518/// combined-AC megakernel.
519pub fn compress_dense_transitions_into(
520    dense: &[u32],
521    state_count: usize,
522    class_map: &[u32],
523    num_classes: u32,
524    out: &mut Vec<u32>,
525) {
526    let num_classes = num_classes as usize;
527    let mut class_representative = vec![0usize; num_classes];
528    let mut seen = vec![false; num_classes];
529    for (byte, &class) in class_map.iter().enumerate() {
530        let class = class as usize;
531        if !seen[class] {
532            seen[class] = true;
533            class_representative[class] = byte;
534        }
535    }
536    for state in 0..state_count {
537        let dense_row = state * ALPHABET_SIZE_USIZE;
538        for &rep_byte in &class_representative {
539            out.push(dense[dense_row + rep_byte]);
540        }
541    }
542}
543
544/// Pack a byte-class-compressed `state_count * num_classes` transition table
545/// (from [`compress_dense_transitions_into`]) into u16 targets stored two per
546/// u32 word: the LOW half holds the even flat index, the HIGH half the odd
547/// index. This halves the device transition footprint and narrows every
548/// transition-table read. Row deduplication does not provide this bound.
549///
550/// FAIL CLOSED (Law 10): every transition target is a state index, so it must
551/// fit `u16`. If ANY target exceeds `u16::MAX` this REFUSES the pack, a silent
552/// `& 0xFFFF` truncation would repoint a next-state at the wrong state, an
553/// invisible recall loss. Callers gate on `state_count <= 65536`; this is the
554/// enforcing check, not an assumption. `out` is cleared/reused so the hot
555/// resident-refresh path allocates nothing.
556///
557/// # Errors
558///
559/// Returns [`PipelineError::Backend`] naming the offending index/target when a
560/// transition target does not fit `u16`.
561pub fn try_pack_u16_transitions_into(
562    compressed: &[u32],
563    out: &mut Vec<u32>,
564) -> Result<(), PipelineError> {
565    for (idx, &target) in compressed.iter().enumerate() {
566        if target > u32::from(u16::MAX) {
567            return Err(PipelineError::Backend(format!(
568                "transition target {target} at index {idx} exceeds u16::MAX ({}); u16 packing \
569                 would silently truncate it and corrupt the automaton. Fix: keep this catalog on \
570                 the u32 transition path (state_count must be <= 65536 for u16 packing).",
571                u16::MAX
572            )));
573        }
574    }
575    out.clear();
576    out.reserve(compressed.len().div_ceil(2));
577    let mut chunks = compressed.chunks_exact(2);
578    for pair in chunks.by_ref() {
579        // Both halves validated <= 0xFFFF above, so no masking is needed.
580        out.push(pair[0] | (pair[1] << 16));
581    }
582    if let [last] = chunks.remainder() {
583        out.push(*last);
584    }
585    Ok(())
586}
587
588/// Unpack the `flat_index`-th u16 transition target from a table packed by
589/// [`try_pack_u16_transitions_into`]. The exact CPU mirror of the kernel's
590/// unpack (`word = packed[idx/2]; (word >> ((idx & 1) * 16)) & 0xFFFF`), so the
591/// round-trip can be proven lossless without a GPU.
592#[must_use]
593pub fn unpack_u16_transition(packed: &[u32], flat_index: usize) -> u32 {
594    let word = packed[flat_index / 2];
595    (word >> ((flat_index & 1) * 16)) & 0xFFFF
596}
597
598fn validate_rule_shape(
599    rule_idx: u32,
600    transitions: &[u32],
601    accept: &[u32],
602    state_count: u32,
603) -> Result<(), PipelineError> {
604    let expected_transitions = usize::try_from(state_count)
605        .ok()
606        .and_then(|count| count.checked_mul(ALPHABET_SIZE_USIZE))
607        .ok_or_else(|| {
608            PipelineError::Backend("rule transition table size overflowed usize".to_string())
609        })?;
610    if transitions.len() != expected_transitions {
611        return Err(PipelineError::Backend(format!(
612            "rule {rule_idx} transition table has {} words, expected {expected_transitions}. Fix: compile a dense state_count * 256 DFA table before batch dispatch.",
613            transitions.len()
614        )));
615    }
616    let state_count_usize = usize::try_from(state_count).map_err(|source| {
617        PipelineError::Backend(format!(
618            "rule {rule_idx} state_count {state_count} cannot fit usize: {source}. Fix: shard the DFA state space before batch dispatch."
619        ))
620    })?;
621    if accept.len() != state_count_usize {
622        return Err(PipelineError::Backend(format!(
623            "rule {rule_idx} accept table has {} words, expected {state_count}. Fix: emit one accept entry per DFA state before batch dispatch.",
624            accept.len()
625        )));
626    }
627    Ok(())
628}
629
630fn rule_fingerprint(rule: &BatchRuleProgram) -> [u8; 32] {
631    let mut hasher = blake3::Hasher::new();
632    hasher.update(&rule.rule_idx.to_le_bytes());
633    hasher.update(bytemuck::cast_slice(&rule.transitions));
634    hasher.update(bytemuck::cast_slice(&rule.accept));
635    hasher.update(&rule.state_count.to_le_bytes());
636    *hasher.finalize().as_bytes()
637}
638
639fn dfa_storage_fingerprint(rule: &BatchRuleProgram) -> [u8; 32] {
640    let mut hasher = blake3::Hasher::new();
641    hasher.update(bytemuck::cast_slice(&rule.transitions));
642    hasher.update(bytemuck::cast_slice(&rule.accept));
643    hasher.update(&rule.state_count.to_le_bytes());
644    *hasher.finalize().as_bytes()
645}
646
647fn mark_addressed(addressed: &mut [bool], rule_idx: u32) {
648    if let Some(index) = usize::try_from(rule_idx)
649        .ok()
650        .filter(|index| *index < addressed.len())
651    {
652        addressed[index] = true;
653    }
654}
655
656fn claim_dense_index(
657    occupied: &mut [bool],
658    rule_idx: u32,
659    slot_count: usize,
660) -> Result<usize, BatchRuleRejection> {
661    let Some(meta_index) = usize::try_from(rule_idx).ok() else {
662        return Err(BatchRuleRejection {
663            rule_idx: Some(rule_idx),
664            reason: "rule_idx exceeds usize. Fix: rebuild the batch with a smaller rule catalog"
665                .to_string(),
666        });
667    };
668    if meta_index >= slot_count {
669        return Err(BatchRuleRejection {
670            rule_idx: Some(rule_idx),
671            reason: format!(
672                "rule_idx {rule_idx} falls outside 0..{slot_count}. Fix: keep the rule catalog dense so the batch work queue can address every rule"
673            ),
674        });
675    }
676    if occupied[meta_index] {
677        return Err(BatchRuleRejection {
678            rule_idx: Some(rule_idx),
679            reason: format!(
680                "duplicate rule_idx {rule_idx}. Fix: keep exactly one rule per dense rule-table slot"
681            ),
682        });
683    }
684    occupied[meta_index] = true;
685    Ok(meta_index)
686}
687
688fn extend_missing_rejections(
689    occupied: &[bool],
690    addressed: &[bool],
691    out: &mut Vec<BatchRuleRejection>,
692) {
693    for (rule_idx, (occupied, addressed)) in occupied
694        .iter()
695        .copied()
696        .zip(addressed.iter().copied())
697        .enumerate()
698    {
699        if !occupied && !addressed {
700            let Ok(rule_idx_u32) = u32::try_from(rule_idx) else {
701                continue;
702            };
703            out.push(BatchRuleRejection {
704                rule_idx: Some(rule_idx_u32),
705                reason: format!(
706                    "rule_idx {rule_idx} has no valid catalog entry. Fix: provide a well-formed DFA for every dense rule slot before batch dispatch"
707                ),
708            });
709        }
710    }
711}
712
713#[cfg(test)]
714mod tests {
715    use super::*;
716
717    /// Resolve the next state the COMPRESSED packed catalog yields for
718    /// `(rule, state, byte)`: mirrors the GPU kernel's index math exactly so
719    /// the parity tests can prove byte-for-byte equivalence to the dense table.
720    fn packed_next_state(
721        packed: &PackedRuleCatalog,
722        meta_index: usize,
723        state: u32,
724        byte: u8,
725    ) -> u32 {
726        let meta = packed.rule_meta[meta_index];
727        let class = packed.class_maps[meta.class_map_base as usize + byte as usize];
728        let idx = meta.transition_base as usize
729            + state as usize * meta.num_classes as usize
730            + class as usize;
731        packed.transitions[idx]
732    }
733
734    #[test]
735    fn duplicate_dfas_share_catalog_storage() {
736        let first = BatchRuleProgram::new(0, vec![0; 256], vec![0], 1).unwrap();
737        let second = BatchRuleProgram::new(1, vec![0; 256], vec![0], 1).unwrap();
738        let packed = pack_rule_catalog(&[first, second]).unwrap();
739        // Identical DFAs share compressed transition, accept AND class-map storage.
740        assert_eq!(
741            packed.rule_meta[0].transition_base,
742            packed.rule_meta[1].transition_base
743        );
744        assert_eq!(
745            packed.rule_meta[0].accept_base,
746            packed.rule_meta[1].accept_base
747        );
748        assert_eq!(
749            packed.rule_meta[0].class_map_base,
750            packed.rule_meta[1].class_map_base
751        );
752        assert_eq!(
753            packed.rule_meta[0].num_classes,
754            packed.rule_meta[1].num_classes
755        );
756        // An all-zero 1-state DFA collapses to a SINGLE byte class (every byte
757        // self-loops to state 0), so its compressed row is exactly one word, not
758        // 256. transition_base points just past the 1-word inert row.
759        assert_eq!(packed.rule_meta[0].num_classes, 1);
760        assert_eq!(packed.rule_meta[0].transition_base, 1);
761        assert_eq!(
762            packed.transitions.len(),
763            packed.rule_meta[0].transition_base as usize + 1
764        );
765        assert_eq!(
766            packed.accept.len(),
767            packed.rule_meta[0].accept_base as usize + 1
768        );
769        assert!(packed.rejected_rules.is_empty());
770    }
771
772    #[test]
773    fn u16_pack_round_trips_losslessly_including_odd_tail() {
774        // Odd element count exercises the lone-remainder path; 65535 is the max
775        // legal u16 target.
776        let compressed: Vec<u32> = vec![0, 1, 2, 65_535, 100, 0, 42];
777        let mut packed = Vec::new();
778        try_pack_u16_transitions_into(&compressed, &mut packed).expect("all targets fit u16");
779        // Two u16 per u32 word, rounding up for the odd tail.
780        assert_eq!(packed.len(), compressed.len().div_ceil(2));
781        // Every flat index unpacks to EXACTLY the original target, proving the
782        // pack→kernel-unpack round-trip changes no transition (Law 6).
783        for (idx, &original) in compressed.iter().enumerate() {
784            assert_eq!(
785                unpack_u16_transition(&packed, idx),
786                original,
787                "u16 round-trip diverged at flat index {idx}",
788            );
789        }
790    }
791
792    #[test]
793    fn u16_pack_fails_closed_on_target_exceeding_u16() {
794        // 70000 > u16::MAX: packing MUST refuse, never silently `& 0xFFFF` it to a
795        // wrong next-state (Law 10 (that truncation is an invisible recall loss)).
796        let compressed: Vec<u32> = vec![0, 1, 70_000, 3];
797        let mut out = Vec::new();
798        let err = try_pack_u16_transitions_into(&compressed, &mut out)
799            .expect_err("a target above u16::MAX must be refused");
800        let msg = err.to_string();
801        assert!(
802            msg.contains("70000") && msg.contains("index 2") && msg.contains("u16"),
803            "error must name the offending target/index and the u16 cause: {msg}",
804        );
805    }
806
807    /// Regression for P2 decoration test: the structural shared-storage checks
808    /// above were not sufficient, a refactor could share the WRONG compressed
809    /// block and still pass the field-equality assertion. This test packs TWO
810    /// identical copies of the non-trivial 3-class DFA from
811    /// `byte_class_compression_is_lossless` and then calls `packed_next_state`
812    /// on BOTH meta indices for every (state, byte) pair, asserting both return
813    /// the same value AND that value matches the dense source table.
814    #[test]
815    fn duplicate_dfas_shared_storage_both_rules_fire_correctly() {
816        // 3-state, 3-class DFA (same fixture as byte_class_compression_is_lossless).
817        let states = 3usize;
818        let mut dense = vec![0u32; states * 256];
819        dense[0 * 256 + 0x41] = 1; // state 0: 'A' -> 1
820        dense[1 * 256 + 0x41] = 2; // state 1: 'A' -> 2
821        dense[1 * 256 + 0x42] = 2; // state 1: 'B' -> 2
822        dense[2 * 256 + 0x41] = 2; // state 2: 'A' -> 2
823        let accept = vec![0u32, 0, 1];
824
825        let rule0 = BatchRuleProgram::new(0, dense.clone(), accept.clone(), states as u32).unwrap();
826        let rule1 = BatchRuleProgram::new(1, dense.clone(), accept.clone(), states as u32).unwrap();
827        let packed = pack_rule_catalog(&[rule0, rule1]).unwrap();
828
829        assert!(packed.rejected_rules.is_empty());
830        // Both rules must share storage.
831        assert_eq!(
832            packed.rule_meta[0].transition_base, packed.rule_meta[1].transition_base,
833            "Fix: duplicate DFAs must share transition storage"
834        );
835        assert_eq!(
836            packed.rule_meta[0].accept_base, packed.rule_meta[1].accept_base,
837            "Fix: duplicate DFAs must share accept storage"
838        );
839
840        // Critical: verify BOTH meta indices yield the correct DFA output for
841        // every (state, byte), structural field sharing is necessary but not
842        // sufficient; the shared block must actually encode the right DFA.
843        for state in 0..states as u32 {
844            for byte in 0u16..256 {
845                let byte = byte as u8;
846                let expected = dense[state as usize * 256 + byte as usize];
847                let got0 = packed_next_state(&packed, 0, state, byte);
848                let got1 = packed_next_state(&packed, 1, state, byte);
849                assert_eq!(
850                    got0, expected,
851                    "Fix: rule0 compressed transition mismatch at state {state} byte {byte:#x}: expected {expected} got {got0}"
852                );
853                assert_eq!(
854                    got1, expected,
855                    "Fix: rule1 compressed transition mismatch at state {state} byte {byte:#x}: expected {expected} got {got1}"
856                );
857            }
858        }
859    }
860
861    #[test]
862    fn duplicate_dfas_do_not_reserve_raw_duplicate_storage() {
863        let rules = (0..32)
864            .map(|rule_idx| BatchRuleProgram::new(rule_idx, vec![0; 256], vec![0], 1).unwrap())
865            .collect::<Vec<_>>();
866
867        let packed = pack_rule_catalog(&rules).unwrap();
868
869        // 1-word inert row + 1-word shared compressed row for all 32 duplicates.
870        assert_eq!(packed.transitions.len(), 2);
871        assert!(
872            packed.transitions.capacity() < ALPHABET_SIZE as usize * rules.len(),
873            "Fix: duplicate DFA catalogs must not reserve memory as if every rule had unique transition storage."
874        );
875        assert_eq!(packed.accept.len(), 2);
876        assert!(
877            packed.accept.capacity() < rules.len(),
878            "Fix: duplicate DFA catalogs must not reserve accept storage for every duplicate rule."
879        );
880        // One inert + one shared class map, not 32.
881        assert_eq!(packed.class_maps.len(), ALPHABET_SIZE as usize * 2);
882    }
883
884    /// The compressed catalog yields byte-for-byte identical next-states to the
885    /// dense `state * 256 + byte` table for EVERY (state, byte) of a non-trivial
886    /// multi-class DFA (the lossless parity contract the GPU kernel depends on).
887    #[test]
888    fn byte_class_compression_is_lossless() {
889        // 3-state DFA. byte 0x41 ('A') advances 0->1->2->2; byte 0x42 ('B')
890        // advances 1->2 only; all other bytes reset to 0. This forces THREE
891        // distinct byte classes (A, B, everything-else) so num_classes < 256
892        // and the compression is exercised, not a degenerate single class.
893        let states = 3usize;
894        let mut dense = vec![0u32; states * 256];
895        // state 0: 'A' -> 1, else -> 0
896        dense[0 * 256 + 0x41] = 1;
897        // state 1: 'A' -> 2, 'B' -> 2, else -> 0
898        dense[1 * 256 + 0x41] = 2;
899        dense[1 * 256 + 0x42] = 2;
900        // state 2: 'A' -> 2, else -> 0
901        dense[2 * 256 + 0x41] = 2;
902        let accept = vec![0u32, 0, 1];
903        let rule = BatchRuleProgram::new(0, dense.clone(), accept, states as u32).unwrap();
904        let packed = pack_rule_catalog(&[rule]).unwrap();
905
906        assert_eq!(packed.rejected_rules.len(), 0);
907        // 'A', 'B', and the rest are three behaviourally-distinct columns.
908        assert_eq!(packed.rule_meta[0].num_classes, 3);
909        assert!(
910            packed.transitions.len() < 1 + states * 256,
911            "compressed transitions must be smaller than the dense table"
912        );
913
914        for state in 0..states as u32 {
915            for byte in 0u16..256 {
916                let byte = byte as u8;
917                let expected = dense[state as usize * 256 + byte as usize];
918                let got = packed_next_state(&packed, 0, state, byte);
919                assert_eq!(
920                    got, expected,
921                    "compressed transition mismatch at state {state} byte {byte:#x}: dense={expected} packed={got}"
922                );
923            }
924        }
925    }
926
927    /// A DFA whose every byte transitions differently in some state must NOT be
928    /// over-compressed: it keeps all 256 classes and still round-trips losslessly.
929    #[test]
930    fn full_alphabet_dfa_keeps_all_classes_and_is_lossless() {
931        // 2-state DFA where state 0 sends byte b -> (b as state is impossible
932        // with 2 states), so instead: state 0 sends EVERY byte to a distinct
933        // value by using state 1 vs 0 based on parity, that only yields 2
934        // classes. To force 256 classes we need 256 distinct columns, which
935        // needs >=256 states. Use a 256-state identity: state s, byte b -> b.
936        let states = 256usize;
937        let mut dense = vec![0u32; states * 256];
938        for s in 0..states {
939            for b in 0..256 {
940                dense[s * 256 + b] = b as u32; // column for byte b is constant = b across all states
941            }
942        }
943        // Every byte's column is the constant vector [b; 256], all distinct, so
944        // 256 classes.
945        let accept = vec![0u32; states];
946        let rule = BatchRuleProgram::new(0, dense.clone(), accept, states as u32).unwrap();
947        let packed = pack_rule_catalog(&[rule]).unwrap();
948        assert_eq!(packed.rule_meta[0].num_classes, 256);
949        for state in 0..states as u32 {
950            for byte in 0u16..256 {
951                let byte = byte as u8;
952                let expected = dense[state as usize * 256 + byte as usize];
953                assert_eq!(packed_next_state(&packed, 0, state, byte), expected);
954            }
955        }
956    }
957
958    #[test]
959    fn accepted_rule_fingerprints_into_reuses_caller_storage() {
960        let rules = (0..8)
961            .map(|rule_idx| BatchRuleProgram::new(rule_idx, vec![0; 256], vec![0], 1).unwrap())
962            .collect::<Vec<_>>();
963        let mut fingerprints = Vec::with_capacity(16);
964        let mut occupied = Vec::with_capacity(16);
965        let mut addressed = Vec::with_capacity(16);
966        let fingerprint_ptr = fingerprints.as_ptr();
967        let occupied_ptr = occupied.as_ptr();
968        let addressed_ptr = addressed.as_ptr();
969
970        let rejections = accepted_rule_fingerprints_into(
971            &rules,
972            &mut fingerprints,
973            &mut occupied,
974            &mut addressed,
975        );
976
977        assert!(rejections.is_empty());
978        assert_eq!(fingerprints.len(), rules.len());
979        assert_eq!(fingerprints.as_ptr(), fingerprint_ptr);
980        assert_eq!(occupied.as_ptr(), occupied_ptr);
981        assert_eq!(addressed.as_ptr(), addressed_ptr);
982    }
983
984    #[test]
985    fn invalid_rules_are_isolated_to_inert_catalog_entries() {
986        let valid = BatchRuleProgram::new(0, vec![0; 256], vec![1], 1).unwrap();
987        let invalid = BatchRuleProgram {
988            rule_idx: 1,
989            transitions: vec![0; 8],
990            accept: vec![0],
991            state_count: 1,
992        };
993
994        let packed = pack_rule_catalog(&[valid, invalid]).unwrap();
995        assert_eq!(packed.rejected_rules.len(), 1);
996        assert_eq!(packed.rejected_rules[0].rule_idx, Some(1));
997        // Valid rule (slot 0) points at a REAL compressed block past the inert
998        // row; the inert/rejected slot 1 points back at the inert row 0.
999        assert_eq!(packed.rule_meta[0].state_count, 1);
1000        assert!(packed.rule_meta[0].transition_base >= 1);
1001        assert_eq!(packed.rule_meta[1].transition_base, 0);
1002        assert_eq!(packed.rule_meta[1].accept_base, 0);
1003        assert_eq!(packed.rule_meta[1].state_count, 1);
1004        assert_eq!(packed.rule_meta[1].class_map_base, 0);
1005        assert_eq!(packed.rule_meta[1].num_classes, 1);
1006        // Inert row 0: a single self-loop word and an all-zero 256-entry class
1007        // map (the rejected slot reads a well-formed no-match DFA).
1008        assert_eq!(packed.transitions[0], 0);
1009        assert_eq!(packed.accept[0], 0);
1010        assert_eq!(
1011            &packed.class_maps[..ALPHABET_SIZE as usize],
1012            &vec![0; ALPHABET_SIZE as usize]
1013        );
1014        // Regression for P2 decoration test: a single-byte spot check is not
1015        // sufficient, a corrupt inert row could have non-zero entries at other
1016        // bytes or at the accept table while still passing b'X'. This loop
1017        // proves the inert slot self-loops to state 0 on EVERY byte value and
1018        // that the accept entry for the inert slot is zero (can never match).
1019        for byte in 0u16..256 {
1020            let byte = byte as u8;
1021            assert_eq!(
1022                packed_next_state(&packed, 1, 0, byte),
1023                0,
1024                "Fix: inert slot must self-loop to state 0 for every byte, failed at byte {byte:#x}"
1025            );
1026        }
1027        // Accept entry for the inert slot at state 0 must be zero (no match).
1028        assert_eq!(
1029            packed.accept[packed.rule_meta[1].accept_base as usize],
1030            0,
1031            "Fix: inert slot accept entry at state 0 must be 0, the inert DFA must never produce a match"
1032        );
1033    }
1034}