entropy_map/
rank.rs

1//! `RankedBits` efficiently handles rank queries on bit vectors.
2//! Optimized for minimal memory usage with ~3.125% overhead and fast lookups, it supports the
3//! crate's focus on low-latency hash maps. For detailed methodology, refer to the related paper:
4//! [Engineering Compact Data Structures for Rank and Select Queries on Bit Vectors](https://arxiv.org/pdf/2206.01149.pdf).
5
6use std::mem::size_of_val;
7
8/// Size of the L2 block in bits.
9const L2_BIT_SIZE: usize = 512;
10/// Size of the L1 block in bits, calculated as a multiple of the L2 block size.
11const L1_BIT_SIZE: usize = 8 * L2_BIT_SIZE;
12
13/// Trait for efficient bit-level operations on ranked bit sequences.
14///
15/// This trait is designed to provide consistent methods for accessing ranked bit sequences in both
16/// their standard and `Archived` formats (utilizing the `rkyv` library).
17pub trait RankedBitsAccess {
18    /// Returns the number of set bits up to `idx`, or `None` if the bit at `idx` is not set.
19    fn rank(&self, idx: usize) -> Option<usize>;
20
21    /// Inner implementation of `rank` with `bits` and `l12_ranks` passed from different implementations.
22    ///
23    /// # Safety
24    /// This method is unsafe because `idx` must be within the bounds of the bits stored in `RankedBitsAccess`.
25    /// An index out of bounds can lead to undefined behavior.
26    #[inline]
27    unsafe fn rank_impl<T: L12RankAccess>(bits: &[u64], l12_ranks: &T, idx: usize) -> Option<usize> {
28        let word_idx = idx / 64;
29        let bit_idx = idx % 64;
30        let word = *bits.get_unchecked(word_idx);
31
32        if (word & (1u64 << bit_idx)) == 0 {
33            return None;
34        }
35
36        let l1_pos = idx / L1_BIT_SIZE;
37        let l2_pos = (idx % L1_BIT_SIZE) / L2_BIT_SIZE;
38
39        let idx_within_l2 = idx % L2_BIT_SIZE;
40        let blocks_num = idx_within_l2 / 64;
41        let offset = (idx / L2_BIT_SIZE) * 8;
42        let block = bits.get_unchecked(offset..offset + blocks_num);
43
44        let block_rank = block.iter().map(|&x| x.count_ones() as usize).sum::<usize>();
45
46        let word = *bits.get_unchecked(offset + blocks_num);
47        let word_mask = ((1u64 << (idx_within_l2 % 64)) - 1) * (idx_within_l2 > 0) as u64;
48        let word_rank = (word & word_mask).count_ones() as usize;
49
50        let (l1_rank, l2_rank) = l12_ranks.l12_ranks(l1_pos, l2_pos);
51        let total_rank = l1_rank + l2_rank + block_rank + word_rank;
52
53        Some(total_rank)
54    }
55}
56
57#[derive(Debug, Default)]
58#[cfg_attr(feature = "rkyv_derive", derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize))]
59#[cfg_attr(feature = "rkyv_derive", archive_attr(derive(rkyv::CheckBytes)))]
60pub struct RankedBits {
61    /// The bit vector represented as an array of u64 integers.
62    pub(crate) bits: Box<[u64]>,
63    /// Precomputed rank information for L1 and L2 blocks.
64    l12_ranks: Box<[L12Rank]>,
65}
66
67/// Custom Serde logic that stores `Box<[u64]>` as a single little-endian byte blob.
68///
69/// Unlike the default `seq` encoding (which in msgpack frames each `u64` as a
70/// `uint64` tag + 8 bytes = 9 bytes/element), this emits one `bin` blob via
71/// `serialize_bytes`, eliminating per-element framing. Endianness is explicit
72/// (little-endian) so the on-wire form is portable across architectures,
73/// consistent with the crate's existing `L12Rank` LE convention.
74#[cfg(feature = "serde")]
75impl serde::Serialize for RankedBits {
76    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
77    where
78        S: serde::Serializer,
79    {
80        let mut out = serde_bytes::ByteBuf::with_capacity(self.bits.len() * 8);
81        for word in &self.bits {
82            out.extend_from_slice(&word.to_le_bytes());
83        }
84        out.serialize(serializer)
85    }
86}
87#[cfg(feature = "serde")]
88impl<'de> serde::Deserialize<'de> for RankedBits {
89    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
90    where
91        D: serde::Deserializer<'de>,
92    {
93        let buf = serde_bytes::ByteBuf::deserialize(deserializer)?;
94        if buf.len() % 8 != 0 {
95            return Err(serde::de::Error::invalid_length(
96                buf.len(),
97                &"a byte length that is a multiple of 8",
98            ));
99        }
100        let bits = buf
101            .as_chunks::<8>()
102            .0
103            .iter()
104            .map(|&c| u64::from_le_bytes(c))
105            .collect::<Box<[u64]>>();
106
107        Ok(RankedBits::new(bits))
108    }
109}
110
111/// L12Rank represents l1 and l2 bit ranks stored inside 16 bytes (little endian).
112/// NB: it's important to use `[u8; 16]` instead of `u128` for `rkyv` versions 0.7.X
113/// because of alignment differences between `x86_64` and `aarch64` architectures.
114/// See https://github.com/rkyv/rkyv/issues/409 for more details.
115#[derive(Debug)]
116#[cfg_attr(feature = "rkyv_derive", derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize))]
117#[cfg_attr(feature = "rkyv_derive", archive_attr(derive(rkyv::CheckBytes)))]
118#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
119pub struct L12Rank(#[cfg_attr(feature = "serde", serde(with = "serde_bytes"))] [u8; 16]);
120
121/// Trait used to access archived and non-archived L1 and L2 ranks
122pub trait L12RankAccess {
123    /// Return `L12Rank` as `u128`
124    fn l12_rank(&self, l1_pos: usize) -> u128;
125
126    /// Return `l1_rank` and `l2_rank`
127    #[inline]
128    fn l12_ranks(&self, l1_pos: usize, l2_pos: usize) -> (usize, usize) {
129        let l12_rank = self.l12_rank(l1_pos);
130        let l1_rank = (l12_rank & 0xFFFFFFFFFFF) as usize;
131        let l2_rank = ((l12_rank >> (32 + 12 * l2_pos)) & 0xFFF) as usize;
132        (l1_rank, l2_rank)
133    }
134}
135
136impl L12RankAccess for Box<[L12Rank]> {
137    #[inline]
138    fn l12_rank(&self, l1_pos: usize) -> u128 {
139        u128::from_le_bytes(unsafe { self.get_unchecked(l1_pos).0 })
140    }
141}
142
143#[cfg(feature = "rkyv_derive")]
144impl L12RankAccess for rkyv::boxed::ArchivedBox<[ArchivedL12Rank]> {
145    #[inline]
146    fn l12_rank(&self, l1_pos: usize) -> u128 {
147        u128::from_le_bytes(unsafe { self.get_unchecked(l1_pos).0 })
148    }
149}
150
151impl From<u128> for L12Rank {
152    #[inline]
153    fn from(v: u128) -> Self {
154        L12Rank(v.to_le_bytes())
155    }
156}
157
158impl RankedBits {
159    /// Returns the number of ones in `bits`
160    pub fn count_ones(&self) -> usize {
161        self.bits.iter().map(|x| x.count_ones() as usize).sum()
162    }
163
164    /// Initializes `RankedBits` with a provided bit vector.
165    pub fn new(bits: Box<[u64]>) -> Self {
166        let (blocks, remainder) = bits.as_chunks::<64>();
167        let mut l12_ranks = Vec::with_capacity(bits.len().div_ceil(64));
168        let mut l1_rank: u128 = 0;
169
170        for block64 in blocks {
171            let mut l12_rank = 0u128;
172            let mut sum = 0u16;
173            for (i, block8) in block64.as_chunks::<8>().0.iter().enumerate() {
174                sum += block8.iter().map(|&x| x.count_ones() as u16).sum::<u16>();
175                l12_rank += (sum as u128) << (i * 12);
176            }
177            l12_rank = (l12_rank << 44) | l1_rank;
178            l12_ranks.push(l12_rank.into());
179            l1_rank += sum as u128;
180        }
181
182        if !remainder.is_empty() {
183            let mut l12_rank = 0u128;
184            let mut sum = 0u16;
185            for (i, block) in remainder.chunks(8).enumerate() {
186                sum += block.iter().map(|&x| x.count_ones() as u16).sum::<u16>();
187                l12_rank += (sum as u128) << (i * 12);
188            }
189            l12_rank = (l12_rank << 44) | l1_rank;
190            l12_ranks.push(l12_rank.into());
191        }
192
193        RankedBits { bits, l12_ranks: l12_ranks.into_boxed_slice() }
194    }
195
196    /// Returns the total number of bytes occupied by `RankedBits`
197    pub fn size(&self) -> usize {
198        size_of_val(self) + size_of_val(self.bits.as_ref()) + size_of_val(self.l12_ranks.as_ref())
199    }
200}
201
202/// Implement `rank` for `Archived` version of `RankedBits` if feature is enabled
203impl RankedBitsAccess for RankedBits {
204    #[inline]
205    fn rank(&self, idx: usize) -> Option<usize> {
206        unsafe { Self::rank_impl(&self.bits, &self.l12_ranks, idx) }
207    }
208}
209
210/// Implement `rank` for `Archived` version of `RankedBits` if feature is enabled
211#[cfg(feature = "rkyv_derive")]
212impl RankedBitsAccess for ArchivedRankedBits {
213    #[inline]
214    fn rank(&self, idx: usize) -> Option<usize> {
215        unsafe { Self::rank_impl(&self.bits, &self.l12_ranks, idx) }
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use bitvec::order::Lsb0;
223    use bitvec::vec::BitVec;
224    use rand::distributions::Standard;
225    use rand::Rng;
226
227    #[test]
228    fn test_rank_and_get() {
229        let bits = vec![
230            0b11001010, // 4 set bits
231            0b00110111, // 5 set bits
232            0b11110000, // 4 set bits
233        ];
234
235        let ranked_bits = RankedBits::new(bits.into_boxed_slice());
236        assert_eq!(ranked_bits.rank(0), None); // No set bits before the first
237        assert_eq!(ranked_bits.rank(7), Some(3)); // 3 set bits set before 7-th bit
238    }
239
240    #[test]
241    fn test_random_bits() {
242        let rng = rand::thread_rng();
243        let bits: Vec<u64> = rng.sample_iter(Standard).take(1001).collect();
244        let ranked_bits = RankedBits::new(bits.clone().into_boxed_slice());
245        let bv = BitVec::<u64, Lsb0>::from_slice(&bits);
246
247        for idx in 0..bv.len() {
248            if bv[idx] {
249                assert_eq!(
250                    ranked_bits.rank(idx).unwrap(),
251                    bv[..idx].count_ones(),
252                    "Rank mismatch at index {}",
253                    idx
254                );
255            }
256        }
257    }
258
259    #[cfg(feature = "serde")]
260    #[test]
261    fn test_serde() {
262        let rng = rand::thread_rng();
263        let bits: Vec<u64> = rng.sample_iter(Standard).take(1001).collect();
264        let ranked_bits = RankedBits::new(bits.clone().into_boxed_slice());
265
266        let bytes = rmp_serde::to_vec(&ranked_bits).unwrap();
267        let de: RankedBits = rmp_serde::from_slice(&bytes).unwrap();
268
269        // The deserialized `RankedBits` must answer every `rank` query identically
270        // to the original (the whole point of persisting the PHF artifacts).
271        for idx in 0..bits.len() * 64 {
272            assert_eq!(ranked_bits.rank(idx), de.rank(idx), "rank mismatch at {}", idx);
273        }
274    }
275}
entropy_map/rank.rs

entropy_map/
rank.rs
