Skip to main content

rolldown_utils/
index_bitset.rs

1use std::marker::PhantomData;
2
3use oxc_index::Idx;
4
5use crate::BitSet;
6
7#[derive(Clone, PartialEq, Eq, Hash)]
8pub struct IndexBitSet<I: Idx> {
9  inner: BitSet,
10  _marker: PhantomData<I>,
11}
12
13impl<I: Idx> IndexBitSet<I> {
14  const _ASSERT_IDX_FITS_U32: () = assert!(I::MAX <= u32::MAX as usize);
15
16  #[inline]
17  fn bit(idx: I) -> u32 {
18    // asserted by the const assertion above
19    #[expect(clippy::cast_possible_truncation)]
20    {
21      idx.index() as u32
22    }
23  }
24
25  pub fn new(capacity: usize) -> Self {
26    let () = Self::_ASSERT_IDX_FITS_U32;
27    debug_assert!(capacity <= I::MAX, "IndexBitSet capacity exceeds maximum index of type");
28    #[expect(clippy::cast_possible_truncation)]
29    let inner = { BitSet::new(capacity as u32) };
30    Self { inner, _marker: PhantomData }
31  }
32
33  pub fn has_bit(&self, idx: I) -> bool {
34    self.inner.has_bit(Self::bit(idx))
35  }
36
37  pub fn set_bit(&mut self, idx: I) -> bool {
38    self.inner.set_bit(Self::bit(idx))
39  }
40
41  pub fn clear_bit(&mut self, idx: I) {
42    self.inner.clear_bit(Self::bit(idx));
43  }
44
45  pub fn is_empty(&self) -> bool {
46    self.inner.is_empty()
47  }
48
49  pub fn bit_count(&self) -> u32 {
50    self.inner.bit_count()
51  }
52
53  pub fn index_of_one(&self) -> impl Iterator<Item = I> + '_ {
54    self.inner.index_of_one().map(|i| I::from_usize(i as usize))
55  }
56
57  pub fn union(&mut self, other: &Self) {
58    self.inner.union(&other.inner);
59  }
60}
61
62impl<I: Idx> Default for IndexBitSet<I> {
63  fn default() -> Self {
64    Self::new(0)
65  }
66}
67
68impl<I: Idx> Extend<I> for IndexBitSet<I> {
69  fn extend<T: IntoIterator<Item = I>>(&mut self, iter: T) {
70    self.inner.extend(iter.into_iter().map(Self::bit));
71  }
72}
73
74impl<I: Idx> IntoIterator for IndexBitSet<I> {
75  type Item = I;
76  type IntoIter = std::vec::IntoIter<I>;
77
78  fn into_iter(self) -> Self::IntoIter {
79    self.inner.into_iter().map(|i| I::from_usize(i as usize)).collect::<Vec<_>>().into_iter()
80  }
81}
82
83impl<I: Idx> FromIterator<I> for IndexBitSet<I> {
84  fn from_iter<T: IntoIterator<Item = I>>(iter: T) -> Self {
85    let () = Self::_ASSERT_IDX_FITS_U32;
86    let inner = iter.into_iter().map(Self::bit).collect::<BitSet>();
87    Self { inner, _marker: PhantomData }
88  }
89}
90
91impl<I: Idx> std::fmt::Debug for IndexBitSet<I> {
92  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
93    f.debug_set().entries(self.index_of_one()).finish()
94  }
95}
96
97#[cfg(test)]
98mod tests {
99  use super::*;
100
101  oxc_index::define_index_type! {
102    struct TestIdx = u32;
103  }
104
105  #[test]
106  fn basic_operations() {
107    let mut set = IndexBitSet::<TestIdx>::new(16);
108    assert!(set.is_empty());
109    assert_eq!(set.bit_count(), 0);
110
111    set.set_bit(TestIdx::from_usize(0));
112    assert!(set.has_bit(TestIdx::from_usize(0)));
113    assert_eq!(set.bit_count(), 1);
114
115    set.set_bit(TestIdx::from_usize(5));
116    set.set_bit(TestIdx::from_usize(10));
117    assert_eq!(set.bit_count(), 3);
118
119    set.clear_bit(TestIdx::from_usize(5));
120    assert!(!set.has_bit(TestIdx::from_usize(5)));
121    assert_eq!(set.bit_count(), 2);
122  }
123
124  #[test]
125  fn index_of_one_and_collect() {
126    let mut set = IndexBitSet::<TestIdx>::new(16);
127    set.set_bit(TestIdx::from_usize(1));
128    set.set_bit(TestIdx::from_usize(3));
129    set.set_bit(TestIdx::from_usize(7));
130
131    let items: Vec<TestIdx> = set.index_of_one().collect();
132    assert_eq!(items, vec![TestIdx::from_usize(1), TestIdx::from_usize(3), TestIdx::from_usize(7)]);
133
134    let collected: IndexBitSet<TestIdx> =
135      vec![TestIdx::from_usize(2), TestIdx::from_usize(4)].into_iter().collect();
136    assert!(collected.has_bit(TestIdx::from_usize(2)));
137    assert!(collected.has_bit(TestIdx::from_usize(4)));
138    assert_eq!(collected.bit_count(), 2);
139  }
140
141  #[test]
142  fn into_iterator() {
143    let mut set = IndexBitSet::<TestIdx>::new(8);
144    set.set_bit(TestIdx::from_usize(2));
145    set.set_bit(TestIdx::from_usize(6));
146
147    let items: Vec<TestIdx> = set.into_iter().collect();
148    assert_eq!(items, vec![TestIdx::from_usize(2), TestIdx::from_usize(6)]);
149  }
150
151  #[test]
152  fn extend_trait() {
153    let mut set = IndexBitSet::<TestIdx>::new(16);
154    set.set_bit(TestIdx::from_usize(0));
155    set.extend(vec![TestIdx::from_usize(3), TestIdx::from_usize(7)]);
156    assert_eq!(set.bit_count(), 3);
157    assert!(set.has_bit(TestIdx::from_usize(3)));
158  }
159
160  #[test]
161  fn union_sets() {
162    let mut a = IndexBitSet::<TestIdx>::new(16);
163    a.set_bit(TestIdx::from_usize(1));
164    a.set_bit(TestIdx::from_usize(3));
165
166    let mut b = IndexBitSet::<TestIdx>::new(16);
167    b.set_bit(TestIdx::from_usize(3));
168    b.set_bit(TestIdx::from_usize(5));
169
170    a.union(&b);
171    assert_eq!(a.bit_count(), 3);
172    assert!(a.has_bit(TestIdx::from_usize(1)));
173    assert!(a.has_bit(TestIdx::from_usize(3)));
174    assert!(a.has_bit(TestIdx::from_usize(5)));
175  }
176}