Skip to main content

columnar/
sums.rs

1//! Containers for enumerations ("sum types") that store variants separately.
2//!
3//! The main work of these types is storing a discriminant and index efficiently,
4//! as containers for each of the variant types can hold the actual data.
5
6/// Stores for maintaining discriminants, and associated sequential indexes.
7///
8/// The sequential indexes are not explicitly maintained, but are supported
9/// by a `rank(index)` function that indicates how many of a certain variant
10/// precede the given index. While this could potentially be done with a scan
11/// of all preceding discriminants, the stores maintain running accumulations
12/// that make the operation constant time (using additional amortized memory).
13pub mod rank_select {
14
15    use alloc::{vec::Vec, string::String};
16    use crate::primitive::Bools;
17
18    use crate::{Borrow, Len, Index, IndexAs, Push, Clear};
19
20    /// Number of 64-bit words per cumulative-popcount chunk. Smaller values
21    /// reduce the worst-case word scan in `select` (and the catch-up scan
22    /// in `rank`) at the cost of more `counts` entries (~6% overhead per
23    /// halving). 16 → 1024 bits/chunk is the historical default; 8 → 512
24    /// bits gives faster random `select` for ~12% counts memory.
25    const WORDS_PER_CHUNK: usize = 16;
26    const BITS_PER_CHUNK: usize = 64 * WORDS_PER_CHUNK;
27
28    /// A store for maintaining `Vec<bool>` with fast `rank` and `select` access.
29    ///
30    /// The design is to have `u64` running counts for each block of 1024 bits,
31    /// which are roughly the size of a cache line. This is roughly 6% overhead,
32    /// above the bits themselves, which seems pretty solid.
33    #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
34    #[derive(Copy, Clone, Debug, Default, PartialEq)]
35    pub struct RankSelect<CC = Vec<u64>, VC = Vec<u64>, WC = [u64; 2]> {
36        /// Counts of the number of cumulative set (true) bits, *after* each block of 1024 bits.
37        pub counts: CC,
38        /// The bits themselves.
39        pub values: Bools<VC, WC>,
40    }
41
42    impl<CC: crate::common::BorrowIndexAs<u64>, VC: crate::common::BorrowIndexAs<u64>> RankSelect<CC, VC> {
43        #[inline(always)]
44        pub fn borrow<'a>(&'a self) -> RankSelect<CC::Borrowed<'a>, VC::Borrowed<'a>, &'a [u64]> {
45            RankSelect {
46                counts: self.counts.borrow(),
47                values: self.values.borrow(),
48            }
49        }
50        #[inline(always)]
51        pub fn reborrow<'b, 'a: 'b>(thing: RankSelect<CC::Borrowed<'a>, VC::Borrowed<'a>, &'a [u64]>) -> RankSelect<CC::Borrowed<'b>, VC::Borrowed<'b>, &'b [u64]> {
52            RankSelect {
53                counts: CC::reborrow(thing.counts),
54                values: Bools::<VC, [u64; 2]>::reborrow(thing.values),
55            }
56        }
57    }
58
59    impl<'a, CC: crate::AsBytes<'a>, VC: crate::AsBytes<'a>> crate::AsBytes<'a> for RankSelect<CC, VC, &'a [u64]> {
60        const SLICE_COUNT: usize = CC::SLICE_COUNT + <Bools<VC, &'a [u64]> as crate::AsBytes<'a>>::SLICE_COUNT;
61        #[inline]
62        fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
63            debug_assert!(index < Self::SLICE_COUNT);
64            if index < CC::SLICE_COUNT {
65                self.counts.get_byte_slice(index)
66            } else {
67                self.values.get_byte_slice(index - CC::SLICE_COUNT)
68            }
69        }
70    }
71    impl<'a, CC: crate::FromBytes<'a>, VC: crate::FromBytes<'a>> crate::FromBytes<'a> for RankSelect<CC, VC, &'a [u64]> {
72        const SLICE_COUNT: usize = CC::SLICE_COUNT + <crate::primitive::Bools<VC, &'a [u64]>>::SLICE_COUNT;
73        #[inline(always)]
74        fn from_bytes(bytes: &mut impl Iterator<Item=&'a [u8]>) -> Self {
75            Self {
76                counts: crate::FromBytes::from_bytes(bytes),
77                values: crate::FromBytes::from_bytes(bytes),
78            }
79        }
80        #[inline(always)]
81        fn from_store(store: &crate::bytes::indexed::DecodedStore<'a>, offset: &mut usize) -> Self {
82            Self {
83                counts: CC::from_store(store, offset),
84                values: <crate::primitive::Bools<VC, &'a [u64]>>::from_store(store, offset),
85            }
86        }
87        #[inline(always)]
88        fn element_sizes(sizes: &mut Vec<usize>) -> Result<(), String> {
89            CC::element_sizes(sizes)?;
90            <crate::primitive::Bools<VC, &'a [u64]>>::element_sizes(sizes)?;
91            Ok(())
92        }
93    }
94
95
96    impl<CC, VC: Len + IndexAs<u64>, WC: IndexAs<u64>> RankSelect<CC, VC, WC> {
97        #[inline(always)]
98        pub fn get(&self, index: usize) -> bool {
99            Index::get(&self.values, index)
100        }
101    }
102    impl<CC: Len + IndexAs<u64>, VC: Len + IndexAs<u64>, WC: IndexAs<u64>> RankSelect<CC, VC, WC> {
103        /// The number of set bits *strictly* preceding `index`.
104        ///
105        /// This number is accumulated first by reading out of `self.counts` at the correct position,
106        /// then by summing the ones in strictly prior `u64` entries, then by counting the ones in the
107        /// masked `u64` in which the bit lives.
108        pub fn rank(&self, index: usize) -> usize {
109            let bit = index % 64;
110            let block = index / 64;
111            let chunk = block / WORDS_PER_CHUNK;
112            let mut count = if chunk > 0 { self.counts.index_as(chunk - 1) as usize } else { 0 };
113            for pos in (WORDS_PER_CHUNK * chunk) .. block {
114                count += self.values.values.index_as(pos).count_ones() as usize;
115            }
116            // TODO: Panic if out of bounds?
117            let intra_word = if block == self.values.values.len() { self.values.tail.index_as(0) } else { self.values.values.index_as(block) };
118            count += (intra_word & ((1 << bit) - 1)).count_ones() as usize;
119            count
120        }
121        /// The position of the `rank`-th set bit (0-indexed), if it exists.
122        ///
123        /// `select(0)` returns the position of the first set bit. In general,
124        /// `select(rank(p)) == p` when `p` is the position of a set bit, mirroring
125        /// the convention of [`Self::rank`] (which counts bits strictly before
126        /// its argument).
127        #[inline]
128        pub fn select(&self, rank: u64) -> Option<usize> {
129            // Step one: find the BITS_PER_CHUNK-bit chunk containing the rank-th set bit.
130            // We want the smallest `chunk` for which `counts[chunk] > rank` — the
131            // chunk whose cumulative count first exceeds `rank`. Equivalent to a
132            // partition point over `counts` on the predicate `counts[i] <= rank`.
133            let chunk = {
134                let mut lo = 0;
135                let mut hi = self.counts.len();
136                while lo < hi {
137                    let mid = lo + (hi - lo) / 2;
138                    if self.counts.index_as(mid) <= rank { lo = mid + 1; } else { hi = mid; }
139                }
140                lo
141            };
142            // Number of set bits strictly before `chunk`'s BITS_PER_CHUNK bits.
143            let mut count = if chunk > 0 { self.counts.index_as(chunk - 1) } else { 0 };
144
145            // Step two: find the 64-bit word within `chunk` containing the rank-th set bit.
146            let mut block = WORDS_PER_CHUNK * chunk;
147            while block < self.values.values.len() {
148                let pop = self.values.values.index_as(block).count_ones() as u64;
149                if count + pop > rank { break; }
150                count += pop;
151                block += 1;
152            }
153
154            // Step three: locate the bit within the chosen word, or return `None`
155            // if `rank` is past the total set-bit count.
156            let (last_word, last_bits) = if block == self.values.values.len() {
157                (self.values.tail.index_as(0), self.values.tail.index_as(1) as usize)
158            } else {
159                (self.values.values.index_as(block), 64)
160            };
161            // Mask off any padding past the last valid bit (only relevant for the tail).
162            let masked = if last_bits == 64 { last_word } else { last_word & ((1u64 << last_bits) - 1) };
163            let k = (rank - count) as u32;
164            let shift = select_in_word(masked, k);
165            if shift >= last_bits as u32 { None } else { Some(64 * block + shift as usize) }
166        }
167    }
168
169    /// Position of the `k`-th set bit (0-indexed) within a 64-bit word. Returns 64 if
170    /// the word has fewer than `k + 1` set bits. Portable, ~6 conditional shifts based
171    /// on byte/half popcounts — independent of bit density, unlike a linear scan.
172    #[inline]
173    fn select_in_word(mut w: u64, mut k: u32) -> u32 {
174        if k >= w.count_ones() { return 64; }
175        let mut pos = 0u32;
176        // Halve the search range repeatedly. At each step, popcount of the low half
177        // tells us whether the k-th bit is in the low half (recurse) or the high half
178        // (subtract that count and shift).
179        let pop = (w & 0xFFFF_FFFF).count_ones();
180        if k >= pop { k -= pop; pos += 32; w >>= 32; }
181        let pop = (w & 0xFFFF).count_ones();
182        if k >= pop { k -= pop; pos += 16; w >>= 16; }
183        let pop = (w & 0xFF).count_ones();
184        if k >= pop { k -= pop; pos += 8; w >>= 8; }
185        let pop = (w & 0xF).count_ones();
186        if k >= pop { k -= pop; pos += 4; w >>= 4; }
187        let pop = (w & 0x3).count_ones();
188        if k >= pop { k -= pop; pos += 2; w >>= 2; }
189        let pop = w & 0x1;
190        if (k as u64) >= pop { pos += 1; }
191        pos
192    }
193
194    /// Forward cursor over a [`RankSelect`].
195    ///
196    /// Use this instead of repeated `rank`/`select` calls when you want to traverse
197    /// the bitvector in order. The cursor caches a current word and running rank,
198    /// so a single word load serves many subsequent operations — no re-probing
199    /// `counts` or rescanning words.
200    ///
201    /// At a high level the cursor maintains an *invariant pair* (`pos`, `rank`):
202    /// `pos` is the next bit to consider, and `rank` is the number of 1-bits in
203    /// `[0, pos)`. Every operation maintains this pair so that callers can read
204    /// either coordinate freely.
205    ///
206    /// Pick the method that matches the question you want to answer:
207    ///
208    /// | Operation         | Use when you want…                                        |
209    /// |-------------------|-----------------------------------------------------------|
210    /// | [`next_one`][n1]  | "Where is the next 1-bit?" — emits 1-bit positions in order. |
211    /// | [`step`][s]       | "What is the bit at the current position?" — read-by-bit traversal. |
212    /// | [`seek_to_pos`][stp] | "Jump to bit position p; what is its rank?" — random forward seek. |
213    /// | [`seek_to_rank`][str] | "Jump to the k-th 1-bit; where is it?" — equivalent to `select`. |
214    ///
215    /// The cursor is **forward-only**: every operation advances `pos` and `rank`
216    /// monotonically. Trying to seek backward triggers a debug-assertion failure.
217    ///
218    /// [n1]: Cursor::next_one
219    /// [s]: Cursor::step
220    /// [stp]: Cursor::seek_to_pos
221    /// [str]: Cursor::seek_to_rank
222    pub struct Cursor<'a, CC, VC, WC> {
223        rs: &'a RankSelect<CC, VC, WC>,
224        /// Index of the current 64-bit word within `values` (or `== values.len()` for tail).
225        word_idx: usize,
226        /// 1-bits remaining in the current word at positions ≥ `bit_pos % 64`.
227        /// Bits at positions strictly below `bit_pos % 64` are cleared.
228        word_remaining: u64,
229        /// Position of the next bit to consider.
230        bit_pos: usize,
231        /// Number of 1-bits in `[0, bit_pos)`.
232        rank: u64,
233        /// Cached total bit count of the underlying RankSelect.
234        total_bits: usize,
235    }
236
237    impl<CC: Len + IndexAs<u64>, VC: Len + IndexAs<u64>, WC: IndexAs<u64>> RankSelect<CC, VC, WC> {
238        /// Create a forward cursor positioned at bit 0.
239        pub fn cursor(&self) -> Cursor<'_, CC, VC, WC> {
240            let total_bits = self.len();
241            let mut c = Cursor {
242                rs: self,
243                word_idx: 0,
244                word_remaining: 0,
245                bit_pos: 0,
246                rank: 0,
247                total_bits,
248            };
249            c.load_word();
250            c
251        }
252    }
253
254    impl<'a, CC: Len + IndexAs<u64>, VC: Len + IndexAs<u64>, WC: IndexAs<u64>> Cursor<'a, CC, VC, WC> {
255        /// Position of the next bit to consider.
256        #[inline] pub fn pos(&self) -> usize { self.bit_pos }
257        /// Number of 1-bits strictly before `pos()`.
258        #[inline] pub fn rank(&self) -> u64 { self.rank }
259        /// Total number of bits in the underlying vector.
260        #[inline] pub fn total_bits(&self) -> usize { self.total_bits }
261
262        /// Load the word at `self.word_idx`, masked to its valid bits, and align
263        /// `bit_pos` to the start of that word. No-op past the end.
264        fn load_word(&mut self) {
265            let vlen = self.rs.values.values.len();
266            if self.word_idx < vlen {
267                self.word_remaining = self.rs.values.values.index_as(self.word_idx);
268                self.bit_pos = self.word_idx * 64;
269            } else if self.word_idx == vlen {
270                let raw = self.rs.values.tail.index_as(0);
271                let valid = self.rs.values.tail.index_as(1);
272                let mask = if valid >= 64 { !0u64 } else if valid == 0 { 0 } else { (1u64 << valid) - 1 };
273                self.word_remaining = raw & mask;
274                self.bit_pos = self.word_idx * 64;
275            } else {
276                self.word_remaining = 0;
277            }
278        }
279
280        /// Emit the next 1-bit position in order; advance past it.
281        ///
282        /// Equivalent to `select(self.rank())` followed by stepping one past that
283        /// position, but amortized: a single word load serves up to 64 results.
284        ///
285        /// **Use this for streaming through every set bit.** The canonical example
286        /// is recovering monotone Vecs bounds from an RS-encoded unary bitvector:
287        /// each call returns the next bound's bit position, no per-call re-probing
288        /// of `counts`. Returns `None` once all set bits have been emitted.
289        pub fn next_one(&mut self) -> Option<usize> {
290            loop {
291                if self.word_remaining != 0 {
292                    let bit_in_word = self.word_remaining.trailing_zeros() as usize;
293                    let pos = self.word_idx * 64 + bit_in_word;
294                    self.word_remaining &= self.word_remaining - 1;
295                    self.bit_pos = pos + 1;
296                    self.rank += 1;
297                    return Some(pos);
298                }
299                if self.bit_pos >= self.total_bits { return None; }
300                self.word_idx += 1;
301                self.load_word();
302                if self.word_idx > self.rs.values.values.len() { return None; }
303            }
304        }
305
306        /// Advance one bit. Return its value: `Some(true)` for 1, `Some(false)` for 0,
307        /// or `None` if past the end.
308        ///
309        /// **Use this when you need to visit every bit in order, regardless of value.**
310        /// Common pattern: the lookup-walk loop in a hash table that probes consecutive
311        /// slots, deciding what to do based on whether each slot is occupied. Each
312        /// call is a single bit-test + an occasional word-boundary crossing.
313        pub fn step(&mut self) -> Option<bool> {
314            if self.bit_pos >= self.total_bits { return None; }
315            let bit_in_word = self.bit_pos & 63;
316            let is_one = (self.word_remaining >> bit_in_word) & 1 == 1;
317            if is_one {
318                self.word_remaining &= !(1u64 << bit_in_word);
319                self.rank += 1;
320            }
321            self.bit_pos += 1;
322            if (self.bit_pos & 63) == 0 && self.bit_pos < self.total_bits {
323                self.word_idx += 1;
324                self.load_word();
325            }
326            Some(is_one)
327        }
328
329        /// Jump forward to bit position `target` and report its `rank` (via the
330        /// cursor's state) without consuming the bit there.
331        ///
332        /// After return, `pos() == target` and `rank() == rank(target)` (the count
333        /// of 1-bits strictly before `target`). A subsequent [`step`][Self::step]
334        /// reads the bit *at* `target`.
335        ///
336        /// **Use this when you have a known target position and want to read or
337        /// continue from there.** Typical use is hash-table lookup: a query's slot
338        /// position is computed from its hash; you want to jump there and then
339        /// walk forward through the probe chain.
340        ///
341        /// Fast path: if `target` lies within the cursor's current word, this is a
342        /// popcount + mask. Otherwise it pays one `RankSelect::rank` call to
343        /// re-anchor.
344        ///
345        /// `target` must be `>= self.pos()` (forward-only; debug-asserts otherwise).
346        /// Returns `false` if `target` is past the end of the bitvector.
347        pub fn seek_to_pos(&mut self, target: usize) -> bool {
348            debug_assert!(target >= self.bit_pos, "seek_to_pos is forward-only");
349            if target >= self.total_bits {
350                self.bit_pos = self.total_bits;
351                self.word_remaining = 0;
352                return false;
353            }
354            let target_word = target / 64;
355            if target_word == self.word_idx {
356                // Same word: count 1-bits we skip past, then mask the word.
357                let cur_bit = (self.bit_pos & 63) as u32;
358                let tgt_bit = (target & 63) as u32;
359                let skipped_mask = if tgt_bit == 64 { !0u64 } else { (1u64 << tgt_bit) - 1 };
360                let skipped = self.word_remaining & skipped_mask;
361                self.rank += skipped.count_ones() as u64;
362                // Clear consumed low bits (positions < target).
363                self.word_remaining &= !skipped_mask;
364                let _ = cur_bit;
365                self.bit_pos = target;
366                return true;
367            }
368            // Different word: re-anchor rank via the standard rank() formula, then
369            // load the target word and mask off bits below `target`.
370            self.rank = self.rs.rank(target) as u64;
371            self.word_idx = target_word;
372            self.load_word();
373            let tgt_bit = target & 63;
374            let mask = if tgt_bit == 0 { !0u64 } else { !((1u64 << tgt_bit) - 1) };
375            self.word_remaining &= mask;
376            self.bit_pos = target;
377            true
378        }
379
380        /// Jump forward to the `target`-th set bit (0-indexed) and return its
381        /// position. Equivalent to [`RankSelect::select`] for the cursor's
382        /// current state.
383        ///
384        /// After return, `rank() == target + 1` (the target bit has been consumed).
385        ///
386        /// **Use this when you have a target rank and want the position.** Typical
387        /// use is "give me bound[k]" against an RS-encoded Vecs-bounds bitvector,
388        /// possibly skipping forward over a stretch of consecutive bounds you
389        /// don't care about.
390        ///
391        /// Fast path: if the target's bit lies within the cursor's current word,
392        /// this is a `select_in_word` + bit-clear. Otherwise it pays a binary
393        /// search over `counts`, mirroring `RankSelect::select`.
394        ///
395        /// `target` must be `>= self.rank()` (forward-only; debug-asserts otherwise).
396        /// Returns `None` if there are fewer than `target + 1` set bits in total.
397        pub fn seek_to_rank(&mut self, target: u64) -> Option<usize> {
398            debug_assert!(target >= self.rank, "seek_to_rank is forward-only");
399            // Cheap path: target is within the current word's remaining 1-bits.
400            let here_pop = self.word_remaining.count_ones() as u64;
401            if target < self.rank + here_pop {
402                let k = (target - self.rank) as u32;
403                let bit_in_word = select_in_word(self.word_remaining, k) as usize;
404                let pos = self.word_idx * 64 + bit_in_word;
405                // Consume up to and including the target bit: clear the lowest k+1 set bits.
406                let mut w = self.word_remaining;
407                for _ in 0..=k { w &= w.wrapping_sub(1); }
408                self.word_remaining = w;
409                self.rank = target + 1;
410                self.bit_pos = pos + 1;
411                return Some(pos);
412            }
413            // Otherwise: jump via chunk binary search, mirroring `select`.
414            let counts = &self.rs.counts;
415            let chunk = {
416                let mut lo = 0;
417                let mut hi = counts.len();
418                while lo < hi {
419                    let mid = lo + (hi - lo) / 2;
420                    if counts.index_as(mid) <= target { lo = mid + 1; } else { hi = mid; }
421                }
422                lo
423            };
424            let mut count = if chunk > 0 { counts.index_as(chunk - 1) } else { 0 };
425            let vlen = self.rs.values.values.len();
426            let mut block = WORDS_PER_CHUNK * chunk;
427            while block < vlen {
428                let pop = self.rs.values.values.index_as(block).count_ones() as u64;
429                if count + pop > target { break; }
430                count += pop;
431                block += 1;
432            }
433            self.word_idx = block;
434            self.load_word();
435            self.rank = count;
436            // Now finish within the freshly loaded word.
437            if target >= count + self.word_remaining.count_ones() as u64 {
438                // Exhausted.
439                self.bit_pos = self.total_bits;
440                self.word_remaining = 0;
441                return None;
442            }
443            let k = (target - count) as u32;
444            let bit_in_word = select_in_word(self.word_remaining, k) as usize;
445            let pos = self.word_idx * 64 + bit_in_word;
446            let mut w = self.word_remaining;
447            for _ in 0..=k { w &= w.wrapping_sub(1); }
448            self.word_remaining = w;
449            self.rank = target + 1;
450            self.bit_pos = pos + 1;
451            Some(pos)
452        }
453    }
454
455    impl<CC, VC: Len, WC: IndexAs<u64>> RankSelect<CC, VC, WC> {
456        pub fn len(&self) -> usize {
457            self.values.len()
458        }
459    }
460
461    // This implementation probably only works for `Vec<u64>` and `Vec<u64>`, but we could fix that.
462    // Partly, it's hard to name the `Index` flavor that allows one to get back a `u64`.
463    impl<CC: for<'a> Push<&'a u64> + Len + IndexAs<u64>, VC: for<'a> Push<&'a u64> + Len + IndexAs<u64>> RankSelect<CC, VC> {
464        #[inline]
465        pub fn push(&mut self, bit: bool) {
466            self.values.push(&bit);
467            self.update_counts();
468        }
469        /// Appends the bits of `other` in `range`, a word at a time, maintaining the cumulative counts.
470        pub fn extend_from_self<CC2, VC2: Len + IndexAs<u64>, WC2: IndexAs<u64>>(&mut self, other: &RankSelect<CC2, VC2, WC2>, range: core::ops::Range<usize>) {
471            self.values.extend_bits_from(&other.values, range);
472            self.update_counts();
473        }
474        /// Extends `counts` to cover each newly completed chunk of `values`.
475        #[inline]
476        fn update_counts(&mut self) {
477            while self.counts.len() < self.values.len() / BITS_PER_CHUNK {
478                let mut count = self.counts.last().unwrap_or(0);
479                let lower = WORDS_PER_CHUNK * self.counts.len();
480                let upper = lower + WORDS_PER_CHUNK;
481                for i in lower .. upper {
482                    count += self.values.values.index_as(i).count_ones() as u64;
483                }
484                self.counts.push(&count);
485            }
486        }
487    }
488    impl RankSelect {
489        /// Reserves capacity for the counts and words of `selves`.
490        pub fn reserve_for<'a, I>(&mut self, selves: I) where I: Iterator<Item = RankSelect<&'a [u64], &'a [u64], &'a [u64]>> + Clone {
491            self.counts.reserve(selves.clone().map(|x| x.counts.len()).sum::<usize>());
492            self.values.values.reserve(selves.map(|x| x.values.values.len()).sum::<usize>());
493        }
494    }
495    impl<CC: Clear, VC: Clear> Clear for RankSelect<CC, VC> {
496        fn clear(&mut self) {
497            self.counts.clear();
498            self.values.clear();
499        }
500    }
501
502    #[cfg(test)]
503    mod tests {
504        use alloc::{vec, vec::Vec};
505        use super::RankSelect;
506
507        fn build(bits: &[bool]) -> RankSelect {
508            let mut rs: RankSelect = RankSelect::default();
509            for &b in bits { rs.push(b); }
510            rs
511        }
512
513        /// All true bits are recovered by `select(rank(p))` for every set position `p`,
514        /// and `rank` agrees with a naive count.
515        fn check_round_trip(bits: &[bool]) {
516            let rs = build(bits);
517            let mut expected = 0u64;
518            for (i, &b) in bits.iter().enumerate() {
519                assert_eq!(rs.rank(i), expected as usize, "rank({}) on pattern of len {}", i, bits.len());
520                if b {
521                    let pos = rs.select(expected).unwrap_or_else(|| panic!("select({}) returned None for set bit at {}", expected, i));
522                    assert_eq!(pos, i, "select({}) on pattern of len {}", expected, bits.len());
523                    expected += 1;
524                }
525            }
526            // Out-of-range select returns None.
527            assert!(rs.select(expected).is_none());
528        }
529
530        #[test]
531        fn select_first_bit() {
532            // Bit 0 set, nothing else.
533            let mut bits = vec![false; 2048];
534            bits[0] = true;
535            check_round_trip(&bits);
536        }
537
538        #[test]
539        fn select_small_dense() {
540            // First five bits set in a 2048-bit vector (spans two chunks).
541            let mut bits = vec![false; 2048];
542            for i in 0..5 { bits[i] = true; }
543            check_round_trip(&bits);
544        }
545
546        #[test]
547        fn select_chunk_boundary() {
548            // Set bits exactly at the chunk boundary positions 1023 and 1024.
549            let mut bits = vec![false; 4096];
550            bits[1023] = true;
551            bits[1024] = true;
552            check_round_trip(&bits);
553        }
554
555        #[test]
556        fn select_in_tail() {
557            // Bits 0..3 set, then nothing through bit 1100. Pattern length 1100 puts
558            // the final bits in the tail (not a complete 1024-bit chunk).
559            let mut bits = vec![false; 1100];
560            bits[0] = true;
561            bits[1099] = true;
562            check_round_trip(&bits);
563        }
564
565        #[test]
566        fn select_every_other() {
567            // Dense, multiple chunks.
568            let bits: Vec<bool> = (0..3000).map(|i| i % 2 == 0).collect();
569            check_round_trip(&bits);
570        }
571
572        #[test]
573        fn select_sparse_multi_chunk() {
574            // One set bit per 1024-bit chunk, six chunks.
575            let mut bits = vec![false; 6 * 1024];
576            for c in 0..6 { bits[1024 * c + 17] = true; }
577            check_round_trip(&bits);
578        }
579
580        #[test]
581        fn select_out_of_range() {
582            let rs = build(&[true, false, true]);
583            assert_eq!(rs.select(0), Some(0));
584            assert_eq!(rs.select(1), Some(2));
585            assert_eq!(rs.select(2), None);
586            assert_eq!(rs.select(1000), None);
587        }
588
589        #[test]
590        fn cursor_next_one_matches_select() {
591            // For each set bit in a 3000-bit pattern, walking next_one gives the
592            // same positions as repeated select calls.
593            let bits: Vec<bool> = (0..3000).map(|i| i % 7 == 0).collect();
594            let rs = build(&bits);
595            let mut cur = rs.cursor();
596            let mut k = 0u64;
597            while let Some(pos) = cur.next_one() {
598                assert_eq!(Some(pos), rs.select(k));
599                assert_eq!(cur.rank(), k + 1);
600                assert_eq!(cur.pos(), pos + 1);
601                k += 1;
602            }
603            assert_eq!(k, rs.rank(rs.len()) as u64);
604        }
605
606        #[test]
607        fn cursor_step_walks_every_bit() {
608            let bits: Vec<bool> = (0..2050).map(|i| i % 3 == 0).collect();
609            let rs = build(&bits);
610            let mut cur = rs.cursor();
611            let mut expected_rank = 0u64;
612            for (i, &b) in bits.iter().enumerate() {
613                assert_eq!(cur.pos(), i);
614                assert_eq!(cur.rank(), expected_rank);
615                assert_eq!(cur.step(), Some(b));
616                if b { expected_rank += 1; }
617            }
618            assert_eq!(cur.step(), None);
619        }
620
621        #[test]
622        fn cursor_seek_to_rank_skips_far_forward() {
623            // 3 chunks worth of bits, with set bits clustered. Seek directly to
624            // the 100-th set bit from a fresh cursor; result matches select(100).
625            let bits: Vec<bool> = (0..3200).map(|i| i % 3 == 1).collect();
626            let rs = build(&bits);
627            let mut cur = rs.cursor();
628            let pos = cur.seek_to_rank(100).unwrap();
629            assert_eq!(Some(pos), rs.select(100));
630            assert_eq!(cur.rank(), 101);
631            // A subsequent next_one continues correctly.
632            let next = cur.next_one().unwrap();
633            assert_eq!(Some(next), rs.select(101));
634        }
635
636        #[test]
637        fn cursor_seek_to_rank_within_current_word() {
638            // Set bits in a known pattern within the first word; cursor should
639            // satisfy seek_to_rank from the current-word fast path without re-probing.
640            let bits: Vec<bool> = (0..64).map(|i| matches!(i, 1 | 5 | 13 | 30 | 50)).collect();
641            let rs = build(&bits);
642            let mut cur = rs.cursor();
643            assert_eq!(cur.seek_to_rank(0), Some(1));
644            assert_eq!(cur.seek_to_rank(2), Some(13));
645            assert_eq!(cur.seek_to_rank(4), Some(50));
646            assert_eq!(cur.next_one(), None);
647        }
648
649        #[test]
650        fn cursor_seek_to_pos_same_word_and_cross_word() {
651            let bits: Vec<bool> = (0..256).map(|i| matches!(i % 5, 0 | 2)).collect();
652            let rs = build(&bits);
653            // Same-word seek.
654            let mut cur = rs.cursor();
655            assert!(cur.seek_to_pos(10));
656            assert_eq!(cur.pos(), 10);
657            assert_eq!(cur.rank(), rs.rank(10) as u64);
658            assert!(cur.seek_to_pos(40));
659            assert_eq!(cur.rank(), rs.rank(40) as u64);
660            // Cross-word seek.
661            assert!(cur.seek_to_pos(200));
662            assert_eq!(cur.pos(), 200);
663            assert_eq!(cur.rank(), rs.rank(200) as u64);
664            // Step reads the bit at the target.
665            assert_eq!(cur.step(), Some(bits[200]));
666        }
667
668        #[test]
669        fn cursor_seek_to_pos_out_of_range() {
670            let bits: Vec<bool> = (0..100).map(|_| true).collect();
671            let rs = build(&bits);
672            let mut cur = rs.cursor();
673            assert!(!cur.seek_to_pos(1000));
674            assert_eq!(cur.step(), None);
675        }
676
677        #[test]
678        fn cursor_seek_to_rank_out_of_range() {
679            let bits: Vec<bool> = (0..200).map(|i| i % 10 == 0).collect();
680            let rs = build(&bits);
681            let mut cur = rs.cursor();
682            assert_eq!(cur.seek_to_rank(10_000), None);
683        }
684
685        /// `extend_from_self` agrees with pushing one bit at a time, and the ranks and selects of
686        /// the result agree with a naive count, for random sources, targets and ranges.
687        #[test]
688        fn extend_from_self_matches_push() {
689            use crate::sums::test_util::Lcg;
690            let mut rng = Lcg(7);
691            for _ in 0..200 {
692                let source_bits: Vec<bool> = (0..(rng.next() % 5000)).map(|_| rng.next() % 4 == 0).collect();
693                let source = build(&source_bits);
694                let mut expected: Vec<bool> = (0..(rng.next() % 3000)).map(|_| rng.next() % 4 == 0).collect();
695                let mut target = build(&expected);
696                let lower = (rng.next() as usize) % (source_bits.len() + 1);
697                let upper = lower + (rng.next() as usize) % (source_bits.len() - lower + 1);
698                target.extend_from_self(&source, lower .. upper);
699                expected.extend_from_slice(&source_bits[lower .. upper]);
700                assert_eq!(target.len(), expected.len());
701                assert_eq!(target.counts.len(), expected.len() / super::BITS_PER_CHUNK);
702                let mut ones = 0u64;
703                for (i, &b) in expected.iter().enumerate() {
704                    assert_eq!(target.rank(i), ones as usize, "rank({i})");
705                    assert_eq!(target.get(i), b, "get({i})");
706                    if b {
707                        assert_eq!(target.select(ones), Some(i), "select({ones})");
708                        ones += 1;
709                    }
710                }
711                assert_eq!(target.rank(expected.len()), ones as usize);
712                assert_eq!(target.select(ones), None);
713            }
714        }
715
716        #[test]
717        fn select_in_word_basic() {
718            use super::select_in_word;
719            // 0b10110: set bits at positions 1, 2, 4.
720            assert_eq!(select_in_word(0b10110, 0), 1);
721            assert_eq!(select_in_word(0b10110, 1), 2);
722            assert_eq!(select_in_word(0b10110, 2), 4);
723            assert_eq!(select_in_word(0b10110, 3), 64);
724            // Edges of the word.
725            assert_eq!(select_in_word(1u64 << 63, 0), 63);
726            assert_eq!(select_in_word(u64::MAX, 0), 0);
727            assert_eq!(select_in_word(u64::MAX, 63), 63);
728            assert_eq!(select_in_word(u64::MAX, 64), 64);
729            assert_eq!(select_in_word(0, 0), 64);
730        }
731    }
732}
733
734pub mod result {
735
736    use alloc::{vec::Vec, string::String};
737
738    use crate::{Clear, Columnar, Container, Len, IndexMut, Index, IndexAs, Push, Borrow};
739    use crate::RankSelect;
740
741    #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
742    #[derive(Copy, Clone, Debug, Default, PartialEq)]
743    pub struct Results<SC, TC, CC=Vec<u64>, VC=Vec<u64>, WC=[u64; 2]> {
744        /// Bits set to `true` correspond to `Ok` variants.
745        pub indexes: RankSelect<CC, VC, WC>,
746        pub oks: SC,
747        pub errs: TC,
748    }
749
750    impl<S: Columnar, T: Columnar> Columnar for Result<S, T> {
751        fn copy_from<'a>(&mut self, other: crate::Ref<'a, Self>) {
752            match (&mut *self, other) {
753                (Ok(x), Ok(y)) => x.copy_from(y),
754                (Err(x), Err(y)) => x.copy_from(y),
755                (_, other) => { *self = Self::into_owned(other); },
756            }
757        }
758        fn into_owned<'a>(other: crate::Ref<'a, Self>) -> Self {
759            match other {
760                Ok(y) => Ok(S::into_owned(y)),
761                Err(y) => Err(T::into_owned(y)),
762            }
763        }
764        type Container = Results<S::Container, T::Container>;
765    }
766
767    impl<SC: Borrow, TC: Borrow> Borrow for Results<SC, TC> {
768        type Ref<'a> = Result<SC::Ref<'a>, TC::Ref<'a>> where SC: 'a, TC: 'a;
769        type Borrowed<'a> = Results<SC::Borrowed<'a>, TC::Borrowed<'a>, &'a [u64], &'a [u64], &'a [u64]> where SC: 'a, TC: 'a;
770        fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
771            Results {
772                indexes: self.indexes.borrow(),
773                oks: self.oks.borrow(),
774                errs: self.errs.borrow(),
775            }
776        }
777        #[inline(always)]
778        fn reborrow<'b, 'a: 'b>(thing: Self::Borrowed<'a>) -> Self::Borrowed<'b> where SC: 'a, TC: 'a {
779            Results {
780                indexes: RankSelect::<Vec<u64>, Vec<u64>>::reborrow(thing.indexes),
781                oks: SC::reborrow(thing.oks),
782                errs: TC::reborrow(thing.errs),
783            }
784        }
785        #[inline(always)]
786        fn reborrow_ref<'b, 'a: 'b>(thing: Self::Ref<'a>) -> Self::Ref<'b> where Self: 'a {
787            match thing {
788                Ok(y) => Ok(SC::reborrow_ref(y)),
789                Err(y) => Err(TC::reborrow_ref(y)),
790            }
791        }
792    }
793
794    impl<SC: Container, TC: Container> Container for Results<SC, TC> {
795        #[inline(always)]
796        fn extend_from_self(&mut self, other: Self::Borrowed<'_>, range: core::ops::Range<usize>) {
797            if !range.is_empty() {
798                // The `Ok` and `Err` variants in `range` each occupy a contiguous range of their container,
799                // located by the number of `Ok` variants before the start and end of `range`.
800                let oks_start = other.indexes.rank(range.start);
801                let oks_end = other.indexes.rank(range.end);
802                let errs_start = range.start - oks_start;
803                let errs_end = range.end - oks_end;
804
805                self.indexes.extend_from_self(&other.indexes, range);
806                self.oks.extend_from_self(other.oks, oks_start .. oks_end);
807                self.errs.extend_from_self(other.errs, errs_start .. errs_end);
808            }
809        }
810
811        fn reserve_for<'a, I>(&mut self, selves: I) where Self: 'a, I: Iterator<Item = Self::Borrowed<'a>> + Clone {
812            self.indexes.reserve_for(selves.clone().map(|x| x.indexes));
813            self.oks.reserve_for(selves.clone().map(|x| x.oks));
814            self.errs.reserve_for(selves.map(|x| x.errs));
815        }
816    }
817
818    impl<'a, SC: crate::AsBytes<'a>, TC: crate::AsBytes<'a>, CC: crate::AsBytes<'a>, VC: crate::AsBytes<'a>> crate::AsBytes<'a> for Results<SC, TC, CC, VC, &'a [u64]> {
819        const SLICE_COUNT: usize = <RankSelect<CC, VC, &'a [u64]> as crate::AsBytes<'a>>::SLICE_COUNT + SC::SLICE_COUNT + TC::SLICE_COUNT;
820        #[inline]
821        fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
822            debug_assert!(index < Self::SLICE_COUNT);
823            let idx_count = <RankSelect<CC, VC, &'a [u64]> as crate::AsBytes<'a>>::SLICE_COUNT;
824            if index < idx_count {
825                self.indexes.get_byte_slice(index)
826            } else if index < idx_count + SC::SLICE_COUNT {
827                self.oks.get_byte_slice(index - idx_count)
828            } else {
829                self.errs.get_byte_slice(index - idx_count - SC::SLICE_COUNT)
830            }
831        }
832    }
833    impl<'a, SC: crate::FromBytes<'a>, TC: crate::FromBytes<'a>, CC: crate::FromBytes<'a>, VC: crate::FromBytes<'a>> crate::FromBytes<'a> for Results<SC, TC, CC, VC, &'a [u64]> {
834        const SLICE_COUNT: usize = <RankSelect<CC, VC, &'a [u64]>>::SLICE_COUNT + SC::SLICE_COUNT + TC::SLICE_COUNT;
835        #[inline(always)]
836        fn from_bytes(bytes: &mut impl Iterator<Item=&'a [u8]>) -> Self {
837            Self {
838                indexes: crate::FromBytes::from_bytes(bytes),
839                oks: crate::FromBytes::from_bytes(bytes),
840                errs: crate::FromBytes::from_bytes(bytes),
841            }
842        }
843        #[inline(always)]
844        fn from_store(store: &crate::bytes::indexed::DecodedStore<'a>, offset: &mut usize) -> Self {
845            Self {
846                indexes: crate::FromBytes::from_store(store, offset),
847                oks: SC::from_store(store, offset),
848                errs: TC::from_store(store, offset),
849            }
850        }
851        #[inline(always)]
852        fn element_sizes(sizes: &mut Vec<usize>) -> Result<(), String> {
853            <RankSelect<CC, VC, &'a [u64]>>::element_sizes(sizes)?;
854            SC::element_sizes(sizes)?;
855            TC::element_sizes(sizes)?;
856            Ok(())
857        }
858    }
859
860    impl<SC, TC, CC, VC: Len, WC: IndexAs<u64>> Len for Results<SC, TC, CC, VC, WC> {
861        #[inline(always)] fn len(&self) -> usize { self.indexes.len() }
862    }
863
864    impl<SC, TC, CC, VC, WC> Index for Results<SC, TC, CC, VC, WC>
865    where
866        SC: Index,
867        TC: Index,
868        CC: IndexAs<u64> + Len,
869        VC: IndexAs<u64> + Len,
870        WC: IndexAs<u64>,
871    {
872        type Ref = Result<SC::Ref, TC::Ref>;
873        #[inline(always)]
874        fn get(&self, index: usize) -> Self::Ref {
875            if self.indexes.get(index) {
876                Ok(self.oks.get(self.indexes.rank(index)))
877            } else {
878                Err(self.errs.get(index - self.indexes.rank(index)))
879            }
880        }
881    }
882    impl<'a, SC, TC, CC, VC, WC> Index for &'a Results<SC, TC, CC, VC, WC>
883    where
884        &'a SC: Index,
885        &'a TC: Index,
886        CC: IndexAs<u64> + Len,
887        VC: IndexAs<u64> + Len,
888        WC: IndexAs<u64>,
889    {
890        type Ref = Result<<&'a SC as Index>::Ref, <&'a TC as Index>::Ref>;
891        #[inline(always)]
892        fn get(&self, index: usize) -> Self::Ref {
893            if self.indexes.get(index) {
894                Ok((&self.oks).get(self.indexes.rank(index)))
895            } else {
896                Err((&self.errs).get(index - self.indexes.rank(index)))
897            }
898        }
899    }
900
901    // NB: You are not allowed to change the variant, but can change its contents.
902    impl<SC: IndexMut, TC: IndexMut, CC: IndexAs<u64> + Len, VC: IndexAs<u64> + Len> IndexMut for Results<SC, TC, CC, VC> {
903        type IndexMut<'a> = Result<SC::IndexMut<'a>, TC::IndexMut<'a>> where SC: 'a, TC: 'a, CC: 'a, VC: 'a;
904        #[inline(always)]
905        fn get_mut(&mut self, index: usize) -> Self::IndexMut<'_> {
906            if self.indexes.get(index) {
907                Ok(self.oks.get_mut(self.indexes.rank(index)))
908            } else {
909                Err(self.errs.get_mut(index - self.indexes.rank(index)))
910            }
911        }
912    }
913
914    impl<S, SC: Push<S>, T, TC: Push<T>> Push<Result<S, T>> for Results<SC, TC> {
915        #[inline]
916        fn push(&mut self, item: Result<S, T>) {
917            match item {
918                Ok(item) => {
919                    self.indexes.push(true);
920                    self.oks.push(item);
921                }
922                Err(item) => {
923                    self.indexes.push(false);
924                    self.errs.push(item);
925                }
926            }
927        }
928    }
929    impl<'a, S, SC: Push<&'a S>, T, TC: Push<&'a T>> Push<&'a Result<S, T>> for Results<SC, TC> {
930        #[inline]
931        fn push(&mut self, item: &'a Result<S, T>) {
932            match item {
933                Ok(item) => {
934                    self.indexes.push(true);
935                    self.oks.push(item);
936                }
937                Err(item) => {
938                    self.indexes.push(false);
939                    self.errs.push(item);
940                }
941            }
942        }
943    }
944
945    impl<SC: Clear, TC: Clear> Clear for Results<SC, TC> {
946        fn clear(&mut self) {
947            self.indexes.clear();
948            self.oks.clear();
949            self.errs.clear();
950        }
951    }
952
953    impl<SC, TC, CC, VC, WC> Results<SC, TC, CC, VC, WC> {
954        /// Returns ok values if no errors exist.
955        pub fn unwrap(self) -> SC where TC: Len {
956            assert!(self.errs.is_empty());
957            self.oks
958        }
959        /// Returns error values if no oks exist.
960        pub fn unwrap_err(self) -> TC where SC: Len {
961            assert!(self.oks.is_empty());
962            self.errs
963        }
964        /// Returns ok values if no errors exist, or `None`.
965        pub fn try_unwrap(self) -> Option<SC> where TC: Len {
966            if self.errs.is_empty() { Some(self.oks) } else { None }
967        }
968        /// Returns error values if no oks exist, or `None`.
969        pub fn try_unwrap_err(self) -> Option<TC> where SC: Len {
970            if self.oks.is_empty() { Some(self.errs) } else { None }
971        }
972    }
973    #[cfg(test)]
974    mod test {
975        #[test]
976        fn round_trip() {
977
978            use crate::common::{Index, Push, Len};
979
980            let mut column: crate::ContainerOf<Result<u64, u64>> = Default::default();
981            for i in 0..100 {
982                column.push(Ok::<u64, u64>(i));
983                column.push(Err::<u64, u64>(i));
984            }
985
986            assert_eq!(column.len(), 200);
987
988            for i in 0..100 {
989                assert_eq!(column.get(2*i+0), Ok(i as u64));
990                assert_eq!(column.get(2*i+1), Err(i as u64));
991            }
992
993            let mut column: crate::ContainerOf<Result<u64, u8>> = Default::default();
994            for i in 0..100 {
995                column.push(Ok::<u64, u8>(i as u64));
996                column.push(Err::<u64, u8>(i as u8));
997            }
998
999            assert_eq!(column.len(), 200);
1000
1001            for i in 0..100 {
1002                assert_eq!(column.get(2*i+0), Ok(i as u64));
1003                assert_eq!(column.get(2*i+1), Err(i as u8));
1004            }
1005        }
1006
1007        /// `extend_from_self` agrees with pushing one element at a time, for random
1008        /// sources, targets and ranges.
1009        #[test]
1010        fn extend_from_self_random() {
1011            use alloc::vec::Vec;
1012            use crate::common::{Index, Len};
1013            use crate::{Borrow, Columnar, Container};
1014            use crate::sums::test_util::Lcg;
1015            let mut rng = Lcg(5);
1016            for _ in 0..200 {
1017                let source: Vec<Result<u32, u8>> = (0..(rng.next() % 3000)).map(|i| if rng.next() % 3 == 0 { Ok(i as u32) } else { Err(i as u8) }).collect();
1018                let mut expected: Vec<Result<u32, u8>> = (0..(rng.next() % 1500)).map(|i| if rng.next() % 2 == 0 { Ok(100_000 + i as u32) } else { Err(!(i as u8)) }).collect();
1019                let source_column: crate::ContainerOf<Result<u32, u8>> = Columnar::into_columns(source.iter().copied());
1020                let mut target: crate::ContainerOf<Result<u32, u8>> = Columnar::into_columns(expected.iter().copied());
1021                let lower = (rng.next() as usize) % (source.len() + 1);
1022                let upper = lower + (rng.next() as usize) % (source.len() - lower + 1);
1023                target.extend_from_self(source_column.borrow(), lower .. upper);
1024                expected.extend_from_slice(&source[lower .. upper]);
1025                assert_eq!(target.len(), expected.len());
1026                for (index, item) in expected.iter().enumerate() {
1027                    assert_eq!(target.get(index), *item, "item {index} of {}", expected.len());
1028                }
1029            }
1030        }
1031    }
1032}
1033
1034pub mod option {
1035
1036    use alloc::{vec::Vec, string::String};
1037
1038    use crate::{Clear, Columnar, Container, Len, IndexMut, Index, IndexAs, Push, Borrow};
1039    use crate::RankSelect;
1040
1041#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
1042    #[derive(Copy, Clone, Debug, Default, PartialEq)]
1043    pub struct Options<TC, CC=Vec<u64>, VC=Vec<u64>, WC=[u64; 2]> {
1044        /// Uses two bits for each item, one to indicate the variant and one (amortized)
1045        /// to enable efficient rank determination.
1046        pub indexes: RankSelect<CC, VC, WC>,
1047        pub somes: TC,
1048    }
1049
1050    impl<T: Columnar> Columnar for Option<T> {
1051        fn copy_from<'a>(&mut self, other: crate::Ref<'a, Self>) {
1052            match (&mut *self, other) {
1053                (Some(x), Some(y)) => { x.copy_from(y); }
1054                (_, other) => { *self = Self::into_owned(other); }
1055            }
1056        }
1057        fn into_owned<'a>(other: crate::Ref<'a, Self>) -> Self {
1058            other.map(|x| T::into_owned(x))
1059        }
1060        type Container = Options<T::Container>;
1061    }
1062
1063    impl<TC: Borrow> Borrow for Options<TC> {
1064        type Ref<'a> = Option<TC::Ref<'a>> where TC: 'a;
1065        type Borrowed<'a> = Options<TC::Borrowed<'a>, &'a [u64], &'a [u64], &'a [u64]> where TC: 'a;
1066        fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
1067            Options {
1068                indexes: self.indexes.borrow(),
1069                somes: self.somes.borrow(),
1070            }
1071        }
1072        #[inline(always)]
1073        fn reborrow<'b, 'a: 'b>(thing: Self::Borrowed<'a>) -> Self::Borrowed<'b> where TC: 'a {
1074            Options {
1075                indexes: RankSelect::<Vec<u64>, Vec<u64>>::reborrow(thing.indexes),
1076                somes: TC::reborrow(thing.somes),
1077            }
1078        }
1079        #[inline(always)]
1080        fn reborrow_ref<'b, 'a: 'b>(thing: Self::Ref<'a>) -> Self::Ref<'b> where Self: 'a {
1081            thing.map(TC::reborrow_ref)
1082        }
1083    }
1084
1085    impl<TC: Container> Container for Options<TC> {
1086        #[inline(always)]
1087        fn extend_from_self(&mut self, other: Self::Borrowed<'_>, range: core::ops::Range<usize>) {
1088            if !range.is_empty() {
1089                // The `Some` variants in `range` occupy a contiguous range of `other.somes`,
1090                // located by the number of `Some` variants before the start and end of `range`.
1091                let somes_start = other.indexes.rank(range.start);
1092                let somes_end = other.indexes.rank(range.end);
1093
1094                self.indexes.extend_from_self(&other.indexes, range);
1095                self.somes.extend_from_self(other.somes, somes_start .. somes_end);
1096            }
1097        }
1098
1099        fn reserve_for<'a, I>(&mut self, selves: I) where Self: 'a, I: Iterator<Item = Self::Borrowed<'a>> + Clone {
1100            self.indexes.reserve_for(selves.clone().map(|x| x.indexes));
1101            self.somes.reserve_for(selves.map(|x| x.somes));
1102        }
1103    }
1104
1105    impl<'a, TC: crate::AsBytes<'a>, CC: crate::AsBytes<'a>, VC: crate::AsBytes<'a>> crate::AsBytes<'a> for Options<TC, CC, VC, &'a [u64]> {
1106        const SLICE_COUNT: usize = <RankSelect<CC, VC, &'a [u64]> as crate::AsBytes<'a>>::SLICE_COUNT + TC::SLICE_COUNT;
1107        #[inline]
1108        fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
1109            debug_assert!(index < Self::SLICE_COUNT);
1110            let idx_count = <RankSelect<CC, VC, &'a [u64]> as crate::AsBytes<'a>>::SLICE_COUNT;
1111            if index < idx_count {
1112                self.indexes.get_byte_slice(index)
1113            } else {
1114                self.somes.get_byte_slice(index - idx_count)
1115            }
1116        }
1117    }
1118
1119    impl <'a, TC: crate::FromBytes<'a>, CC: crate::FromBytes<'a>, VC: crate::FromBytes<'a>> crate::FromBytes<'a> for Options<TC, CC, VC, &'a [u64]> {
1120        const SLICE_COUNT: usize = <RankSelect<CC, VC, &'a [u64]>>::SLICE_COUNT + TC::SLICE_COUNT;
1121        #[inline(always)]
1122        fn from_bytes(bytes: &mut impl Iterator<Item=&'a [u8]>) -> Self {
1123            Self {
1124                indexes: crate::FromBytes::from_bytes(bytes),
1125                somes: crate::FromBytes::from_bytes(bytes),
1126            }
1127        }
1128        #[inline(always)]
1129        fn from_store(store: &crate::bytes::indexed::DecodedStore<'a>, offset: &mut usize) -> Self {
1130            Self {
1131                indexes: crate::FromBytes::from_store(store, offset),
1132                somes: TC::from_store(store, offset),
1133            }
1134        }
1135        #[inline(always)]
1136        fn element_sizes(sizes: &mut Vec<usize>) -> Result<(), String> {
1137            <RankSelect<CC, VC, &'a [u64]>>::element_sizes(sizes)?;
1138            TC::element_sizes(sizes)?;
1139            Ok(())
1140        }
1141    }
1142
1143    impl<T, CC, VC: Len, WC: IndexAs<u64>> Len for Options<T, CC, VC, WC> {
1144        #[inline(always)] fn len(&self) -> usize { self.indexes.len() }
1145    }
1146
1147    impl<TC: Index, CC: IndexAs<u64> + Len, VC: IndexAs<u64> + Len, WC: IndexAs<u64>> Index for Options<TC, CC, VC, WC> {
1148        type Ref = Option<TC::Ref>;
1149        #[inline(always)]
1150        fn get(&self, index: usize) -> Self::Ref {
1151            if self.indexes.get(index) {
1152                Some(self.somes.get(self.indexes.rank(index)))
1153            } else {
1154                None
1155            }
1156        }
1157    }
1158    impl<'a, TC, CC: IndexAs<u64> + Len, VC: IndexAs<u64> + Len, WC: IndexAs<u64>> Index for &'a Options<TC, CC, VC, WC>
1159    where &'a TC: Index
1160    {
1161        type Ref = Option<<&'a TC as Index>::Ref>;
1162        #[inline(always)]
1163        fn get(&self, index: usize) -> Self::Ref {
1164            if self.indexes.get(index) {
1165                Some((&self.somes).get(self.indexes.rank(index)))
1166            } else {
1167                None
1168            }
1169        }
1170    }
1171    impl<TC: IndexMut, CC: IndexAs<u64> + Len, VC: IndexAs<u64> + Len> IndexMut for Options<TC, CC, VC> {
1172        type IndexMut<'a> = Option<TC::IndexMut<'a>> where TC: 'a, CC: 'a, VC: 'a;
1173        #[inline(always)]
1174        fn get_mut(&mut self, index: usize) -> Self::IndexMut<'_> {
1175            if self.indexes.get(index) {
1176                Some(self.somes.get_mut(self.indexes.rank(index)))
1177            } else {
1178                None
1179            }
1180        }
1181    }
1182
1183    impl<T, TC: Push<T> + Len> Push<Option<T>> for Options<TC> {
1184        #[inline]
1185        fn push(&mut self, item: Option<T>) {
1186            match item {
1187                Some(item) => {
1188                    self.indexes.push(true);
1189                    self.somes.push(item);
1190                }
1191                None => {
1192                    self.indexes.push(false);
1193                }
1194            }
1195        }
1196    }
1197    impl<'a, T, TC: Push<&'a T> + Len> Push<&'a Option<T>> for Options<TC> {
1198        #[inline]
1199        fn push(&mut self, item: &'a Option<T>) {
1200            match item {
1201                Some(item) => {
1202                    self.indexes.push(true);
1203                    self.somes.push(item);
1204                }
1205                None => {
1206                    self.indexes.push(false);
1207                }
1208            }
1209        }
1210    }
1211
1212    impl<TC, CC, VC, WC> Options<TC, CC, VC, WC> {
1213        /// Returns the inner container if all elements are `Some`, or `None`.
1214        pub fn try_unwrap(self) -> Option<TC> where TC: Len, VC: Len, WC: IndexAs<u64> {
1215            if self.somes.len() == self.indexes.len() { Some(self.somes) } else { None }
1216        }
1217        /// True if all elements are `None`.
1218        pub fn is_all_none(&self) -> bool where TC: Len {
1219            self.somes.is_empty()
1220        }
1221    }
1222
1223    impl<TC: Clear> Clear for Options<TC> {
1224        fn clear(&mut self) {
1225            self.indexes.clear();
1226            self.somes.clear();
1227        }
1228    }
1229
1230    #[cfg(test)]
1231    mod test {
1232        use alloc::vec::Vec;
1233
1234        use crate::Columnar;
1235        use crate::common::{Index, Len};
1236        use crate::Options;
1237
1238        #[test]
1239        fn round_trip_some() {
1240            // Type annotation is important to avoid some inference overflow.
1241            let store: Options<Vec<i32>> = Columnar::into_columns((0..100).map(Some));
1242            assert_eq!(store.len(), 100);
1243            assert!((&store).index_iter().zip(0..100).all(|(a, b)| a == Some(&b)));
1244        }
1245
1246        #[test]
1247        fn round_trip_none() {
1248            let store = Columnar::into_columns((0..100).map(|_x| None::<i32>));
1249            assert_eq!(store.len(), 100);
1250            let foo = &store;
1251            assert!(foo.index_iter().zip(0..100).all(|(a, _b)| a == None));
1252        }
1253
1254        #[test]
1255        fn round_trip_mixed() {
1256            // Type annotation is important to avoid some inference overflow.
1257            let store: Options<Vec<i32>>  = Columnar::into_columns((0..100).map(|x| if x % 2 == 0 { Some(x) } else { None }));
1258            assert_eq!(store.len(), 100);
1259            assert!((&store).index_iter().zip(0..100).all(|(a, b)| a == if b % 2 == 0 { Some(&b) } else { None }));
1260        }
1261
1262        /// `extend_from_self` agrees with pushing one element at a time, for random
1263        /// sources, targets and ranges.
1264        #[test]
1265        fn extend_from_self_random() {
1266            use crate::{Borrow, Container};
1267            use crate::sums::test_util::Lcg;
1268            let mut rng = Lcg(3);
1269            for _ in 0..200 {
1270                let source: Vec<Option<u32>> = (0..(rng.next() % 3000)).map(|i| if rng.next() % 3 == 0 { Some(i as u32) } else { None }).collect();
1271                let mut expected: Vec<Option<u32>> = (0..(rng.next() % 1500)).map(|i| if rng.next() % 2 == 0 { Some(100_000 + i as u32) } else { None }).collect();
1272                let source_column: Options<Vec<u32>> = Columnar::into_columns(source.iter().copied());
1273                let mut target: Options<Vec<u32>> = Columnar::into_columns(expected.iter().copied());
1274                let lower = (rng.next() as usize) % (source.len() + 1);
1275                let upper = lower + (rng.next() as usize) % (source.len() - lower + 1);
1276                target.extend_from_self(source_column.borrow(), lower .. upper);
1277                expected.extend_from_slice(&source[lower .. upper]);
1278                assert_eq!(target.len(), expected.len());
1279                for (index, item) in expected.iter().enumerate() {
1280                    assert_eq!(target.get(index), *item, "item {index} of {}", expected.len());
1281                }
1282            }
1283        }
1284    }
1285}
1286
1287#[cfg(test)]
1288pub(crate) mod test_util {
1289    /// A small deterministic generator for randomized tests.
1290    pub struct Lcg(pub u64);
1291    impl Lcg {
1292        pub fn next(&mut self) -> u64 {
1293            self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
1294            self.0 >> 11
1295        }
1296    }
1297}
1298
1299pub mod discriminant {
1300
1301    use alloc::{vec::Vec, string::String};
1302    use crate::{Clear, Container, Len, Index, IndexAs, Borrow};
1303
1304    /// Tracks variant discriminants and offsets for enum containers.
1305    ///
1306    /// Uses two arrays (`variant` and `offset`) with three states:
1307    /// - **Empty**: both arrays empty, length is 0.
1308    /// - **Homogeneous**: `variant` is empty, `offset` holds `[tag, count]` where
1309    ///   `tag = variant_index + 1`. All elements share a single variant with
1310    ///   identity offsets (element `i` maps to offset `i`).
1311    /// - **Heterogeneous**: `variant` has per-element discriminants (`u8`),
1312    ///   `offset` has per-element offsets into variant containers (`u64`).
1313    #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
1314    #[derive(Clone, Debug, Default, PartialEq)]
1315    pub struct Discriminant<CVar = Vec<u8>, COff = Vec<u64>> {
1316        /// Per-element variant discriminants; empty when homogeneous.
1317        pub variant: CVar,
1318        /// Per-element offsets (heterogeneous), or `[tag, count]` (homogeneous), or empty.
1319        pub offset: COff,
1320    }
1321
1322    impl<CVar: Copy, COff: Copy> Copy for Discriminant<CVar, COff> {}
1323
1324    impl Discriminant {
1325        /// Push a variant discriminant and the offset into its variant container.
1326        #[inline]
1327        pub fn push(&mut self, variant: u8, offset: u64) {
1328            let tag = variant as u64 + 1;
1329            if self.variant.is_empty() {
1330                if self.offset.is_empty() {
1331                    // Empty → start homogeneous: offset = [tag, 1].
1332                    self.offset.push(tag);
1333                    self.offset.push(1);
1334                } else if self.offset[0] == tag {
1335                    // Same variant; stay homogeneous, increment count.
1336                    self.offset[1] += 1;
1337                } else {
1338                    // Different variant; transition to heterogeneous.
1339                    let prev = (self.offset[0] - 1) as u8;
1340                    let count = self.offset[1];
1341                    self.variant.reserve(count as usize + 1);
1342                    self.offset.clear();
1343                    self.offset.reserve(count as usize + 1);
1344                    for i in 0..count {
1345                        self.variant.push(prev);
1346                        self.offset.push(i);
1347                    }
1348                    self.variant.push(variant);
1349                    self.offset.push(offset);
1350                }
1351            } else {
1352                // Already heterogeneous.
1353                self.variant.push(variant);
1354                self.offset.push(offset);
1355            }
1356        }
1357
1358        /// Pre-allocate for the given borrowed discriminants.
1359        pub fn reserve_for<'a>(&mut self, selves: impl Iterator<Item = Discriminant<&'a [u8], &'a [u64]>> + Clone) {
1360            self.variant.reserve_for(selves.clone().map(|x| x.variant));
1361            self.offset.reserve_for(selves.map(|x| x.offset));
1362        }
1363    }
1364
1365    impl<CVar: Len, COff: Len> Discriminant<CVar, COff> {
1366        /// True if elements have mixed variants, with per-element discriminants and offsets.
1367        #[inline]
1368        pub fn is_heterogeneous(&self) -> bool {
1369            !self.variant.is_empty()
1370        }
1371        /// Returns `Some(variant)` if all elements share a single variant.
1372        #[inline]
1373        pub fn homogeneous(&self) -> Option<u8> where COff: IndexAs<u64> {
1374            if self.variant.is_empty() && self.offset.len() >= 2 {
1375                Some((self.offset.index_as(0) - 1) as u8)
1376            } else {
1377                None
1378            }
1379        }
1380        /// Returns `(variant, offset)` for the element at `index`.
1381        #[inline(always)]
1382        pub fn get(&self, index: usize) -> (u8, u64) where CVar: IndexAs<u8>, COff: IndexAs<u64> {
1383            if self.is_heterogeneous() {
1384                (self.variant.index_as(index), self.offset.index_as(index))
1385            } else {
1386                let tag: u64 = self.offset.index_as(0);
1387                ((tag - 1) as u8, index as u64)
1388            }
1389        }
1390    }
1391
1392    impl<CVar: Len, COff: Len + IndexAs<u64>> Len for Discriminant<CVar, COff> {
1393        #[inline(always)]
1394        fn len(&self) -> usize {
1395            if self.is_heterogeneous() { self.variant.len() }
1396            else if self.offset.len() >= 2 { self.offset.index_as(1) as usize }
1397            else { 0 }
1398        }
1399    }
1400
1401    // Index for the borrowed form: returns (variant, offset).
1402    impl<'a> Index for Discriminant<&'a [u8], &'a [u64]> {
1403        type Ref = (u8, u64);
1404        #[inline(always)]
1405        fn get(&self, index: usize) -> (u8, u64) {
1406            if self.is_heterogeneous() {
1407                (self.variant.index_as(index), self.offset.index_as(index))
1408            } else {
1409                ((self.offset[0] - 1) as u8, index as u64)
1410            }
1411        }
1412    }
1413
1414    // Borrow
1415    impl Borrow for Discriminant {
1416        type Ref<'a> = (u8, u64);
1417        type Borrowed<'a> = Discriminant<&'a [u8], &'a [u64]>;
1418        #[inline(always)]
1419        fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
1420            Discriminant {
1421                variant: &self.variant[..],
1422                offset: &self.offset[..],
1423            }
1424        }
1425        #[inline(always)]
1426        fn reborrow<'b, 'a: 'b>(thing: Self::Borrowed<'a>) -> Self::Borrowed<'b> {
1427            Discriminant {
1428                variant: thing.variant,
1429                offset: thing.offset,
1430            }
1431        }
1432        #[inline(always)]
1433        fn reborrow_ref<'b, 'a: 'b>(thing: Self::Ref<'a>) -> Self::Ref<'b> { thing }
1434    }
1435
1436    impl<CVar: Clear, COff: Clear> Clear for Discriminant<CVar, COff> {
1437        #[inline(always)]
1438        fn clear(&mut self) {
1439            self.variant.clear();
1440            self.offset.clear();
1441        }
1442    }
1443
1444
1445    // AsBytes for Discriminant, generic over container types.
1446    impl<'a, CVar: crate::AsBytes<'a>, COff: crate::AsBytes<'a>> crate::AsBytes<'a> for Discriminant<CVar, COff> {
1447        const SLICE_COUNT: usize = CVar::SLICE_COUNT + COff::SLICE_COUNT;
1448        #[inline]
1449        fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
1450            debug_assert!(index < Self::SLICE_COUNT);
1451            if index < CVar::SLICE_COUNT {
1452                self.variant.get_byte_slice(index)
1453            } else {
1454                self.offset.get_byte_slice(index - CVar::SLICE_COUNT)
1455            }
1456        }
1457    }
1458
1459    // FromBytes for borrowed form
1460    impl<'a> crate::FromBytes<'a> for Discriminant<&'a [u8], &'a [u64]> {
1461        const SLICE_COUNT: usize = <&'a [u8]>::SLICE_COUNT + <&'a [u64]>::SLICE_COUNT;
1462        #[inline(always)]
1463        fn from_bytes(bytes: &mut impl Iterator<Item=&'a [u8]>) -> Self {
1464            let variant = crate::FromBytes::from_bytes(bytes);
1465            let offset = crate::FromBytes::from_bytes(bytes);
1466            Self { variant, offset }
1467        }
1468        #[inline(always)]
1469        fn from_store(store: &crate::bytes::indexed::DecodedStore<'a>, offset: &mut usize) -> Self {
1470            let variant = crate::FromBytes::from_store(store, offset);
1471            let offset_field = crate::FromBytes::from_store(store, offset);
1472            Self { variant, offset: offset_field }
1473        }
1474        #[inline(always)]
1475        fn element_sizes(sizes: &mut Vec<usize>) -> Result<(), String> {
1476            <&[u8]>::element_sizes(sizes)?;
1477            <&[u64]>::element_sizes(sizes)?;
1478            Ok(())
1479        }
1480    }
1481
1482    #[cfg(test)]
1483    mod test {
1484        use crate::Len;
1485
1486        #[test]
1487        fn homogeneous_push() {
1488            let mut d = super::Discriminant::default();
1489            d.push(2, 0);
1490            d.push(2, 1);
1491            d.push(2, 2);
1492            assert_eq!(d.len(), 3);
1493            assert_eq!(d.homogeneous(), Some(2));
1494            assert!(d.variant.is_empty());
1495            // offset holds [tag, count] = [3, 3] in homogeneous mode.
1496            assert_eq!(d.offset, vec![3, 3]);
1497        }
1498
1499        #[test]
1500        fn heterogeneous_transition() {
1501            let mut d = super::Discriminant::default();
1502            d.push(0, 0);
1503            d.push(0, 1);
1504            d.push(1, 0); // transition
1505            assert_eq!(d.len(), 3);
1506            assert_eq!(d.homogeneous(), None);
1507            assert_eq!(d.variant, vec![0, 0, 1]);
1508            assert_eq!(d.offset, vec![0, 1, 0]);
1509        }
1510
1511        #[test]
1512        fn clear_resets() {
1513            use crate::Clear;
1514            let mut d = super::Discriminant::default();
1515            d.push(1, 0);
1516            d.push(1, 1);
1517            d.clear();
1518            assert_eq!(d.len(), 0);
1519            // After clear, first push starts homogeneous again.
1520            d.push(3, 0);
1521            assert_eq!(d.homogeneous(), Some(3));
1522            assert_eq!(d.len(), 1);
1523        }
1524
1525        #[test]
1526        fn borrow_index() {
1527            use crate::Borrow;
1528            let mut d = super::Discriminant::default();
1529            d.push(2, 0);
1530            d.push(2, 1);
1531            d.push(2, 2);
1532            let b = d.borrow();
1533            assert_eq!(b.get(0), (2, 0));
1534            assert_eq!(b.get(1), (2, 1));
1535            assert_eq!(b.get(2), (2, 2));
1536        }
1537
1538        #[test]
1539        fn borrow_index_heterogeneous() {
1540            use crate::Borrow;
1541            let mut d = super::Discriminant::default();
1542            d.push(0, 0);
1543            d.push(1, 0);
1544            d.push(0, 1);
1545            let b = d.borrow();
1546            assert_eq!(b.get(0), (0, 0));
1547            assert_eq!(b.get(1), (1, 0));
1548            assert_eq!(b.get(2), (0, 1));
1549        }
1550    }
1551}