cubecl_ir/
enum_set.rs

1//! A compact set of enum variants, stored as a bitmask.
2//!
3//! [`EnumSet`] is a drop-in stand-in for the `enumset` crate's type of the same name, kept in the
4//! tree so the derive can share `cubecl-macros-internal`'s `darling`/`syn` rather than pulling a
5//! second, older copy of both through `enumset_derive`.
6
7use core::fmt::{self, Debug};
8use core::hash::{Hash, Hasher};
9use core::marker::PhantomData;
10use core::ops::{
11    BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not, Sub, SubAssign,
12};
13
14/// An enum whose variants can be collected into an [`EnumSet`].
15///
16/// Derive it with `#[derive(EnumSetType)]`, alongside `Clone`, `Copy`, `PartialEq` and `Eq`.
17///
18/// The set is backed by a `u64`, so an enum may have at most 64 variants, and it must be a plain
19/// C-like enum with no fields.
20pub trait EnumSetType: Copy + Eq + 'static {
21    /// Number of variants in the enum.
22    const VARIANTS: u32;
23
24    /// Index of the bit representing this variant.
25    fn to_bit(self) -> u32;
26
27    /// The variant a bit index stands for.
28    ///
29    /// # Panics
30    /// If `bit` is not below [`Self::VARIANTS`].
31    fn from_bit(bit: u32) -> Self;
32}
33
34/// A set of enum variants, stored as a bitmask.
35pub struct EnumSet<T: EnumSetType> {
36    bits: u64,
37    _ty: PhantomData<T>,
38}
39
40impl<T: EnumSetType> EnumSet<T> {
41    /// Mask with every valid bit for `T` set.
42    const MASK: u64 = if T::VARIANTS >= 64 {
43        u64::MAX
44    } else {
45        (1u64 << T::VARIANTS) - 1
46    };
47
48    /// The empty set.
49    pub const fn empty() -> Self {
50        Self {
51            bits: 0,
52            _ty: PhantomData,
53        }
54    }
55
56    /// The set of every variant of `T`.
57    pub const fn all() -> Self {
58        Self {
59            bits: Self::MASK,
60            _ty: PhantomData,
61        }
62    }
63
64    /// A set holding only `value`.
65    pub fn only(value: T) -> Self {
66        Self {
67            bits: 1 << value.to_bit(),
68            _ty: PhantomData,
69        }
70    }
71
72    /// The raw bitmask.
73    pub const fn as_u64(self) -> u64 {
74        self.bits
75    }
76
77    /// Build a set from a raw bitmask, discarding bits that name no variant.
78    pub const fn from_u64_truncated(bits: u64) -> Self {
79        Self {
80            bits: bits & Self::MASK,
81            _ty: PhantomData,
82        }
83    }
84
85    /// Number of variants in the set.
86    pub const fn len(self) -> usize {
87        self.bits.count_ones() as usize
88    }
89
90    /// Whether the set holds no variants.
91    pub const fn is_empty(self) -> bool {
92        self.bits == 0
93    }
94
95    /// Whether `value` is in the set.
96    pub fn contains(self, value: T) -> bool {
97        self.bits & (1 << value.to_bit()) != 0
98    }
99
100    /// Add `value`, returning whether it was newly added.
101    pub fn insert(&mut self, value: T) -> bool {
102        let added = !self.contains(value);
103        self.bits |= 1 << value.to_bit();
104        added
105    }
106
107    /// Remove `value`, returning whether it had been present.
108    pub fn remove(&mut self, value: T) -> bool {
109        let present = self.contains(value);
110        self.bits &= !(1 << value.to_bit());
111        present
112    }
113
114    /// Remove every variant.
115    pub fn clear(&mut self) {
116        self.bits = 0;
117    }
118
119    /// Variants in either set.
120    pub const fn union(self, other: Self) -> Self {
121        Self {
122            bits: self.bits | other.bits,
123            _ty: PhantomData,
124        }
125    }
126
127    /// Variants in both sets.
128    pub const fn intersection(self, other: Self) -> Self {
129        Self {
130            bits: self.bits & other.bits,
131            _ty: PhantomData,
132        }
133    }
134
135    /// Variants in `self` but not `other`.
136    pub const fn difference(self, other: Self) -> Self {
137        Self {
138            bits: self.bits & !other.bits,
139            _ty: PhantomData,
140        }
141    }
142
143    /// Variants in exactly one of the two sets.
144    pub const fn symmetrical_difference(self, other: Self) -> Self {
145        Self {
146            bits: self.bits ^ other.bits,
147            _ty: PhantomData,
148        }
149    }
150
151    /// Every variant not in the set.
152    pub const fn complement(self) -> Self {
153        Self {
154            bits: !self.bits & Self::MASK,
155            _ty: PhantomData,
156        }
157    }
158
159    /// Whether every variant of `self` is in `other`.
160    pub const fn is_subset(self, other: Self) -> bool {
161        self.bits & other.bits == self.bits
162    }
163
164    /// Whether every variant of `other` is in `self`.
165    pub const fn is_superset(self, other: Self) -> bool {
166        other.is_subset(self)
167    }
168
169    /// Whether the two sets share no variants.
170    pub const fn is_disjoint(self, other: Self) -> bool {
171        self.bits & other.bits == 0
172    }
173
174    /// Iterate the variants in the set, in declaration order.
175    pub fn iter(self) -> EnumSetIter<T> {
176        EnumSetIter {
177            bits: self.bits,
178            _ty: PhantomData,
179        }
180    }
181}
182
183impl<T: EnumSetType> Clone for EnumSet<T> {
184    fn clone(&self) -> Self {
185        *self
186    }
187}
188
189impl<T: EnumSetType> Copy for EnumSet<T> {}
190
191impl<T: EnumSetType> Default for EnumSet<T> {
192    fn default() -> Self {
193        Self::empty()
194    }
195}
196
197impl<T: EnumSetType> PartialEq for EnumSet<T> {
198    fn eq(&self, other: &Self) -> bool {
199        self.bits == other.bits
200    }
201}
202
203impl<T: EnumSetType> Eq for EnumSet<T> {}
204
205impl<T: EnumSetType> PartialOrd for EnumSet<T> {
206    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
207        Some(self.cmp(other))
208    }
209}
210
211impl<T: EnumSetType> Ord for EnumSet<T> {
212    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
213        self.bits.cmp(&other.bits)
214    }
215}
216
217impl<T: EnumSetType> Hash for EnumSet<T> {
218    fn hash<H: Hasher>(&self, state: &mut H) {
219        self.bits.hash(state);
220    }
221}
222
223impl<T: EnumSetType + Debug> Debug for EnumSet<T> {
224    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
225        write!(f, "EnumSet(")?;
226        for (i, value) in self.iter().enumerate() {
227            if i > 0 {
228                write!(f, " | ")?;
229            }
230            write!(f, "{value:?}")?;
231        }
232        write!(f, ")")
233    }
234}
235
236impl<T: EnumSetType> From<T> for EnumSet<T> {
237    fn from(value: T) -> Self {
238        Self::only(value)
239    }
240}
241
242impl<T: EnumSetType> FromIterator<T> for EnumSet<T> {
243    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
244        let mut set = Self::empty();
245        set.extend(iter);
246        set
247    }
248}
249
250impl<T: EnumSetType> Extend<T> for EnumSet<T> {
251    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
252        for value in iter {
253            self.insert(value);
254        }
255    }
256}
257
258impl<T: EnumSetType> IntoIterator for EnumSet<T> {
259    type Item = T;
260    type IntoIter = EnumSetIter<T>;
261
262    fn into_iter(self) -> Self::IntoIter {
263        self.iter()
264    }
265}
266
267/// Iterator over the variants of an [`EnumSet`], in declaration order.
268pub struct EnumSetIter<T: EnumSetType> {
269    bits: u64,
270    _ty: PhantomData<T>,
271}
272
273impl<T: EnumSetType> Iterator for EnumSetIter<T> {
274    type Item = T;
275
276    fn next(&mut self) -> Option<T> {
277        if self.bits == 0 {
278            return None;
279        }
280        let bit = self.bits.trailing_zeros();
281        self.bits &= self.bits - 1;
282        Some(T::from_bit(bit))
283    }
284
285    fn size_hint(&self) -> (usize, Option<usize>) {
286        let len = self.bits.count_ones() as usize;
287        (len, Some(len))
288    }
289}
290
291impl<T: EnumSetType> ExactSizeIterator for EnumSetIter<T> {}
292
293/// Generates the set/set and set/variant halves of a binary operator.
294macro_rules! impl_op {
295    ($op: ident, $method: ident, $assign_op: ident, $assign_method: ident, $call: ident) => {
296        impl<T: EnumSetType> $op<EnumSet<T>> for EnumSet<T> {
297            type Output = EnumSet<T>;
298
299            fn $method(self, other: EnumSet<T>) -> EnumSet<T> {
300                self.$call(other)
301            }
302        }
303
304        impl<T: EnumSetType> $op<T> for EnumSet<T> {
305            type Output = EnumSet<T>;
306
307            fn $method(self, other: T) -> EnumSet<T> {
308                self.$call(EnumSet::only(other))
309            }
310        }
311
312        impl<T: EnumSetType> $assign_op<EnumSet<T>> for EnumSet<T> {
313            fn $assign_method(&mut self, other: EnumSet<T>) {
314                *self = self.$call(other);
315            }
316        }
317
318        impl<T: EnumSetType> $assign_op<T> for EnumSet<T> {
319            fn $assign_method(&mut self, other: T) {
320                *self = self.$call(EnumSet::only(other));
321            }
322        }
323    };
324}
325
326impl_op!(BitOr, bitor, BitOrAssign, bitor_assign, union);
327impl_op!(BitAnd, bitand, BitAndAssign, bitand_assign, intersection);
328impl_op!(
329    BitXor,
330    bitxor,
331    BitXorAssign,
332    bitxor_assign,
333    symmetrical_difference
334);
335impl_op!(Sub, sub, SubAssign, sub_assign, difference);
336
337impl<T: EnumSetType> Not for EnumSet<T> {
338    type Output = EnumSet<T>;
339
340    fn not(self) -> EnumSet<T> {
341        self.complement()
342    }
343}
344
345#[cfg(feature = "serde")]
346impl<T: EnumSetType> serde::Serialize for EnumSet<T> {
347    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
348        serializer.serialize_u64(self.bits)
349    }
350}
351
352#[cfg(feature = "serde")]
353impl<'de, T: EnumSetType> serde::Deserialize<'de> for EnumSet<T> {
354    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
355        use serde::de::Error;
356
357        let bits = u64::deserialize(deserializer)?;
358        if bits & !Self::MASK != 0 {
359            return Err(D::Error::custom(
360                "bitmask holds bits that name no variant of the enum",
361            ));
362        }
363        Ok(Self {
364            bits,
365            _ty: PhantomData,
366        })
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373    use crate::EnumSetType;
374
375    #[derive(Debug, Clone, Copy, PartialEq, Eq, EnumSetType)]
376    enum Usage {
377        Load,
378        Store,
379        Arithmetic,
380    }
381
382    #[test]
383    fn all_holds_every_variant() {
384        let all = EnumSet::<Usage>::all();
385        assert_eq!(all.len(), 3);
386        assert!(all.contains(Usage::Load));
387        assert!(all.contains(Usage::Store));
388        assert!(all.contains(Usage::Arithmetic));
389    }
390
391    #[test]
392    fn empty_is_the_default() {
393        assert_eq!(EnumSet::<Usage>::default(), EnumSet::empty());
394        assert!(EnumSet::<Usage>::empty().is_empty());
395        assert_eq!(EnumSet::<Usage>::empty().len(), 0);
396    }
397
398    #[test]
399    fn or_on_variants_builds_a_set() {
400        let set = Usage::Load | Usage::Store;
401        assert_eq!(set.len(), 2);
402        assert!(set.contains(Usage::Load));
403        assert!(!set.contains(Usage::Arithmetic));
404    }
405
406    #[test]
407    fn insert_and_remove_report_whether_the_set_changed() {
408        let mut set = EnumSet::empty();
409        assert!(set.insert(Usage::Load));
410        assert!(!set.insert(Usage::Load));
411        assert!(set.remove(Usage::Load));
412        assert!(!set.remove(Usage::Load));
413    }
414
415    #[test]
416    fn iter_yields_declaration_order() {
417        let set = Usage::Arithmetic | Usage::Load;
418        let seen: alloc::vec::Vec<_> = set.iter().collect();
419        assert_eq!(seen, alloc::vec![Usage::Load, Usage::Arithmetic]);
420    }
421
422    #[test]
423    fn complement_stays_inside_the_variant_mask() {
424        let set = EnumSet::only(Usage::Load);
425        assert_eq!(!set, Usage::Store | Usage::Arithmetic);
426        assert_eq!((!EnumSet::<Usage>::empty()), EnumSet::all());
427    }
428
429    #[test]
430    fn set_algebra() {
431        let a = Usage::Load | Usage::Store;
432        let b = Usage::Store | Usage::Arithmetic;
433        assert_eq!(a.intersection(b), EnumSet::only(Usage::Store));
434        assert_eq!(a.difference(b), EnumSet::only(Usage::Load));
435        assert_eq!(a.union(b), EnumSet::all());
436        assert_eq!(a.symmetrical_difference(b), Usage::Load | Usage::Arithmetic);
437        assert!(EnumSet::only(Usage::Load).is_subset(a));
438        assert!(a.is_superset(EnumSet::only(Usage::Load)));
439        assert!(EnumSet::only(Usage::Load).is_disjoint(EnumSet::only(Usage::Store)));
440    }
441}
cubecl_ir/enum_set.rs

cubecl_ir/
enum_set.rs
