1 //! Scalar bitsets.
2 
3 use core::mem::size_of;
4 use core::ops::{Add, BitAnd, BitOr, Not, Shl, Shr, Sub};
5 
6 /// A small bitset built on top of a single primitive integer type.
7 ///
8 /// # Example
9 ///
10 /// ```
11 /// use cranelift_bitset::ScalarBitSet;
12 ///
13 /// // Create a new bitset backed with a `u32`.
14 /// let mut bitset = ScalarBitSet::<u32>::new();
15 ///
16 /// // Bitsets are initially empty.
17 /// assert!(bitset.is_empty());
18 /// assert_eq!(bitset.len(), 0);
19 ///
20 /// // Insert into the bitset.
21 /// bitset.insert(4);
22 /// bitset.insert(5);
23 /// bitset.insert(6);
24 ///
25 /// // Now the bitset is not empty.
26 /// assert_eq!(bitset.len(), 3);
27 /// assert!(!bitset.is_empty());
28 /// assert!(bitset.contains(4));
29 /// assert!(bitset.contains(5));
30 /// assert!(bitset.contains(6));
31 ///
32 /// // Remove an element from the bitset.
33 /// let was_present = bitset.remove(6);
34 /// assert!(was_present);
35 /// assert!(!bitset.contains(6));
36 /// assert_eq!(bitset.len(), 2);
37 ///
38 /// // Can iterate over the elements in the set.
39 /// let elems: Vec<_> = bitset.iter().collect();
40 /// assert_eq!(elems, [4, 5]);
41 /// ```
42 #[derive(Clone, Copy, PartialEq, Eq)]
43 #[cfg_attr(
44     feature = "enable-serde",
45     derive(serde_derive::Serialize, serde_derive::Deserialize)
46 )]
47 pub struct ScalarBitSet<T>(pub T);
48 
49 impl<T> core::fmt::Debug for ScalarBitSet<T>
50 where
51     T: ScalarBitSetStorage,
52 {
fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result53     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
54         let mut s = f.debug_struct(core::any::type_name::<Self>());
55         for i in 0..Self::capacity() {
56             use alloc::string::ToString;
57             s.field(&i.to_string(), &self.contains(i));
58         }
59         s.finish()
60     }
61 }
62 
63 impl<T> Default for ScalarBitSet<T>
64 where
65     T: ScalarBitSetStorage,
66 {
67     #[inline]
default() -> Self68     fn default() -> Self {
69         Self::new()
70     }
71 }
72 
73 impl<T> ScalarBitSet<T>
74 where
75     T: ScalarBitSetStorage,
76 {
77     /// Create a new, empty bitset.
78     ///
79     /// # Example
80     ///
81     /// ```
82     /// use cranelift_bitset::ScalarBitSet;
83     ///
84     /// let bitset = ScalarBitSet::<u64>::new();
85     ///
86     /// assert!(bitset.is_empty());
87     /// ```
88     #[inline]
new() -> Self89     pub fn new() -> Self {
90         Self(T::from(0))
91     }
92 
93     /// Construct a bitset with the half-open range `[lo, hi)` inserted.
94     ///
95     /// # Example
96     ///
97     /// ```
98     /// use cranelift_bitset::ScalarBitSet;
99     ///
100     /// let bitset = ScalarBitSet::<u64>::from_range(3, 6);
101     ///
102     /// assert_eq!(bitset.len(), 3);
103     ///
104     /// assert!(bitset.contains(3));
105     /// assert!(bitset.contains(4));
106     /// assert!(bitset.contains(5));
107     /// ```
108     ///
109     /// # Panics
110     ///
111     /// Panics if `lo > hi` or if `hi > Self::capacity()`.
112     ///
113     /// ```should_panic
114     /// use cranelift_bitset::ScalarBitSet;
115     ///
116     /// // The lower bound may not be larger than the upper bound.
117     /// let bitset = ScalarBitSet::<u64>::from_range(6, 3);
118     /// ```
119     ///
120     /// ```should_panic
121     /// use cranelift_bitset::ScalarBitSet;
122     ///
123     /// // The bounds must fit within the backing scalar type.
124     /// let bitset = ScalarBitSet::<u64>::from_range(3, 69);
125     /// ```
126     #[inline]
from_range(lo: u8, hi: u8) -> Self127     pub fn from_range(lo: u8, hi: u8) -> Self {
128         assert!(lo <= hi);
129         assert!(hi <= Self::capacity());
130 
131         let one = T::from(1);
132 
133         // We can't just do (one << hi) - one here as the shift may overflow
134         let hi_rng = if hi >= 1 {
135             (one << (hi - 1)) + ((one << (hi - 1)) - one)
136         } else {
137             T::from(0)
138         };
139 
140         let lo_rng = (one << lo) - one;
141 
142         Self(hi_rng - lo_rng)
143     }
144 
145     /// The maximum number of bits that can be stored in this bitset.
146     ///
147     /// If you need more bits than this, use a
148     /// [`CompoundBitSet`][crate::CompoundBitSet] instead of a `ScalarBitSet`.
149     ///
150     /// # Example
151     ///
152     /// ```
153     /// use cranelift_bitset::ScalarBitSet;
154     ///
155     /// assert_eq!(ScalarBitSet::<u8>::capacity(), 8);
156     /// assert_eq!(ScalarBitSet::<u64>::capacity(), 64);
157     /// ```
158     #[inline]
capacity() -> u8159     pub fn capacity() -> u8 {
160         u8::try_from(size_of::<T>()).unwrap() * 8
161     }
162 
163     /// Get the number of elements in this set.
164     ///
165     /// # Example
166     ///
167     /// ```
168     /// use cranelift_bitset::ScalarBitSet;
169     ///
170     /// let mut bitset = ScalarBitSet::<u64>::new();
171     ///
172     /// assert_eq!(bitset.len(), 0);
173     ///
174     /// bitset.insert(24);
175     /// bitset.insert(13);
176     /// bitset.insert(36);
177     ///
178     /// assert_eq!(bitset.len(), 3);
179     /// ```
180     #[inline]
len(&self) -> u8181     pub fn len(&self) -> u8 {
182         self.0.count_ones()
183     }
184 
185     /// Is this bitset empty?
186     ///
187     /// # Example
188     ///
189     /// ```
190     /// use cranelift_bitset::ScalarBitSet;
191     ///
192     /// let mut bitset = ScalarBitSet::<u16>::new();
193     ///
194     /// assert!(bitset.is_empty());
195     ///
196     /// bitset.insert(10);
197     ///
198     /// assert!(!bitset.is_empty());
199     /// ```
200     #[inline]
is_empty(&self) -> bool201     pub fn is_empty(&self) -> bool {
202         self.0 == T::from(0)
203     }
204 
205     /// Check whether this bitset contains `i`.
206     ///
207     /// # Example
208     ///
209     /// ```
210     /// use cranelift_bitset::ScalarBitSet;
211     ///
212     /// let mut bitset = ScalarBitSet::<u8>::new();
213     ///
214     /// assert!(!bitset.contains(7));
215     ///
216     /// bitset.insert(7);
217     ///
218     /// assert!(bitset.contains(7));
219     /// ```
220     ///
221     /// # Panics
222     ///
223     /// Panics if `i` is greater than or equal to [`Self::capacity()`][Self::capacity].
224     ///
225     /// ```should_panic
226     /// use cranelift_bitset::ScalarBitSet;
227     ///
228     /// let mut bitset = ScalarBitSet::<u8>::new();
229     ///
230     /// // A `ScalarBitSet<u8>` can only hold the elements `0..=7`, so `8` is
231     /// // out of bounds and will trigger a panic.
232     /// bitset.contains(8);
233     /// ```
234     #[inline]
contains(&self, i: u8) -> bool235     pub fn contains(&self, i: u8) -> bool {
236         assert!(i < Self::capacity());
237         self.0 & (T::from(1) << i) != T::from(0)
238     }
239 
240     /// Insert `i` into this bitset.
241     ///
242     /// Returns whether the value was newly inserted. That is:
243     ///
244     /// * If the set did not previously contain `i` then `true` is returned.
245     ///
246     /// * If the set already contained `i` then `false` is returned.
247     ///
248     /// # Example
249     ///
250     /// ```
251     /// use cranelift_bitset::ScalarBitSet;
252     ///
253     /// let mut bitset = ScalarBitSet::<u8>::new();
254     ///
255     /// // When an element is inserted that was not already present in the set,
256     /// // then `true` is returned.
257     /// let is_new = bitset.insert(7);
258     /// assert!(is_new);
259     ///
260     /// // The element is now present in the set.
261     /// assert!(bitset.contains(7));
262     ///
263     /// // And when the element is already in the set, `false` is returned from
264     /// // `insert`.
265     /// let is_new = bitset.insert(7);
266     /// assert!(!is_new);
267     /// ```
268     ///
269     /// # Panics
270     ///
271     /// Panics if `i` is greater than or equal to [`Self::capacity()`][Self::capacity].
272     ///
273     /// ```should_panic
274     /// use cranelift_bitset::ScalarBitSet;
275     ///
276     /// let mut bitset = ScalarBitSet::<u32>::new();
277     ///
278     /// // A `ScalarBitSet<u32>` can only hold the elements `0..=31`, so `42` is
279     /// // out of bounds and will trigger a panic.
280     /// bitset.insert(42);
281     /// ```
282     #[inline]
insert(&mut self, i: u8) -> bool283     pub fn insert(&mut self, i: u8) -> bool {
284         let is_new = !self.contains(i);
285         self.0 = self.0 | (T::from(1) << i);
286         is_new
287     }
288 
289     /// Remove `i` from this bitset.
290     ///
291     /// Returns whether `i` was previously in this set or not.
292     ///
293     /// # Example
294     ///
295     /// ```
296     /// use cranelift_bitset::ScalarBitSet;
297     ///
298     /// let mut bitset = ScalarBitSet::<u128>::new();
299     ///
300     /// // Removing an element that was not present in the set returns `false`.
301     /// let was_present = bitset.remove(100);
302     /// assert!(!was_present);
303     ///
304     /// // And when the element was in the set, `true` is returned.
305     /// bitset.insert(100);
306     /// let was_present = bitset.remove(100);
307     /// assert!(was_present);
308     /// ```
309     ///
310     /// # Panics
311     ///
312     /// Panics if `i` is greater than or equal to [`Self::capacity()`][Self::capacity].
313     ///
314     /// ```should_panic
315     /// use cranelift_bitset::ScalarBitSet;
316     ///
317     /// let mut bitset = ScalarBitSet::<u16>::new();
318     ///
319     /// // A `ScalarBitSet<u16>` can only hold the elements `0..=15`, so `20` is
320     /// // out of bounds and will trigger a panic.
321     /// bitset.remove(20);
322     /// ```
323     #[inline]
remove(&mut self, i: u8) -> bool324     pub fn remove(&mut self, i: u8) -> bool {
325         let was_present = self.contains(i);
326         self.0 = self.0 & !(T::from(1) << i);
327         was_present
328     }
329 
330     /// Remove all entries from this bitset.
331     ///
332     /// # Example
333     ///
334     /// ```
335     /// use cranelift_bitset::ScalarBitSet;
336     ///
337     /// let mut bitset = ScalarBitSet::<u32>::new();
338     ///
339     /// bitset.insert(10);
340     /// bitset.insert(20);
341     /// bitset.insert(30);
342     ///
343     /// bitset.clear();
344     ///
345     /// assert!(bitset.is_empty());
346     /// ```
347     #[inline]
clear(&mut self)348     pub fn clear(&mut self) {
349         self.0 = T::from(0);
350     }
351 
352     /// Remove and return the smallest value in the bitset.
353     ///
354     /// # Example
355     ///
356     /// ```
357     /// use cranelift_bitset::ScalarBitSet;
358     ///
359     /// let mut bitset = ScalarBitSet::<u64>::new();
360     ///
361     /// bitset.insert(0);
362     /// bitset.insert(24);
363     /// bitset.insert(13);
364     /// bitset.insert(36);
365     ///
366     /// assert_eq!(bitset.pop_min(), Some(0));
367     /// assert_eq!(bitset.pop_min(), Some(13));
368     /// assert_eq!(bitset.pop_min(), Some(24));
369     /// assert_eq!(bitset.pop_min(), Some(36));
370     /// assert_eq!(bitset.pop_min(), None);
371     /// ```
372     #[inline]
pop_min(&mut self) -> Option<u8>373     pub fn pop_min(&mut self) -> Option<u8> {
374         let min = self.min()?;
375         self.remove(min);
376         Some(min)
377     }
378 
379     /// Remove and return the largest value in the bitset.
380     ///
381     /// # Example
382     ///
383     /// ```
384     /// use cranelift_bitset::ScalarBitSet;
385     ///
386     /// let mut bitset = ScalarBitSet::<u64>::new();
387     ///
388     /// bitset.insert(0);
389     /// bitset.insert(24);
390     /// bitset.insert(13);
391     /// bitset.insert(36);
392     ///
393     /// assert_eq!(bitset.pop_max(), Some(36));
394     /// assert_eq!(bitset.pop_max(), Some(24));
395     /// assert_eq!(bitset.pop_max(), Some(13));
396     /// assert_eq!(bitset.pop_max(), Some(0));
397     /// assert_eq!(bitset.pop_max(), None);
398     /// ```
399     #[inline]
pop_max(&mut self) -> Option<u8>400     pub fn pop_max(&mut self) -> Option<u8> {
401         let max = self.max()?;
402         self.remove(max);
403         Some(max)
404     }
405 
406     /// Return the smallest number contained in this bitset or `None` if this
407     /// bitset is empty.
408     ///
409     /// # Example
410     ///
411     /// ```
412     /// use cranelift_bitset::ScalarBitSet;
413     ///
414     /// let mut bitset = ScalarBitSet::<u64>::new();
415     ///
416     /// // When the bitset is empty, `min` returns `None`.
417     /// assert_eq!(bitset.min(), None);
418     ///
419     /// bitset.insert(28);
420     /// bitset.insert(1);
421     /// bitset.insert(63);
422     ///
423     /// // When the bitset is not empty, it returns the smallest element.
424     /// assert_eq!(bitset.min(), Some(1));
425     /// ```
426     #[inline]
min(&self) -> Option<u8>427     pub fn min(&self) -> Option<u8> {
428         if self.0 == T::from(0) {
429             None
430         } else {
431             Some(self.0.trailing_zeros())
432         }
433     }
434 
435     /// Return the largest number contained in the bitset or None if this bitset
436     /// is empty
437     ///
438     /// # Example
439     ///
440     /// ```
441     /// use cranelift_bitset::ScalarBitSet;
442     ///
443     /// let mut bitset = ScalarBitSet::<u64>::new();
444     ///
445     /// // When the bitset is empty, `max` returns `None`.
446     /// assert_eq!(bitset.max(), None);
447     ///
448     /// bitset.insert(0);
449     /// bitset.insert(36);
450     /// bitset.insert(49);
451     ///
452     /// // When the bitset is not empty, it returns the smallest element.
453     /// assert_eq!(bitset.max(), Some(49));
454     /// ```
455     #[inline]
max(&self) -> Option<u8>456     pub fn max(&self) -> Option<u8> {
457         if self.0 == T::from(0) {
458             None
459         } else {
460             let leading_zeroes = self.0.leading_zeros();
461             Some(Self::capacity() - leading_zeroes - 1)
462         }
463     }
464 
465     /// Iterate over the items in this set.
466     ///
467     /// Items are always yielded in sorted order.
468     ///
469     /// # Example
470     ///
471     /// ```
472     /// use cranelift_bitset::ScalarBitSet;
473     ///
474     /// let mut bitset = ScalarBitSet::<u64>::new();
475     ///
476     /// bitset.insert(19);
477     /// bitset.insert(3);
478     /// bitset.insert(63);
479     /// bitset.insert(0);
480     ///
481     /// assert_eq!(
482     ///     bitset.iter().collect::<Vec<_>>(),
483     ///     [0, 3, 19, 63],
484     /// );
485     /// ```
486     #[inline]
iter(self) -> Iter<T>487     pub fn iter(self) -> Iter<T> {
488         Iter { bitset: self }
489     }
490 }
491 
492 impl<T> IntoIterator for ScalarBitSet<T>
493 where
494     T: ScalarBitSetStorage,
495 {
496     type Item = u8;
497 
498     type IntoIter = Iter<T>;
499 
500     #[inline]
into_iter(self) -> Self::IntoIter501     fn into_iter(self) -> Self::IntoIter {
502         self.iter()
503     }
504 }
505 
506 impl<'a, T> IntoIterator for &'a ScalarBitSet<T>
507 where
508     T: ScalarBitSetStorage,
509 {
510     type Item = u8;
511 
512     type IntoIter = Iter<T>;
513 
514     #[inline]
into_iter(self) -> Self::IntoIter515     fn into_iter(self) -> Self::IntoIter {
516         self.iter()
517     }
518 }
519 
520 impl<T: ScalarBitSetStorage> From<T> for ScalarBitSet<T> {
from(bits: T) -> Self521     fn from(bits: T) -> Self {
522         Self(bits)
523     }
524 }
525 
526 /// A trait implemented by all integers that can be used as the backing storage
527 /// for a [`ScalarBitSet`].
528 ///
529 /// You shouldn't have to implement this yourself, it is already implemented for
530 /// `u{8,16,32,64,128}` and if you need more bits than that, then use
531 /// [`CompoundBitSet`][crate::CompoundBitSet] instead.
532 pub trait ScalarBitSetStorage:
533     Default
534     + From<u8>
535     + Shl<u8, Output = Self>
536     + Shr<u8, Output = Self>
537     + BitAnd<Output = Self>
538     + BitOr<Output = Self>
539     + Not<Output = Self>
540     + Sub<Output = Self>
541     + Add<Output = Self>
542     + PartialEq
543     + Copy
544 {
545     /// Count the number of leading zeros.
leading_zeros(self) -> u8546     fn leading_zeros(self) -> u8;
547 
548     /// Count the number of trailing zeros.
trailing_zeros(self) -> u8549     fn trailing_zeros(self) -> u8;
550 
551     /// Count the number of bits set in this integer.
count_ones(self) -> u8552     fn count_ones(self) -> u8;
553 }
554 
555 macro_rules! impl_storage {
556     ( $int:ty ) => {
557         impl ScalarBitSetStorage for $int {
558             #[inline]
559             fn leading_zeros(self) -> u8 {
560                 u8::try_from(self.leading_zeros()).unwrap()
561             }
562 
563             #[inline]
564             fn trailing_zeros(self) -> u8 {
565                 u8::try_from(self.trailing_zeros()).unwrap()
566             }
567 
568             #[inline]
569             fn count_ones(self) -> u8 {
570                 u8::try_from(self.count_ones()).unwrap()
571             }
572         }
573     };
574 }
575 
576 impl_storage!(u8);
577 impl_storage!(u16);
578 impl_storage!(u32);
579 impl_storage!(u64);
580 impl_storage!(u128);
581 impl_storage!(usize);
582 
583 /// An iterator over the elements in a [`ScalarBitSet`].
584 pub struct Iter<T> {
585     bitset: ScalarBitSet<T>,
586 }
587 
588 impl<T> Iterator for Iter<T>
589 where
590     T: ScalarBitSetStorage,
591 {
592     type Item = u8;
593 
594     #[inline]
next(&mut self) -> Option<u8>595     fn next(&mut self) -> Option<u8> {
596         self.bitset.pop_min()
597     }
598 }
599 
600 impl<T> DoubleEndedIterator for Iter<T>
601 where
602     T: ScalarBitSetStorage,
603 {
604     #[inline]
next_back(&mut self) -> Option<Self::Item>605     fn next_back(&mut self) -> Option<Self::Item> {
606         self.bitset.pop_max()
607     }
608 }
609 
610 impl<T> ExactSizeIterator for Iter<T>
611 where
612     T: ScalarBitSetStorage,
613 {
614     #[inline]
len(&self) -> usize615     fn len(&self) -> usize {
616         usize::from(self.bitset.len())
617     }
618 }
619 
620 #[cfg(feature = "arbitrary")]
621 impl<'a, T> arbitrary::Arbitrary<'a> for ScalarBitSet<T>
622 where
623     T: ScalarBitSetStorage + arbitrary::Arbitrary<'a>,
624 {
arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self>625     fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
626         T::arbitrary(u).map(Self)
627     }
628 }
629