dtrees_rs/cover/
reversible_cover.rs

1//! A reversible sparse bitset of instances.
2
3use crate::bitsets::Bitset;
4use search_trail::{
5    ReversibleU64, ReversibleUsize, SaveAndRestore, StateManager, U64Manager, UsizeManager,
6};
7use std::cmp::Ordering;
8use std::ops::Sub;
9
10/// A set of instances that supports intersection and constant-time undo.
11///
12/// Words that become zero are moved out of the active part of
13/// `non_zero_words`, so later operations only visit non-empty words.
14pub struct SparseBitset {
15    words: Vec<ReversibleU64>,
16    non_zero_words: Vec<usize>,
17    nb_non_zero: ReversibleUsize,
18
19    state_manager: StateManager,
20}
21
22/// A non-reversible snapshot of a [`SparseBitset`]: every word, with the
23/// inactive ones as 0.
24#[derive(Debug)]
25pub struct ShallowBitset {
26    words: Vec<u64>,
27}
28
29/// Sizes of the two set differences between a [`SparseBitset`] and a
30/// [`ShallowBitset`].
31#[derive(Default)]
32pub struct Difference {
33    pub(crate) in_count: usize,
34    pub(crate) out_count: usize,
35}
36
37impl SparseBitset {
38    /// The full set `0..n`.
39    pub fn new(n: usize) -> Self {
40        let mut state_manager = StateManager::default();
41
42        let nb_words = n.div_ceil(64);
43        let mut words = Vec::with_capacity(nb_words);
44        for _ in 0..nb_words {
45            words.push(state_manager.manage_u64(u64::MAX));
46        }
47
48        let mask = if n % 64 == 0 {
49            u64::MAX
50        } else {
51            (1u64 << (n % 64)) - 1
52        };
53        if let Some(last) = words.last_mut() {
54            state_manager.set_u64(*last, mask);
55        }
56        let non_zero_words = (0..nb_words).collect();
57        let nb_non_zero = state_manager.manage_usize(nb_words);
58
59        state_manager.save_state();
60
61        Self {
62            words,
63            non_zero_words,
64            nb_non_zero,
65            state_manager,
66        }
67    }
68
69    /// Number of elements.
70    pub fn count(&self) -> usize {
71        let mut count = 0;
72        let nb_non_zero = self.state_manager.get_usize(self.nb_non_zero);
73        for i in (0..nb_non_zero).rev() {
74            count += self
75                .state_manager
76                .get_u64(self.words[self.non_zero_words[i]])
77                .count_ones();
78        }
79        count as usize
80    }
81
82    /// Whether the set is empty.
83    pub fn is_empty(&self) -> bool {
84        self.state_manager.get_usize(self.nb_non_zero) == 0
85    }
86
87    /// Intersects with `other` (with its complement when `invert`), saving the
88    /// current state for [`Self::restore`]. Returns the new size.
89    pub fn intersect_with(&mut self, other: &Bitset, invert: bool) -> usize {
90        self.state_manager.save_state();
91
92        let mut size = self.state_manager.get_usize(self.nb_non_zero);
93        let mut count = 0;
94        for i in (0..size).rev() {
95            let idx = self.non_zero_words[i];
96            let intersect = self.state_manager.get_u64(self.words[idx])
97                & if invert { !other[idx] } else { other[idx] };
98            if intersect == 0 {
99                size -= 1;
100                self.non_zero_words[i] = self.non_zero_words[size];
101                self.non_zero_words[size] = idx;
102            } else {
103                self.state_manager.set_u64(self.words[idx], intersect);
104                count += intersect.count_ones();
105            }
106        }
107        self.state_manager.set_usize(self.nb_non_zero, size);
108
109        count as usize
110    }
111
112    /// Size of the intersection with `other` (or its complement), without
113    /// changing the set.
114    pub fn count_intersect_with(&self, other: &Bitset, invert: bool) -> usize {
115        let size = self.state_manager.get_usize(self.nb_non_zero);
116        let mut count = 0;
117        for i in (0..size).rev() {
118            let idx = self.non_zero_words[i];
119            let intersect = self.state_manager.get_u64(self.words[idx])
120                & if invert { !other[idx] } else { other[idx] };
121            if intersect != 0 {
122                count += intersect.count_ones();
123            }
124        }
125        count as usize
126    }
127
128    /// Size of the intersection with each of `others`.
129    pub fn count_intersect_with_many(&self, others: &[Bitset]) -> Vec<usize> {
130        let mut counts = vec![0; others.len()];
131        let size = self.state_manager.get_usize(self.nb_non_zero);
132        for i in (0..size).rev() {
133            let idx = self.non_zero_words[i];
134            let word = self.state_manager.get_u64(self.words[idx]);
135            for (bid, other) in others.iter().enumerate() {
136                counts[bid] += (word & other[idx]).count_ones() as usize
137            }
138        }
139        counts
140    }
141
142    /// The elements of the set.
143    pub fn to_vec(&self) -> Vec<usize> {
144        let mut result = Vec::new();
145        let nb_non_zero = self.state_manager.get_usize(self.nb_non_zero);
146
147        for i in 0..nb_non_zero {
148            let word_idx = self.non_zero_words[i];
149            let mut word = self.state_manager.get_u64(self.words[word_idx]);
150            let base_idx = word_idx * 64;
151
152            while word != 0 {
153                let bit_pos = word.trailing_zeros() as usize;
154                result.push(base_idx + bit_pos);
155                word &= word - 1;
156            }
157        }
158        result
159    }
160
161    /// Undoes the last intersection.
162    #[inline]
163    pub fn restore(&mut self) {
164        self.state_manager.restore_state();
165    }
166
167    /// Every word of the set, with the inactive ones as 0.
168    fn active_words(&self) -> Vec<u64> {
169        let mut words = Vec::new();
170        self.write_active_words(&mut words);
171        words
172    }
173
174    /// Writes every word of the set into `out`, with the inactive ones as 0.
175    ///
176    /// A word that becomes empty is moved out of the active part of
177    /// `non_zero_words` without being overwritten, so its stored value is
178    /// stale and must not be read.
179    pub fn write_active_words(&self, out: &mut Vec<u64>) {
180        out.clear();
181        out.resize(self.words.len(), 0);
182        let nb_non_zero = self.state_manager.get_usize(self.nb_non_zero);
183        for &idx in &self.non_zero_words[..nb_non_zero] {
184            out[idx] = self.state_manager.get_u64(self.words[idx]);
185        }
186    }
187}
188
189impl From<&SparseBitset> for ShallowBitset {
190    fn from(val: &SparseBitset) -> Self {
191        ShallowBitset {
192            words: val.active_words(),
193        }
194    }
195}
196
197/// `in_count` counts the elements of `self` missing from `rhs`, and
198/// `out_count` the elements of `rhs` missing from `self`.
199impl Sub<&ShallowBitset> for &SparseBitset {
200    type Output = Difference;
201    fn sub(self, rhs: &ShallowBitset) -> Self::Output {
202        let words = self.active_words();
203        let in_count = words
204            .iter()
205            .zip(&rhs.words)
206            .map(|(&own, &other)| (own & !other).count_ones() as usize)
207            .sum();
208        let out_count = words
209            .iter()
210            .zip(&rhs.words)
211            .map(|(&own, &other)| (other & !own).count_ones() as usize)
212            .sum();
213        Difference {
214            in_count,
215            out_count,
216        }
217    }
218}
219
220impl Sub<ShallowBitset> for SparseBitset {
221    type Output = Difference;
222
223    fn sub(self, rhs: ShallowBitset) -> Self::Output {
224        &self - &rhs
225    }
226}
227impl Sub<&ShallowBitset> for SparseBitset {
228    type Output = Difference;
229
230    fn sub(self, rhs: &ShallowBitset) -> Self::Output {
231        &self - rhs
232    }
233}
234
235impl Sub<ShallowBitset> for &SparseBitset {
236    type Output = Difference;
237
238    fn sub(self, rhs: ShallowBitset) -> Self::Output {
239        self - &rhs
240    }
241}
242
243impl Ord for Difference {
244    fn cmp(&self, other: &Self) -> Ordering {
245        (self.in_count + self.out_count).cmp(&(other.in_count + other.out_count))
246    }
247}
248
249impl PartialOrd for Difference {
250    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
251        Some(self.cmp(other))
252    }
253}
254
255impl PartialEq for Difference {
256    fn eq(&self, other: &Self) -> bool {
257        (self.in_count + self.out_count) == (other.in_count + other.out_count)
258    }
259}
260
261impl Eq for Difference {}
262
263#[cfg(test)]
264mod sparse_test {
265    use crate::bitsets::{BitCollection, Bitset, BitsetInit};
266    use crate::cover::reversible_cover::{ShallowBitset, SparseBitset};
267    use search_trail::UsizeManager;
268
269    #[test]
270    fn create() {
271        let mut cover = SparseBitset::new(10);
272
273        let mut feature = Bitset::new(BitsetInit::Empty(10));
274
275        feature.set(9);
276
277        println!("{:?}", cover.to_vec());
278        assert_eq!(cover.count(), 10);
279
280        cover.intersect_with(&feature, false);
281        println!("{:?}", cover.to_vec());
282        println!("xx {:?}", cover.state_manager.get_usize(cover.nb_non_zero));
283
284        let shallow_cover: ShallowBitset = (&cover).into();
285        println!("Shaloow {:?}", shallow_cover);
286        cover.restore();
287        println!("{:?}", cover.to_vec());
288
289        let _ = &cover - shallow_cover;
290    }
291
292    #[test]
293    fn a_word_emptied_by_an_intersection_counts_as_empty() {
294        // Rows 64..128 only: the intersection empties word 0, which keeps its
295        // old value in storage but is no longer active.
296        let mut feature = Bitset::new(BitsetInit::Empty(128));
297        for row in 64..128 {
298            feature.set(row);
299        }
300        let mut cover = SparseBitset::new(128);
301        let full: ShallowBitset = (&cover).into();
302        cover.intersect_with(&feature, false);
303
304        let difference = &cover - &full;
305        assert_eq!(difference.in_count, 0);
306        assert_eq!(difference.out_count, 64, "rows 0..64 are only in `full`");
307
308        let half: ShallowBitset = (&cover).into();
309        cover.restore();
310        let difference = &cover - &half;
311        assert_eq!(difference.in_count, 64, "rows 0..64 are only in `cover`");
312        assert_eq!(difference.out_count, 0);
313    }
314}
dtrees_rs/cover/reversible_cover.rs

dtrees_rs/cover/
reversible_cover.rs
