bitset_const_macros/
lib.rs

1use derive_syn_parse::Parse;
2use proc_macro::TokenStream;
3use quote::quote;
4use syn::{parse_macro_input, Ident, LitInt, Token, Visibility};
5
6#[proc_macro]
7pub fn bitset(item: TokenStream) -> TokenStream {
8    let BitSet { vis, name, bits } = parse_macro_input!(item as BitSet);
9    let bits = match bits.base10_parse::<usize>() {
10        Ok(bits) => bits,
11        Err(err) => return err.into_compile_error().into(),
12    };
13    const BYTE_SIZE: usize = u128::BITS as usize;
14    let mut bytes = bits / BYTE_SIZE;
15    if bits % BYTE_SIZE > 0 {
16        bytes += 1;
17    }
18
19    let expanded = quote! {
20        #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
21        #[repr(transparent)]
22        #vis struct #name([u128; #bytes]);
23
24        impl #name {
25            #[inline]
26            pub const fn new() -> Self {
27                Self([0; #bytes])
28            }
29
30            #[inline]
31            pub const fn get(&self, i: usize) -> bool {
32                let (i, mask) = Self::__calc_offset(i);
33                self.0[i] & mask != 0
34            }
35
36            #[inline]
37            pub fn set(&mut self, i: usize, value: bool) {
38                let (i, mask) = Self::__calc_offset(i);
39                self.0[i] |= mask;
40            }
41
42            #[inline]
43            pub fn count(&self) -> u64 {
44                self.0.iter().fold(0u64, |mut acc, curr| {
45                    acc += curr.count_ones() as u64;
46                    acc
47                })
48            }
49
50            #[doc(hidden)]
51            #[inline]
52            const fn __calc_offset(i: usize) -> (usize, u128) {
53                (i / #BYTE_SIZE, 1u128 << (i % #BYTE_SIZE))
54            }
55        }
56
57        impl core::fmt::Binary for #name {
58            #[inline]
59            fn fmt(&self, f: &mut core::fmt::Formatter<'_> ) -> core::fmt::Result {
60                for byte in self.0 {
61                    core::fmt::Binary::fmt(&byte, f)?;
62                }
63                Ok(())
64            }
65        }
66
67        impl core::ops::BitAnd for #name {
68            type Output = Self;
69            #[inline]
70            fn bitand(mut self, rhs: Self) -> Self::Output {
71                self &= rhs;
72                self
73            }
74        }
75
76        impl core::ops::BitAndAssign for #name {
77            #[inline]
78            fn bitand_assign(&mut self, rhs: Self) {
79                for (this, other) in self.0.iter_mut().zip(rhs.0.iter()) {
80                    *this &= other;
81                }
82            }
83        }
84
85        impl core::ops::BitOr for #name {
86            type Output = Self;
87            #[inline]
88            fn bitor(mut self, rhs: Self) -> Self::Output {
89                self |= rhs;
90                self
91            }
92        }
93
94        impl core::ops::BitOrAssign for #name {
95            #[inline]
96            fn bitor_assign(&mut self, rhs: Self) {
97                for (this, other) in self.0.iter_mut().zip(rhs.0.iter()) {
98                    *this |= other;
99                }
100            }
101        }
102    };
103
104    TokenStream::from(expanded)
105}
106
107#[derive(Parse)]
108struct BitSet {
109    vis: Visibility,
110    #[prefix(Token![struct])]
111    name: Ident,
112    #[prefix(Token![:])]
113    bits: LitInt,
114}