Skip to main content

vyre_libs/scan/
builders.rs

1//! IR LEGO BLOCKS for matching dialects.
2//!
3//! Exposes granular primitives that can be composed into custom
4//! scanning engines (e.g. combined DFA + ML, decoder-aware scanners).
5
6use vyre_foundation::ir::{Expr, Node};
7
8/// LEGO BLOCK: Load a byte from a packed U32 haystack.
9///
10/// Returns `(let_bind_node, byte_expr)`. The caller must push the
11/// `let_bind_node` into its surrounding Block before evaluating
12/// `byte_expr` - the let-bind gives the optimiser a CSE handle for
13/// the underlying word load when the byte is referenced multiple
14/// times.
15pub fn load_packed_byte(haystack: &str, idx: Expr) -> (Node, Expr) {
16    let word_idx = Expr::div(idx.clone(), Expr::u32(4));
17    let byte_offset = Expr::mul(Expr::rem(idx, Expr::u32(4)), Expr::u32(8));
18
19    let node = Node::let_bind("_byte_word", Expr::load(haystack, word_idx));
20    let byte_expr = Expr::bitand(
21        Expr::shr(Expr::var("_byte_word"), byte_offset),
22        Expr::u32(0xFF),
23    );
24
25    (node, byte_expr)
26}
27
28/// LEGO BLOCK: Pure-expression form of `load_packed_byte`.
29///
30/// Inlines the word load directly into the shift+mask, returning a
31/// single `Expr` with no `Node::Let` side effect. Use this when the
32/// byte is consumed in a single expression context (e.g. an `If`
33/// condition or `Node::ne` predicate inside a `Loop` body) where
34/// hoisting a let-bind would either inject it at the wrong scope or
35/// require restructuring the surrounding IR.
36///
37/// Trade-off: no CSE handle for the word load - if the same byte
38/// position is referenced more than once, prefer `load_packed_byte`
39/// and bind the result. Single-use call sites (the cursor body in
40/// `nfa::nfa_scan_with_plan` and the literal-compare inside
41/// `literal_set::literal_set_program`) take this form because they
42/// reference the byte exactly once per iteration.
43pub fn load_packed_byte_expr(haystack: &str, idx: Expr) -> Expr {
44    Expr::bitand(
45        Expr::shr(
46            Expr::load(haystack, Expr::div(idx.clone(), Expr::u32(4))),
47            Expr::mul(Expr::rem(idx, Expr::u32(4)), Expr::u32(8)),
48        ),
49        Expr::u32(0xFF),
50    )
51}
52
53/// LEGO BLOCK: Append a match to a standardized hit buffer.
54///
55/// Use \`append_match_subgroup\` for production paths that benefit from
56/// subgroup-coalesced atomics (Innovation I.17).
57pub fn append_match(
58    hits_buffer: &str,
59    count_buffer: &str,
60    tag: impl Into<Expr>,
61    start: impl Into<Expr>,
62    end: impl Into<Expr>,
63) -> Node {
64    let max_hits = Expr::div(Expr::buf_len(hits_buffer), Expr::u32(3));
65
66    Node::Block(vec![
67        Node::let_bind(
68            "_vyre_match_slot",
69            Expr::atomic_add(count_buffer, Expr::u32(0), Expr::u32(1)),
70        ),
71        Node::if_then(
72            Expr::lt(Expr::var("_vyre_match_slot"), max_hits),
73            vec![
74                Node::store(
75                    hits_buffer,
76                    Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
77                    tag.into(),
78                ),
79                Node::store(
80                    hits_buffer,
81                    Expr::add(
82                        Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
83                        Expr::u32(1),
84                    ),
85                    start.into(),
86                ),
87                Node::store(
88                    hits_buffer,
89                    Expr::add(
90                        Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
91                        Expr::u32(2),
92                    ),
93                    end.into(),
94                ),
95            ],
96        ),
97    ])
98}
99
100/// Innovation I.17: Subgroup-Coalesced Match Append.
101///
102/// Uses subgroup-ballot and subgroup-shuffle to perform a single
103/// \`atomic_add\` per subgroup, drastically reducing global memory
104/// serialization on high-hit-rate workloads.
105pub fn append_match_subgroup(
106    hits_buffer: &str,
107    count_buffer: &str,
108    tag: impl Into<Expr>,
109    start: impl Into<Expr>,
110    end: impl Into<Expr>,
111    cond: Expr,
112) -> Vec<Node> {
113    let tag = tag.into();
114    let start = start.into();
115    let end = end.into();
116    let max_hits = Expr::div(Expr::buf_len(hits_buffer), Expr::u32(3));
117    let lane_mask = Expr::sub(
118        Expr::shl(Expr::u32(1), Expr::subgroup_local_id()),
119        Expr::u32(1),
120    );
121    let rank = Expr::popcount(Expr::bitand(Expr::var("_vyre_match_ballot"), lane_mask));
122    let leader_pred = Expr::and(
123        cond.clone(),
124        Expr::eq(Expr::var("_vyre_match_rank"), Expr::u32(0)),
125    );
126    let slot = Expr::add(
127        Expr::subgroup_shuffle(
128            Expr::var("_vyre_match_leader_base"),
129            Expr::var("_vyre_match_leader"),
130        ),
131        Expr::var("_vyre_match_rank"),
132    );
133    let ballot_cond = cond.clone();
134    let bounded_hit = Expr::and(cond, Expr::lt(Expr::var("_vyre_match_slot"), max_hits));
135
136    vec![
137        Node::let_bind("_vyre_match_ballot", Expr::subgroup_ballot(ballot_cond)),
138        Node::let_bind("_vyre_match_rank", rank),
139        Node::let_bind(
140            "_vyre_match_count",
141            Expr::popcount(Expr::var("_vyre_match_ballot")),
142        ),
143        Node::let_bind(
144            "_vyre_match_leader",
145            Expr::select(
146                Expr::eq(Expr::var("_vyre_match_count"), Expr::u32(0)),
147                Expr::u32(0),
148                Expr::ctz(Expr::var("_vyre_match_ballot")), // Fixed: relative to subgroup,
149            ),
150        ),
151        Node::let_bind("_vyre_match_leader_base", Expr::u32(0)),
152        Node::if_then(
153            leader_pred,
154            vec![Node::assign(
155                "_vyre_match_leader_base",
156                Expr::atomic_add(count_buffer, Expr::u32(0), Expr::var("_vyre_match_count")),
157            )],
158        ),
159        Node::let_bind("_vyre_match_slot", slot),
160        Node::if_then(
161            bounded_hit,
162            vec![
163                Node::store(
164                    hits_buffer,
165                    Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
166                    tag,
167                ),
168                Node::store(
169                    hits_buffer,
170                    Expr::add(
171                        Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
172                        Expr::u32(1),
173                    ),
174                    start,
175                ),
176                Node::store(
177                    hits_buffer,
178                    Expr::add(
179                        Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
180                        Expr::u32(2),
181                    ),
182                    end,
183                ),
184            ],
185        ),
186    ]
187}
188
189#[cfg(test)]
190mod subgroup_append_shape_tests {
191    use super::*;
192
193    #[test]
194    fn subgroup_append_bounds_hits_with_bound_slot_variable() {
195        let nodes = append_match_subgroup(
196            "matches",
197            "match_count",
198            Expr::u32(7),
199            Expr::u32(11),
200            Expr::u32(13),
201            Expr::var("hit"),
202        );
203
204        let bounded_cond = match nodes.last() {
205            Some(Node::If { cond, .. }) => format!("{cond:?}"),
206            other => panic!("expected bounded-hit If as final subgroup append node, got {other:?}"),
207        };
208
209        assert!(
210            bounded_cond.contains("_vyre_match_slot"),
211            "bounded-hit predicate must use the already-bound slot variable: {bounded_cond}"
212        );
213        assert!(
214            !bounded_cond.contains("_vyre_match_leader"),
215            "bounded-hit predicate must not re-inline subgroup shuffle leader expressions: {bounded_cond}"
216        );
217    }
218}
219
220#[cfg(test)]
221mod packed_byte_dedup_lock {
222    //! Regression gate for the canonical-packed-byte LEGO primitive.
223    //!
224    //! Six prior duplications of the `Expr::shr(Expr::load(buf,
225    //! word_idx), byte_offset) & 0xFF` byte-extract pattern landed in
226    //! vyre-libs over time (scan/nfa, scan/literal_set, parsing/c/
227    //! preprocess/gpu_if_expression/byte_load,
228    //! parsing/c/preprocess/gpu_filter/program_helpers). Tasks #21,
229    //! #22, #26 were marked completed previously while three of those
230    //! copies were still alive. This test prevents the next
231    //! regression: it walks `vyre-libs/src/**/*.rs` for the
232    //! divrem-shr-and(0xFF) shape and fails if it appears outside
233    //! `scan/builders.rs`.
234    //!
235    //! Detection is text-based and conservative - false positives are
236    //! tolerable (just add a `// allow-packed-byte-dup:` reason
237    //! comment on the line to suppress). False negatives (missing a
238    //! real duplicate) are the failure mode that matters; the shape
239    //! is distinctive enough that grep is sufficient.
240    use std::path::{Path, PathBuf};
241
242    fn vyre_libs_src() -> PathBuf {
243        let crate_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
244        crate_root.join("src")
245    }
246
247    fn walk(dir: &Path, out: &mut Vec<PathBuf>) {
248        let Ok(entries) = std::fs::read_dir(dir) else {
249            return;
250        };
251        for entry in entries.flatten() {
252            let path = entry.path();
253            if path.is_dir() {
254                walk(&path, out);
255            } else if path.extension().and_then(|e| e.to_str()) == Some("rs") {
256                out.push(path);
257            }
258        }
259    }
260}