ty_python_core/
rank.rs

1//! A boxed bit slice that supports a constant-time `rank` operation.
2
3use bitvec::prelude::{BitBox, BitVec, Msb0, bitvec};
4use get_size2::GetSize;
5
6/// A boxed bit slice that supports a constant-time `rank` operation.
7///
8/// This can be used to "shrink" a large vector, where you only need to keep certain elements, and
9/// you want to continue to use the index in the large vector to identify each element.
10///
11/// First you create a new smaller vector, keeping only the elements of the large vector that you
12/// care about. Now you need a way to translate an index into the large vector (which no longer
13/// exists) into the corresponding index into the smaller vector. To do that, you create a bit
14/// slice, containing a bit for every element of the original large vector. Each bit in the bit
15/// slice indicates whether that element of the large vector was kept in the smaller vector. And
16/// the `rank` of the bit gives us the index of the element in the smaller vector.
17///
18/// However, the naive implementation of `rank` is O(n) in the size of the bit slice. To address
19/// that, we use a standard trick: we divide the bit slice into 64-bit chunks, and when
20/// constructing the bit slice, precalculate the rank of the first bit in each chunk. Then, to
21/// calculate the rank of an arbitrary bit, we first grab the precalculated rank of the chunk that
22/// bit belongs to, and add the rank of the bit within its (fixed-sized) chunk.
23///
24/// This trick adds O(1.5) bits of overhead per large vector element on 64-bit platforms, and O(2)
25/// bits of overhead on 32-bit platforms.
26#[derive(Clone, Debug, Eq, Hash, PartialEq, GetSize)]
27pub struct RankBitBox {
28    #[get_size(size_fn = bit_box_size)]
29    bits: RankBitBoxStorage,
30    chunk_ranks: Box<[u32]>,
31}
32
33pub type RankBitBoxStorage = BitBox<Chunk, Msb0>;
34pub type RankBitBoxVec = BitVec<Chunk, Msb0>;
35
36fn bit_box_size(bits: &RankBitBoxStorage) -> usize {
37    std::mem::size_of_val(bits.as_raw_slice())
38}
39
40// bitvec does not support `u64` as a Store type on 32-bit platforms
41#[cfg(target_pointer_width = "64")]
42type Chunk = u64;
43#[cfg(not(target_pointer_width = "64"))]
44type Chunk = u32;
45
46const CHUNK_SIZE: usize = Chunk::BITS as usize;
47
48impl RankBitBox {
49    pub fn bits_with_capacity(cap: usize) -> RankBitBoxVec {
50        bitvec![Chunk, Msb0; 0; cap]
51    }
52
53    pub fn from_bits(bits: RankBitBoxVec) -> Self {
54        let chunk_ranks = bits
55            .as_raw_slice()
56            .iter()
57            .scan(0u32, |rank, chunk| {
58                let result = *rank;
59                *rank += chunk.count_ones();
60                Some(result)
61            })
62            .collect();
63        let bits = bits.into();
64        Self { bits, chunk_ranks }
65    }
66
67    #[inline]
68    pub fn len(&self) -> usize {
69        self.bits.len()
70    }
71
72    #[inline]
73    pub fn is_empty(&self) -> bool {
74        self.bits.is_empty()
75    }
76
77    #[inline]
78    pub fn get_bit(&self, index: usize) -> Option<bool> {
79        self.bits.get(index).map(|bit| *bit)
80    }
81
82    #[inline]
83    pub fn iter_ones(&self) -> impl DoubleEndedIterator<Item = usize> + '_ {
84        self.bits.iter_ones()
85    }
86
87    /// Returns the number of bits _before_ (and not including) the given index that are set.
88    #[inline]
89    pub fn rank(&self, index: usize) -> u32 {
90        let chunk_index = index / CHUNK_SIZE;
91        let index_within_chunk = index % CHUNK_SIZE;
92        let chunk_rank = self.chunk_ranks[chunk_index];
93        if index_within_chunk == 0 {
94            return chunk_rank;
95        }
96
97        // To calculate the rank within the bit's chunk, we zero out the requested bit and every
98        // bit to the right, then count the number of 1s remaining (i.e., to the left of the
99        // requested bit).
100        let chunk = self.bits.as_raw_slice()[chunk_index];
101        let chunk_mask = Chunk::MAX << (CHUNK_SIZE - index_within_chunk);
102        let rank_within_chunk = (chunk & chunk_mask).count_ones();
103        chunk_rank + rank_within_chunk
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use std::mem::size_of;
110
111    use get_size2::GetSize;
112
113    use super::{CHUNK_SIZE, Chunk, RankBitBox};
114
115    #[test]
116    fn heap_size_includes_bits_and_chunk_ranks() {
117        let bit_count = CHUNK_SIZE + 1;
118        let bits = RankBitBox::from_bits(RankBitBox::bits_with_capacity(bit_count));
119        let chunk_count = bit_count.div_ceil(CHUNK_SIZE);
120
121        assert_eq!(
122            bits.get_heap_size(),
123            chunk_count * (size_of::<Chunk>() + size_of::<u32>())
124        );
125    }
126}
ty_python_core/rank.rs

ty_python_core/
rank.rs
