blob: b5b69abdeeb62d86b941d1a34bae2c411c163033 (
plain)
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
|
//! Optimized bitset implementation.
type Entry = u128;
const ENTRY_NUM_BITS: usize = Entry::BITS as usize;
// TODO(max): Make a `SmallBitSet` and `LargeBitSet` and switch between them if `num_bits` fits in
// `Entry`.
#[derive(Clone)]
pub struct BitSet<T: Into<usize> + Copy> {
entries: Vec<Entry>,
num_bits: usize,
phantom: std::marker::PhantomData<T>,
}
impl<T: Into<usize> + Copy> BitSet<T> {
pub fn with_capacity(num_bits: usize) -> Self {
let num_entries = num_bits.div_ceil(ENTRY_NUM_BITS);
Self { entries: vec![0; num_entries], num_bits, phantom: Default::default() }
}
/// Returns whether the value was newly inserted: true if the set did not originally contain
/// the bit, and false otherwise.
pub fn insert(&mut self, idx: T) -> bool {
debug_assert!(idx.into() < self.num_bits);
let entry_idx = idx.into() / ENTRY_NUM_BITS;
let bit_idx = idx.into() % ENTRY_NUM_BITS;
let newly_inserted = (self.entries[entry_idx] & (1 << bit_idx)) == 0;
self.entries[entry_idx] |= 1 << bit_idx;
newly_inserted
}
/// Set all bits to 1.
pub fn insert_all(&mut self) {
for i in 0..self.entries.len() {
self.entries[i] = !0;
}
}
pub fn get(&self, idx: T) -> bool {
debug_assert!(idx.into() < self.num_bits);
let entry_idx = idx.into() / ENTRY_NUM_BITS;
let bit_idx = idx.into() % ENTRY_NUM_BITS;
(self.entries[entry_idx] & (1 << bit_idx)) != 0
}
/// Modify `self` to only have bits set if they are also set in `other`. Returns true if `self`
/// was modified, and false otherwise.
/// `self` and `other` must have the same number of bits.
pub fn intersect_with(&mut self, other: &Self) -> bool {
assert_eq!(self.num_bits, other.num_bits);
let mut changed = false;
for i in 0..self.entries.len() {
let before = self.entries[i];
self.entries[i] &= other.entries[i];
changed |= self.entries[i] != before;
}
changed
}
}
#[cfg(test)]
mod tests {
use super::BitSet;
#[test]
#[should_panic]
fn get_over_capacity_panics() {
let set = BitSet::with_capacity(0);
assert!(!set.get(0usize));
}
#[test]
fn with_capacity_defaults_to_zero() {
let set = BitSet::with_capacity(4);
assert!(!set.get(0usize));
assert!(!set.get(1usize));
assert!(!set.get(2usize));
assert!(!set.get(3usize));
}
#[test]
fn insert_sets_bit() {
let mut set = BitSet::with_capacity(4);
assert!(set.insert(1usize));
assert!(set.get(1usize));
}
#[test]
fn insert_with_set_bit_returns_false() {
let mut set = BitSet::with_capacity(4);
assert!(set.insert(1usize));
assert!(!set.insert(1usize));
}
#[test]
fn insert_all_sets_all_bits() {
let mut set = BitSet::with_capacity(4);
set.insert_all();
assert!(set.get(0usize));
assert!(set.get(1usize));
assert!(set.get(2usize));
assert!(set.get(3usize));
}
#[test]
#[should_panic]
fn intersect_with_panics_with_different_num_bits() {
let mut left: BitSet<usize> = BitSet::with_capacity(3);
let right = BitSet::with_capacity(4);
left.intersect_with(&right);
}
#[test]
fn intersect_with_keeps_only_common_bits() {
let mut left = BitSet::with_capacity(3);
let mut right = BitSet::with_capacity(3);
left.insert(0usize);
left.insert(1usize);
right.insert(1usize);
right.insert(2usize);
left.intersect_with(&right);
assert!(!left.get(0usize));
assert!(left.get(1usize));
assert!(!left.get(2usize));
}
}
|
