1 //! Compound bit sets.
2 
3 use crate::scalar::{self, ScalarBitSet, ScalarBitSetStorage};
4 use alloc::boxed::Box;
5 use core::{cmp, iter, mem};
6 
7 /// A large bit set backed by dynamically-sized storage.
8 ///
9 /// # Example
10 ///
11 /// ```
12 /// use cranelift_bitset::CompoundBitSet;
13 ///
14 /// // Create a new bitset.
15 /// let mut bitset = CompoundBitSet::new();
16 ///
17 /// // Bitsets are initially empty.
18 /// assert!(bitset.is_empty());
19 /// assert_eq!(bitset.len(), 0);
20 ///
21 /// // Insert into the bitset.
22 /// bitset.insert(444);
23 /// bitset.insert(555);
24 /// bitset.insert(666);
25 ///
26 /// // Now the bitset is not empty.
27 /// assert_eq!(bitset.len(), 3);
28 /// assert!(!bitset.is_empty());
29 /// assert!(bitset.contains(444));
30 /// assert!(bitset.contains(555));
31 /// assert!(bitset.contains(666));
32 ///
33 /// // Remove an element from the bitset.
34 /// let was_present = bitset.remove(666);
35 /// assert!(was_present);
36 /// assert!(!bitset.contains(666));
37 /// assert_eq!(bitset.len(), 2);
38 ///
39 /// // Can iterate over the elements in the set.
40 /// let elems: Vec<_> = bitset.iter().collect();
41 /// assert_eq!(elems, [444, 555]);
42 /// ```
43 #[derive(Clone, Default, PartialEq, Eq)]
44 #[cfg_attr(
45     feature = "enable-serde",
46     derive(serde_derive::Serialize, serde_derive::Deserialize)
47 )]
48 pub struct CompoundBitSet<T = usize> {
49     elems: Box<[ScalarBitSet<T>]>,
50     max: Option<u32>,
51 }
52 
53 impl core::fmt::Debug for CompoundBitSet {
54     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
55         write!(f, "CompoundBitSet ")?;
56         f.debug_set().entries(self.iter()).finish()
57     }
58 }
59 
60 impl CompoundBitSet {
61     /// Construct a new, empty bit set.
62     ///
63     /// # Example
64     ///
65     /// ```
66     /// use cranelift_bitset::CompoundBitSet;
67     ///
68     /// let bitset = CompoundBitSet::new();
69     ///
70     /// assert!(bitset.is_empty());
71     /// ```
72     #[inline]
73     pub fn new() -> Self {
74         CompoundBitSet::default()
75     }
76 }
77 
78 impl<T: ScalarBitSetStorage> CompoundBitSet<T> {
79     const BITS_PER_SCALAR: usize = mem::size_of::<T>() * 8;
80 
81     /// Construct a new, empty bit set with space reserved to store any element
82     /// `x` such that `x < capacity`.
83     ///
84     /// The actual capacity reserved may be greater than that requested.
85     ///
86     /// # Example
87     ///
88     /// ```
89     /// use cranelift_bitset::CompoundBitSet;
90     ///
91     /// let bitset = CompoundBitSet::<u32>::with_capacity(4096);
92     ///
93     /// assert!(bitset.is_empty());
94     /// assert!(bitset.capacity() >= 4096);
95     /// ```
96     #[inline]
97     pub fn with_capacity(capacity: usize) -> Self {
98         let mut bitset = Self::default();
99         bitset.ensure_capacity(capacity);
100         bitset
101     }
102 
103     /// Get the number of elements in this bitset.
104     ///
105     /// # Example
106     ///
107     /// ```
108     /// use cranelift_bitset::CompoundBitSet;
109     ///
110     /// let mut bitset = CompoundBitSet::new();
111     ///
112     /// assert_eq!(bitset.len(), 0);
113     ///
114     /// bitset.insert(24);
115     /// bitset.insert(130);
116     /// bitset.insert(3600);
117     ///
118     /// assert_eq!(bitset.len(), 3);
119     /// ```
120     #[inline]
121     pub fn len(&self) -> usize {
122         self.elems.iter().map(|sub| usize::from(sub.len())).sum()
123     }
124 
125     /// Get `n + 1` where `n` is the largest value that can be stored inside
126     /// this set without growing the backing storage.
127     ///
128     /// That is, this set can store any value `x` such that `x <
129     /// bitset.capacity()` without growing the backing storage.
130     ///
131     /// # Example
132     ///
133     /// ```
134     /// use cranelift_bitset::CompoundBitSet;
135     ///
136     /// let mut bitset = CompoundBitSet::new();
137     ///
138     /// // New bitsets have zero capacity -- they allocate lazily.
139     /// assert_eq!(bitset.capacity(), 0);
140     ///
141     /// // Insert into the bitset, growing its capacity.
142     /// bitset.insert(999);
143     ///
144     /// // The bitset must now have capacity for at least `999` elements,
145     /// // perhaps more.
146     /// assert!(bitset.capacity() >= 999);
147     ///```
148     pub fn capacity(&self) -> usize {
149         self.elems.len() * Self::BITS_PER_SCALAR
150     }
151 
152     /// Is this bitset empty?
153     ///
154     /// # Example
155     ///
156     /// ```
157     /// use cranelift_bitset::CompoundBitSet;
158     ///
159     /// let mut bitset = CompoundBitSet::new();
160     ///
161     /// assert!(bitset.is_empty());
162     ///
163     /// bitset.insert(1234);
164     ///
165     /// assert!(!bitset.is_empty());
166     /// ```
167     #[inline]
168     pub fn is_empty(&self) -> bool {
169         self.len() == 0
170     }
171 
172     /// Convert an element `i` into the `word` that can be used to index into
173     /// `self.elems` and the `bit` that can be tested in the
174     /// `ScalarBitSet<usize>` at `self.elems[word]`.
175     #[inline]
176     fn word_and_bit(i: usize) -> (usize, u8) {
177         let word = i / Self::BITS_PER_SCALAR;
178         let bit = i % Self::BITS_PER_SCALAR;
179         let bit = u8::try_from(bit).unwrap();
180         (word, bit)
181     }
182 
183     /// The opposite of `word_and_bit`: convert the pair of an index into
184     /// `self.elems` and associated bit index into a set element.
185     #[inline]
186     fn elem(word: usize, bit: u8) -> usize {
187         let bit = usize::from(bit);
188         debug_assert!(bit < Self::BITS_PER_SCALAR);
189         word * Self::BITS_PER_SCALAR + bit
190     }
191 
192     /// Is `i` contained in this bitset?
193     ///
194     /// # Example
195     ///
196     /// ```
197     /// use cranelift_bitset::CompoundBitSet;
198     ///
199     /// let mut bitset = CompoundBitSet::new();
200     ///
201     /// assert!(!bitset.contains(666));
202     ///
203     /// bitset.insert(666);
204     ///
205     /// assert!(bitset.contains(666));
206     /// ```
207     #[inline]
208     pub fn contains(&self, i: usize) -> bool {
209         let (word, bit) = Self::word_and_bit(i);
210         if word < self.elems.len() {
211             self.elems[word].contains(bit)
212         } else {
213             false
214         }
215     }
216 
217     /// Ensure there is space in this bitset for the values `0..n`, growing the
218     /// backing storage if necessary.
219     ///
220     /// After calling `bitset.ensure_capacity(n)`, inserting any element `i`
221     /// where `i < n` is guaranteed to succeed without growing the bitset's
222     /// backing storage.
223     ///
224     /// # Example
225     ///
226     /// ```
227     /// # use cranelift_bitset::CompoundBitSet;
228     /// # let mut bitset = CompoundBitSet::new();
229     /// // We are going to do a series of inserts where `1000` will be the
230     /// // maximum value inserted. Make sure that our bitset has capacity for
231     /// // these elements once up front, to avoid growing the backing storage
232     /// // multiple times incrementally.
233     /// bitset.ensure_capacity(1001);
234     ///
235     /// for i in 0..=1000 {
236     ///     if i % 2 == 0 {
237     ///         // Inserting this value should not require growing the backing
238     ///         // storage.
239     ///         assert!(bitset.capacity() > i);
240     ///         bitset.insert(i);
241     ///     }
242     /// }
243     /// ```
244     #[inline]
245     pub fn ensure_capacity(&mut self, n: usize) {
246         // Subtract one from the capacity to get the maximum bit that we might
247         // set. If `n` is 0 then nothing need be done as no capacity needs to be
248         // allocated.
249         let (word, _bit) = Self::word_and_bit(match n.checked_sub(1) {
250             None => return,
251             Some(n) => n,
252         });
253         if word >= self.elems.len() {
254             assert!(word < usize::try_from(isize::MAX).unwrap());
255 
256             let delta = word - self.elems.len();
257             let to_grow = delta + 1;
258 
259             // Amortize the cost of growing.
260             let to_grow = cmp::max(to_grow, self.elems.len() * 2);
261             // Don't make ridiculously small allocations.
262             let to_grow = cmp::max(to_grow, 4);
263 
264             let new_elems = self
265                 .elems
266                 .iter()
267                 .copied()
268                 .chain(iter::repeat(ScalarBitSet::new()).take(to_grow))
269                 .collect::<Box<[_]>>();
270             self.elems = new_elems;
271         }
272     }
273 
274     /// Insert `i` into this bitset.
275     ///
276     /// Returns whether the value was newly inserted. That is:
277     ///
278     /// * If the set did not previously contain `i` then `true` is returned.
279     ///
280     /// * If the set already contained `i` then `false` is returned.
281     ///
282     /// # Example
283     ///
284     /// ```
285     /// use cranelift_bitset::CompoundBitSet;
286     ///
287     /// let mut bitset = CompoundBitSet::new();
288     ///
289     /// // When an element is inserted that was not already present in the set,
290     /// // then `true` is returned.
291     /// let is_new = bitset.insert(1234);
292     /// assert!(is_new);
293     ///
294     /// // The element is now present in the set.
295     /// assert!(bitset.contains(1234));
296     ///
297     /// // And when the element is already in the set, `false` is returned from
298     /// // `insert`.
299     /// let is_new = bitset.insert(1234);
300     /// assert!(!is_new);
301     /// ```
302     #[inline]
303     pub fn insert(&mut self, i: usize) -> bool {
304         self.ensure_capacity(i + 1);
305 
306         let (word, bit) = Self::word_and_bit(i);
307         let is_new = self.elems[word].insert(bit);
308 
309         let i = u32::try_from(i).unwrap();
310         self.max = self.max.map(|max| cmp::max(max, i)).or(Some(i));
311 
312         is_new
313     }
314 
315     /// Remove `i` from this bitset.
316     ///
317     /// Returns whether `i` was previously in this set or not.
318     ///
319     /// # Example
320     ///
321     /// ```
322     /// use cranelift_bitset::CompoundBitSet;
323     ///
324     /// let mut bitset = CompoundBitSet::new();
325     ///
326     /// // Removing an element that was not present in the set returns `false`.
327     /// let was_present = bitset.remove(456);
328     /// assert!(!was_present);
329     ///
330     /// // And when the element was in the set, `true` is returned.
331     /// bitset.insert(456);
332     /// let was_present = bitset.remove(456);
333     /// assert!(was_present);
334     /// ```
335     #[inline]
336     pub fn remove(&mut self, i: usize) -> bool {
337         let (word, bit) = Self::word_and_bit(i);
338         if word < self.elems.len() {
339             let sub = &mut self.elems[word];
340             let was_present = sub.remove(bit);
341             if was_present && self.max() == Some(i) {
342                 self.update_max(word);
343             }
344             was_present
345         } else {
346             false
347         }
348     }
349 
350     /// Update the `self.max` field, based on the old word index of `self.max`.
351     fn update_max(&mut self, word_of_old_max: usize) {
352         self.max = self.elems[0..word_of_old_max + 1]
353             .iter()
354             .enumerate()
355             .rev()
356             .filter_map(|(word, sub)| {
357                 let bit = sub.max()?;
358                 Some(u32::try_from(Self::elem(word, bit)).unwrap())
359             })
360             .next();
361     }
362 
363     /// Get the largest value in this set, or `None` if this set is empty.
364     ///
365     /// # Example
366     ///
367     /// ```
368     /// use cranelift_bitset::CompoundBitSet;
369     ///
370     /// let mut bitset = CompoundBitSet::new();
371     ///
372     /// // Returns `None` if the bitset is empty.
373     /// assert!(bitset.max().is_none());
374     ///
375     /// bitset.insert(123);
376     /// bitset.insert(987);
377     /// bitset.insert(999);
378     ///
379     /// // Otherwise, it returns the largest value in the set.
380     /// assert_eq!(bitset.max(), Some(999));
381     /// ```
382     #[inline]
383     pub fn max(&self) -> Option<usize> {
384         self.max.map(|m| usize::try_from(m).unwrap())
385     }
386 
387     /// Removes and returns the largest value in this set.
388     ///
389     /// Returns `None` if this set is empty.
390     ///
391     /// # Example
392     ///
393     /// ```
394     /// use cranelift_bitset::CompoundBitSet;
395     ///
396     /// let mut bitset = CompoundBitSet::new();
397     ///
398     /// bitset.insert(111);
399     /// bitset.insert(222);
400     /// bitset.insert(333);
401     /// bitset.insert(444);
402     /// bitset.insert(555);
403     ///
404     /// assert_eq!(bitset.pop(), Some(555));
405     /// assert_eq!(bitset.pop(), Some(444));
406     /// assert_eq!(bitset.pop(), Some(333));
407     /// assert_eq!(bitset.pop(), Some(222));
408     /// assert_eq!(bitset.pop(), Some(111));
409     /// assert_eq!(bitset.pop(), None);
410     /// ```
411     #[inline]
412     pub fn pop(&mut self) -> Option<usize> {
413         let max = self.max()?;
414         self.remove(max);
415         Some(max)
416     }
417 
418     /// Remove all elements from this bitset.
419     ///
420     /// # Example
421     ///
422     /// ```
423     /// use cranelift_bitset::CompoundBitSet;
424     ///
425     /// let mut bitset = CompoundBitSet::new();
426     ///
427     /// bitset.insert(100);
428     /// bitset.insert(200);
429     /// bitset.insert(300);
430     ///
431     /// bitset.clear();
432     ///
433     /// assert!(bitset.is_empty());
434     /// ```
435     #[inline]
436     pub fn clear(&mut self) {
437         let max = match self.max() {
438             Some(max) => max,
439             None => return,
440         };
441         let (word, _bit) = Self::word_and_bit(max);
442         debug_assert!(self.elems[word + 1..].iter().all(|sub| sub.is_empty()));
443         for sub in &mut self.elems[..=word] {
444             *sub = ScalarBitSet::new();
445         }
446         self.max = None;
447     }
448 
449     /// Iterate over the elements in this bitset.
450     ///
451     /// The elements are always yielded in sorted order.
452     ///
453     /// # Example
454     ///
455     /// ```
456     /// use cranelift_bitset::CompoundBitSet;
457     ///
458     /// let mut bitset = CompoundBitSet::new();
459     ///
460     /// bitset.insert(0);
461     /// bitset.insert(4096);
462     /// bitset.insert(123);
463     /// bitset.insert(456);
464     /// bitset.insert(789);
465     ///
466     /// assert_eq!(
467     ///     bitset.iter().collect::<Vec<_>>(),
468     ///     [0, 123, 456, 789, 4096],
469     /// );
470     /// ```
471     #[inline]
472     pub fn iter(&self) -> Iter<'_, T> {
473         Iter {
474             bitset: self,
475             word: 0,
476             sub: None,
477         }
478     }
479 
480     /// Returns an iterator over the words of this bit-set or the in-memory
481     /// representation of the bit set.
482     ///
483     /// # Example
484     ///
485     /// ```
486     /// use cranelift_bitset::{CompoundBitSet, ScalarBitSet};
487     ///
488     /// let mut bitset = CompoundBitSet::<u32>::default();
489     ///
490     /// assert_eq!(
491     ///     bitset.iter_scalars().collect::<Vec<_>>(),
492     ///     [],
493     /// );
494     ///
495     /// bitset.insert(0);
496     ///
497     /// assert_eq!(
498     ///     bitset.iter_scalars().collect::<Vec<_>>(),
499     ///     [ScalarBitSet(0x1)],
500     /// );
501     ///
502     /// bitset.insert(1);
503     ///
504     /// assert_eq!(
505     ///     bitset.iter_scalars().collect::<Vec<_>>(),
506     ///     [ScalarBitSet(0x3)],
507     /// );
508     ///
509     /// bitset.insert(32);
510     ///
511     /// assert_eq!(
512     ///     bitset.iter_scalars().collect::<Vec<_>>(),
513     ///     [ScalarBitSet(0x3), ScalarBitSet(0x1)],
514     /// );
515     /// ```
516     pub fn iter_scalars(&self) -> impl Iterator<Item = ScalarBitSet<T>> + '_ {
517         let nwords = match self.max {
518             Some(n) => 1 + (n as usize / Self::BITS_PER_SCALAR),
519             None => 0,
520         };
521         self.elems.iter().copied().take(nwords)
522     }
523 }
524 
525 impl<'a, T: ScalarBitSetStorage> IntoIterator for &'a CompoundBitSet<T> {
526     type Item = usize;
527 
528     type IntoIter = Iter<'a, T>;
529 
530     #[inline]
531     fn into_iter(self) -> Self::IntoIter {
532         self.iter()
533     }
534 }
535 
536 /// An iterator over the elements in a [`CompoundBitSet`].
537 pub struct Iter<'a, T = usize> {
538     bitset: &'a CompoundBitSet<T>,
539     word: usize,
540     sub: Option<scalar::Iter<T>>,
541 }
542 
543 impl<T: ScalarBitSetStorage> Iterator for Iter<'_, T> {
544     type Item = usize;
545 
546     #[inline]
547     fn next(&mut self) -> Option<usize> {
548         loop {
549             if let Some(sub) = &mut self.sub {
550                 if let Some(bit) = sub.next() {
551                     return Some(CompoundBitSet::<T>::elem(self.word, bit));
552                 } else {
553                     self.word += 1;
554                 }
555             }
556 
557             self.sub = Some(self.bitset.elems.get(self.word)?.iter());
558         }
559     }
560 }
561 
562 #[cfg(test)]
563 mod tests {
564     use super::*;
565 
566     #[test]
567     fn zero_capacity_no_allocs() {
568         let set = CompoundBitSet::<u32>::with_capacity(0);
569         assert_eq!(set.capacity(), 0);
570         let set = CompoundBitSet::new();
571         assert_eq!(set.capacity(), 0);
572     }
573 }
574