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