vyre_libs/math/
succinct.rs

1//! Succinct bitvector metadata primitives.
2//!
3//! These ops build the rank side of rank/select navigation for compact token,
4//! AST, and graph bitvectors. They keep hot navigation state as packed `u32`
5//! words plus sparse superblock counters, so GPU kernels trade bandwidth-heavy
6//! pointer chasing for popcount math over coalesced words.
7
8use core::fmt;
9
10use crate::region::{wrap_anonymous, wrap_child};
11use vyre_foundation::ir::model::expr::GeneratorRef;
12use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
13
14const RANK_SUPERBLOCKS_OP_ID: &str = "vyre-libs::math::succinct::rank1_superblocks";
15const RANK_QUERY_OP_ID: &str = "vyre-libs::math::succinct::rank1_query";
16
17/// Build-time errors for succinct bitvector Programs.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub enum SuccinctBuildError {
20    /// Superblock size must be non-zero.
21    ZeroBlockWords,
22    /// The derived superblock output length overflowed `u32`.
23    SuperblockCountOverflow,
24}
25
26impl fmt::Display for SuccinctBuildError {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        match self {
29            Self::ZeroBlockWords => {
30                write!(f, "Fix: rank superblock size must be at least one u32 word")
31            }
32            Self::SuperblockCountOverflow => write!(
33                f,
34                "Fix: rank superblock count overflowed u32; shard the bitvector"
35            ),
36        }
37    }
38}
39
40impl std::error::Error for SuccinctBuildError {}
41
42fn superblock_count(word_count: u32, block_words: u32) -> Result<u32, SuccinctBuildError> {
43    if block_words == 0 {
44        return Err(SuccinctBuildError::ZeroBlockWords);
45    }
46    let full_blocks = word_count / block_words;
47    let has_partial = u32::from(word_count % block_words != 0);
48    full_blocks
49        .checked_add(has_partial)
50        .and_then(|blocks| blocks.checked_add(1))
51        .ok_or(SuccinctBuildError::SuperblockCountOverflow)
52}
53
54/// Build sparse rank1 superblocks for a packed u32 bitvector.
55///
56/// `superblocks[0]` is always zero. Each following entry stores the cumulative
57/// count of set bits before that superblock. The final sentinel stores the
58/// total popcount for the whole bitvector.
59#[must_use]
60pub fn rank1_superblocks(
61    bits: &str,
62    superblocks: &str,
63    word_count: u32,
64    block_words: u32,
65) -> Program {
66    try_rank1_superblocks(bits, superblocks, word_count, block_words).unwrap_or_else(|err| {
67        crate::builder::invalid_builder_trap_program(
68            RANK_SUPERBLOCKS_OP_ID,
69            superblocks,
70            DataType::U32,
71            format!("{err}"),
72        )
73    })
74}
75
76/// Checked builder for [`rank1_superblocks`].
77///
78/// # Errors
79///
80/// Returns [`SuccinctBuildError`] when `block_words` is zero or the derived
81/// metadata length overflows `u32`.
82pub fn try_rank1_superblocks(
83    bits: &str,
84    superblocks: &str,
85    word_count: u32,
86    block_words: u32,
87) -> Result<Program, SuccinctBuildError> {
88    let out_count = superblock_count(word_count, block_words)?;
89    let body = vec![Node::if_then(
90        Expr::eq(Expr::InvocationId { axis: 0 }, Expr::u32(0)),
91        vec![
92            Node::store(superblocks, Expr::u32(0), Expr::u32(0)),
93            Node::let_bind("rank_acc", Expr::u32(0)),
94            Node::loop_for(
95                "rank_word",
96                Expr::u32(0),
97                Expr::u32(word_count),
98                vec![
99                    Node::if_then(
100                        Expr::and(
101                            Expr::gt(Expr::var("rank_word"), Expr::u32(0)),
102                            Expr::eq(
103                                Expr::rem(Expr::var("rank_word"), Expr::u32(block_words)),
104                                Expr::u32(0),
105                            ),
106                        ),
107                        vec![Node::store(
108                            superblocks,
109                            Expr::div(Expr::var("rank_word"), Expr::u32(block_words)),
110                            Expr::var("rank_acc"),
111                        )],
112                    ),
113                    Node::assign(
114                        "rank_acc",
115                        Expr::add(
116                            Expr::var("rank_acc"),
117                            Expr::popcount(Expr::load(bits, Expr::var("rank_word"))),
118                        ),
119                    ),
120                ],
121            ),
122            Node::store(superblocks, Expr::u32(out_count - 1), Expr::var("rank_acc")),
123        ],
124    )];
125    Ok(Program::wrapped(
126        vec![
127            BufferDecl::storage(bits, 0, BufferAccess::ReadOnly, DataType::U32)
128                .with_count(word_count.max(1)),
129            BufferDecl::output(superblocks, 1, DataType::U32).with_count(out_count),
130        ],
131        [1, 1, 1],
132        vec![wrap_anonymous(
133            RANK_SUPERBLOCKS_OP_ID,
134            vec![wrap_child(
135                vyre_primitives::graph::path_reconstruct::OP_ID,
136                GeneratorRef {
137                    name: RANK_SUPERBLOCKS_OP_ID.to_string(),
138                },
139                body,
140            )],
141        )],
142    ))
143}
144
145/// Answer rank1-before-position queries from sparse superblocks.
146///
147/// Each `bit_indices[q]` is a zero-based bit offset. The output is the number
148/// of set bits strictly before that offset. Query offsets must address an
149/// existing packed word; use the final superblock sentinel for total popcount.
150#[must_use]
151pub fn rank1_query(
152    bits: &str,
153    superblocks: &str,
154    bit_indices: &str,
155    out: &str,
156    word_count: u32,
157    query_count: u32,
158    block_words: u32,
159) -> Program {
160    try_rank1_query(
161        bits,
162        superblocks,
163        bit_indices,
164        out,
165        word_count,
166        query_count,
167        block_words,
168    )
169    .unwrap_or_else(|err| {
170        crate::builder::invalid_builder_trap_program(
171            RANK_QUERY_OP_ID,
172            out,
173            DataType::U32,
174            format!("{err}"),
175        )
176    })
177}
178
179/// Checked builder for [`rank1_query`].
180///
181/// # Errors
182///
183/// Returns [`SuccinctBuildError`] when `block_words` is zero or the derived
184/// metadata length overflows `u32`.
185pub fn try_rank1_query(
186    bits: &str,
187    superblocks: &str,
188    bit_indices: &str,
189    out: &str,
190    word_count: u32,
191    query_count: u32,
192    block_words: u32,
193) -> Result<Program, SuccinctBuildError> {
194    let sb_count = superblock_count(word_count, block_words)?;
195    let q = Expr::InvocationId { axis: 0 };
196    let body = vec![Node::if_then(
197        Expr::lt(q.clone(), Expr::u32(query_count)),
198        vec![
199            Node::let_bind("bit_index", Expr::load(bit_indices, q.clone())),
200            Node::let_bind(
201                "word_index",
202                Expr::div(Expr::var("bit_index"), Expr::u32(32)),
203            ),
204            Node::if_then(
205                Expr::ge(Expr::var("word_index"), Expr::u32(word_count)),
206                vec![Node::trap(
207                    Expr::var("bit_index"),
208                    "rank-query-out-of-bounds",
209                )],
210            ),
211            Node::let_bind(
212                "block_index",
213                Expr::div(Expr::var("word_index"), Expr::u32(block_words)),
214            ),
215            Node::let_bind(
216                "rank_acc",
217                Expr::load(superblocks, Expr::var("block_index")),
218            ),
219            Node::let_bind(
220                "block_start_word",
221                Expr::mul(Expr::var("block_index"), Expr::u32(block_words)),
222            ),
223            Node::loop_for(
224                "rank_word",
225                Expr::var("block_start_word"),
226                Expr::var("word_index"),
227                vec![Node::assign(
228                    "rank_acc",
229                    Expr::add(
230                        Expr::var("rank_acc"),
231                        Expr::popcount(Expr::load(bits, Expr::var("rank_word"))),
232                    ),
233                )],
234            ),
235            Node::let_bind(
236                "bit_offset",
237                Expr::rem(Expr::var("bit_index"), Expr::u32(32)),
238            ),
239            Node::let_bind(
240                "partial_mask",
241                Expr::select(
242                    Expr::eq(Expr::var("bit_offset"), Expr::u32(0)),
243                    Expr::u32(0),
244                    Expr::sub(
245                        Expr::shl(Expr::u32(1), Expr::var("bit_offset")),
246                        Expr::u32(1),
247                    ),
248                ),
249            ),
250            Node::assign(
251                "rank_acc",
252                Expr::add(
253                    Expr::var("rank_acc"),
254                    Expr::popcount(Expr::bitand(
255                        Expr::load(bits, Expr::var("word_index")),
256                        Expr::var("partial_mask"),
257                    )),
258                ),
259            ),
260            Node::store(out, q, Expr::var("rank_acc")),
261        ],
262    )];
263    Ok(Program::wrapped(
264        vec![
265            BufferDecl::storage(bits, 0, BufferAccess::ReadOnly, DataType::U32)
266                .with_count(word_count.max(1)),
267            BufferDecl::storage(superblocks, 1, BufferAccess::ReadOnly, DataType::U32)
268                .with_count(sb_count),
269            BufferDecl::storage(bit_indices, 2, BufferAccess::ReadOnly, DataType::U32)
270                .with_count(query_count.max(1)),
271            BufferDecl::output(out, 3, DataType::U32).with_count(query_count.max(1)),
272        ],
273        [64, 1, 1],
274        vec![wrap_anonymous(RANK_QUERY_OP_ID, body)],
275    ))
276}
277
278inventory::submit! {
279    vyre_foundation::operation::OperationRegistration {
280        semantic_version: 1,
281        signature: None,
282        tier: vyre_foundation::operation::OperationTier::Library,
283        laws: &[],
284        tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
285        id: RANK_SUPERBLOCKS_OP_ID,
286        build: Some(|| rank1_superblocks("bits", "superblocks", 4, 2)),
287        test_inputs: Some(|| {
288            let bits = [0b1011u32, 0x8000_0000, 0xFFFF_0000, 0u32];
289            let to_bytes = vyre_primitives::wire::pack_u32_slice;
290            vec![vec![to_bytes(&bits)]]
291        }),
292        expected_output: Some(|| {
293            let expected = [0u32, 4, 20];
294            let bytes = vyre_primitives::wire::pack_u32_slice(&expected);
295            vec![vec![bytes]]
296        }),
297        category: Some("math"),
298    }
299}
300
301inventory::submit! {
302    vyre_foundation::operation::OperationRegistration {
303        semantic_version: 1,
304        signature: None,
305        tier: vyre_foundation::operation::OperationTier::Library,
306        laws: &[],
307        tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
308        id: RANK_QUERY_OP_ID,
309        build: Some(|| rank1_query("bits", "superblocks", "queries", "out", 4, 5, 2)),
310        test_inputs: Some(|| {
311            let bits = [0b1011u32, 0x8000_0000, 0xFFFF_0000, 0u32];
312            let superblocks = [0u32, 4, 20];
313            let queries = [0u32, 1, 4, 63, 80];
314            let to_bytes = vyre_primitives::wire::pack_u32_slice;
315            vec![vec![to_bytes(&bits), to_bytes(&superblocks), to_bytes(&queries)]]
316        }),
317        expected_output: Some(|| {
318            let expected = [0u32, 1, 3, 3, 4];
319            let bytes = vyre_primitives::wire::pack_u32_slice(&expected);
320            vec![vec![bytes]]
321        }),
322        category: Some("math"),
323    }
324}
vyre_libs/math/succinct.rs

vyre_libs/math/
succinct.rs
