dtrees_rs/bitsets/
mod.rs

1//! Fixed-capacity bitsets.
2
3use std::ops::Index;
4
5/// Initial content of a [`Bitset`] of the given capacity.
6pub enum BitsetInit {
7    /// No element set.
8    Empty(usize),
9    /// Every element set.
10    Full(usize),
11}
12
13/// A fixed-capacity set of integers stored as 64-bit words.
14#[derive(Debug, Clone)]
15pub struct Bitset {
16    capacity: usize,
17    words: Vec<u64>,
18}
19
20/// Operations on a bitset.
21pub trait BitCollection {
22    /// A bitset of the given capacity and content.
23    fn new(init: BitsetInit) -> Self;
24
25    /// Number of elements set.
26    fn count(&self) -> usize;
27
28    /// Whether `index` is set.
29    fn test(&self, index: usize) -> bool;
30
31    /// Sets `index`.
32    fn set(&mut self, index: usize);
33
34    /// Clears `index`.
35    fn unset(&mut self, index: usize);
36
37    /// Whether no element is set.
38    fn is_empty(&self) -> bool;
39
40    /// Clears every element.
41    fn clear(&mut self);
42
43    /// Largest number of elements.
44    fn capacity(&self) -> usize;
45
46    /// Changes the capacity.
47    fn resize(&mut self, capacity: usize);
48
49    /// Keeps only the elements also in `other`.
50    fn intersect_with(&mut self, other: &Bitset);
51
52    /// Adds the elements of `other`.
53    fn union_with(&mut self, other: &Bitset);
54
55    /// Size of the intersection with `other`.
56    fn count_intersect_with(&self, other: &Bitset) -> usize;
57
58    /// Size of the intersection with each of `others`.
59    fn count_interest_with_many(&self, others: &[Bitset]) -> Vec<usize>;
60}
61
62impl BitCollection for Bitset {
63    fn new(init: BitsetInit) -> Self {
64        match init {
65            BitsetInit::Empty(n) => {
66                let word_count = n.div_ceil(64);
67                Self {
68                    capacity: n,
69                    words: vec![0u64; word_count],
70                }
71            }
72            BitsetInit::Full(n) => {
73                let word_count = n.div_ceil(64);
74                let mut words = vec![u64::MAX; word_count];
75                if n > 0 && n % 64 != 0 {
76                    if let Some(last) = words.last_mut() {
77                        *last = (1u64 << (n % 64)) - 1;
78                    }
79                }
80                Self { capacity: n, words }
81            }
82        }
83    }
84
85    fn count(&self) -> usize {
86        self.words
87            .iter()
88            .map(|&word| word.count_ones() as usize)
89            .sum()
90    }
91
92    fn test(&self, index: usize) -> bool {
93        debug_assert!(index < self.capacity, "Index out of bounds");
94        (self.words[index / 64] & (1u64 << (index % 64))) != 0
95    }
96
97    fn set(&mut self, index: usize) {
98        debug_assert!(index < self.capacity, "Index out of bounds");
99        self.words[index / 64] |= 1u64 << (index % 64);
100    }
101
102    fn unset(&mut self, index: usize) {
103        debug_assert!(index < self.capacity, "Index out of bounds");
104        self.words[index / 64] &= !(1u64 << (index % 64));
105    }
106
107    fn is_empty(&self) -> bool {
108        self.words.iter().all(|&word| word == 0)
109    }
110
111    fn clear(&mut self) {
112        self.words.fill(0);
113    }
114
115    fn capacity(&self) -> usize {
116        self.capacity
117    }
118
119    fn resize(&mut self, capacity: usize) {
120        let new_words = capacity.div_ceil(64);
121        match new_words.cmp(&self.words.len()) {
122            std::cmp::Ordering::Greater => self.words.resize(new_words, 0),
123            std::cmp::Ordering::Less => self.words.truncate(new_words),
124            std::cmp::Ordering::Equal => {}
125        }
126        self.capacity = 64 * self.words.len();
127    }
128
129    fn intersect_with(&mut self, other: &Bitset) {
130        debug_assert_eq!(
131            self.capacity, other.capacity,
132            "Bitsets must have the same capacity"
133        );
134        for (word, other_word) in self.words.iter_mut().zip(&other.words) {
135            *word &= other_word;
136        }
137    }
138
139    fn union_with(&mut self, other: &Bitset) {
140        debug_assert_eq!(
141            self.capacity, other.capacity,
142            "Bitsets must have the same capacity"
143        );
144        for (word, other_word) in self.words.iter_mut().zip(&other.words) {
145            *word |= other_word;
146        }
147    }
148
149    fn count_intersect_with(&self, other: &Bitset) -> usize {
150        debug_assert_eq!(
151            self.capacity, other.capacity,
152            "Bitsets must have the same capacity"
153        );
154        let mut count = 0;
155        for (word, other_word) in self.words.iter().zip(&other.words) {
156            count += (*word & *other_word).count_ones() as usize;
157        }
158        count
159    }
160
161    fn count_interest_with_many(&self, others: &[Bitset]) -> Vec<usize> {
162        others
163            .iter()
164            .map(|other| {
165                self.words
166                    .iter()
167                    .zip(&other.words)
168                    .map(|(&a, &b)| (a & b).count_ones() as usize)
169                    .sum()
170            })
171            .collect()
172    }
173}
174
175impl Index<usize> for Bitset {
176    type Output = u64;
177
178    fn index(&self, index: usize) -> &Self::Output {
179        debug_assert!(
180            index < self.words.len(),
181            "Index out for number of words bounds"
182        );
183        &self.words[index]
184    }
185}
dtrees_rs/bitsets/mod.rs

dtrees_rs/bitsets/
mod.rs
