Skip to main content

default_vec2/
flag_vec.rs

1use crate::default_vec::DefaultVec;
2use core::fmt::{Debug, Formatter};
3use core::marker::PhantomData;
4
5type Elt = u32;
6
7pub trait FlagLength: Clone + Eq {
8    fn len(&self) -> u32;
9
10    #[inline]
11    fn base_mask(&self) -> Elt {
12        (1 << self.len()) - 1
13    }
14
15    #[inline]
16    fn chunk_size(&self) -> u32 {
17        Elt::BITS / self.len()
18    }
19
20    #[inline]
21    fn split(&self, x: usize) -> (usize, u32) {
22        let chunk_size = self.chunk_size();
23        let offset = (x % chunk_size as usize) as u32 * self.len();
24        (x / chunk_size as usize, offset)
25    }
26}
27
28#[derive(Copy, Clone, Eq, PartialEq)]
29pub struct Static<const N: u32>;
30
31impl<const N: u32> FlagLength for Static<N> {
32    #[inline]
33    fn len(&self) -> u32 {
34        N
35    }
36}
37
38impl FlagLength for u32 {
39    #[inline]
40    fn len(&self) -> u32 {
41        *self
42    }
43}
44
45/// A data structure an indexed collection of `N` bit bit-flag like objects
46/// Behaves like [`DefaultVec<u32>`](DefaultVec) but only considers the last `N` bits
47///
48/// Use [`StaticFlagVec`] when `N` is known at compile time, or [`DynamicFlagVec`] if it isn't
49pub struct FlagVec<N: FlagLength, I = usize>(N, DefaultVec<Elt>, PhantomData<I>);
50
51/// [`FlagVec`] where `N` is known at compile time
52pub type StaticFlagVec<const N: u32, I = usize> = FlagVec<Static<N>, I>;
53
54impl<const N: u32, I> Default for StaticFlagVec<N, I> {
55    fn default() -> Self {
56        FlagVec(Static, DefaultVec::default(), PhantomData)
57    }
58}
59
60/// [`DynamicFlagVec`] where `N` is not known at compile time
61pub type DynamicFlagVec<I = usize> = FlagVec<u32, I>;
62
63impl DynamicFlagVec {
64    pub fn new(flag_bit_len: u32) -> Self {
65        FlagVec(flag_bit_len, DefaultVec::default(), PhantomData)
66    }
67}
68
69impl<N: FlagLength, I> Clone for FlagVec<N, I> {
70    fn clone(&self) -> Self {
71        FlagVec(self.0.clone(), self.1.clone(), PhantomData)
72    }
73
74    fn clone_from(&mut self, source: &Self) {
75        self.0.clone_from(&source.0);
76        self.1.clone_from(&source.1);
77    }
78}
79
80impl<N: FlagLength, I> PartialEq<Self> for FlagVec<N, I> {
81    fn eq(&self, other: &Self) -> bool {
82        self.1 == other.1
83    }
84}
85
86impl<N: FlagLength, I> Eq for FlagVec<N, I> {}
87
88impl<N: FlagLength, I: Into<usize>> FlagVec<N, I> {
89    /// Equivalent to `self.set(x, self.get(x) | v)`
90    ///
91    /// ```
92    /// use default_vec2::StaticFlagVec;
93    /// let mut s: StaticFlagVec<4> = StaticFlagVec::default();
94    /// s.or_assign(0, 3);
95    /// assert_eq!(s.get(0), 3);
96    /// s.or_assign(0, 9);
97    /// assert_eq!(s.get(0), 11);
98    /// ```
99    pub fn or_assign(&mut self, x: I, v: Elt) {
100        let (chunk_idx, offset) = self.0.split(x.into());
101        let mask = (self.0.base_mask() & v) << offset;
102        let chunk = self.1.get_mut(chunk_idx);
103        *chunk |= mask;
104    }
105
106    /// Equivalent to `self.set(x, self.get(x) & v)`
107    ///
108    /// ```
109    /// use default_vec2::StaticFlagVec;
110    /// let mut s: StaticFlagVec<4> = StaticFlagVec::default();
111    /// s.or_assign(0, 11);
112    /// s.and_assign(0, 5);
113    /// assert_eq!(s.get(0), 1);
114    /// ```
115    pub fn and_assign(&mut self, x: I, v: u32) {
116        let (chunk_idx, offset) = self.0.split(x.into());
117        let mask = (self.0.base_mask() & !v) << offset;
118        let chunk = self.1.get_mut(chunk_idx);
119        *chunk &= !mask;
120    }
121
122    /// Sets the element and index `x` to the last `N` bits of `v`
123    ///
124    /// ```
125    /// use default_vec2::StaticFlagVec;
126    /// let mut s: StaticFlagVec<4> = StaticFlagVec::default();
127    /// s.set(0, 3);
128    /// assert_eq!(s.get(0), 3);
129    /// s.set(0, 9);
130    /// assert_eq!(s.get(0), 9);
131    /// s.set(0, 18);
132    /// assert_eq!(s.get(0), 2);
133    /// ```
134    pub fn set(&mut self, x: I, v: u32) {
135        let (chunk_idx, offset) = self.0.split(x.into());
136        let mask_or = (self.0.base_mask() & v) << offset;
137        let mask_and = (self.0.base_mask() & !v) << offset;
138        let chunk = self.1.get_mut(chunk_idx);
139        *chunk |= mask_or;
140        *chunk &= !mask_and;
141    }
142
143    /// Returns the element and index `x`
144    pub fn get(&self, x: I) -> Elt {
145        let (chunk_idx, offset) = self.0.split(x.into());
146        let chunk = self.1.get(chunk_idx);
147        (chunk >> offset) & self.0.base_mask()
148    }
149
150    /// Same as `get` but already reserves space for `x`
151    pub fn get_reserve(&mut self, x: I) -> Elt {
152        let (chunk_idx, offset) = self.0.split(x.into());
153        let chunk = self.1.get_mut(chunk_idx);
154        (*chunk >> offset) & self.0.base_mask()
155    }
156
157    /// Removes all elements from the set
158    pub fn clear(&mut self) {
159        self.1.clear()
160    }
161
162    pub fn capacity(&self) -> usize {
163        self.1.capacity() * self.0.chunk_size() as usize
164    }
165
166    /// Iterate over the elements of `self`
167    /// ```
168    /// use default_vec2::StaticFlagVec;
169    /// let mut s: StaticFlagVec<10> = StaticFlagVec::default();
170    /// s.set(0, 42);
171    /// s.set(2, 999);
172    /// s.set(3, 365);
173    /// let res: Vec<_> = s.iter().collect();
174    /// assert_eq!(&res[..4], [42, 0, 999, 365])
175    /// ```
176    pub fn iter(&self) -> impl Iterator<Item = Elt> + '_ {
177        self.1.iter().flat_map(move |x| {
178            let mut x = *x;
179            (0..self.0.chunk_size()).map(move |_| {
180                let res = x & self.0.base_mask();
181                x >>= self.0.len();
182                res
183            })
184        })
185    }
186}
187
188#[test]
189fn test() {
190    use alloc::vec;
191    use alloc::vec::Vec;
192    let v1 = vec![
193        145, 114, 177, 130, 57, 228, 108, 147, 117, 119, 102, 143, 216, 215, 2, 215, 191, 217, 96,
194        157, 200, 82, 220, 211, 66, 183, 16, 173, 174, 246, 232, 248, 174, 40, 33, 169, 12, 191,
195        171, 24, 32, 196, 104, 101, 216, 155, 132, 91, 32, 95, 122, 149, 64, 56, 218, 129, 41, 26,
196        63, 87, 77, 120, 101, 213, 26, 141, 166, 167, 70, 16, 136, 159, 157, 144, 94, 52, 121, 188,
197        219, 72, 75, 74, 223, 148, 160, 13, 126, 89, 148, 149, 24, 221, 120, 204, 148, 167, 179,
198        120, 93, 126,
199    ];
200    let v2 = vec![
201        177, 234, 192, 132, 135, 63, 26, 136, 95, 252, 137, 0, 200, 214, 60, 61, 93, 45, 218, 173,
202        81, 162, 23, 188, 97, 228, 26, 159, 199, 238, 128, 59, 19, 143, 15, 133, 52, 182, 208, 56,
203        54, 99, 83, 47, 80, 208, 8, 64, 142, 205, 248, 11, 71, 252, 101, 32, 219, 117, 160, 120,
204        217, 111, 3, 69, 215, 49, 122, 147, 147, 7, 199, 157, 69, 73, 159, 250, 241, 136, 85, 101,
205        14, 32, 118, 64, 123, 154, 236, 7, 239, 206, 159, 12, 170, 84, 101, 136, 54, 138, 182, 247,
206    ];
207    assert_eq!(v1.len(), v2.len());
208    let vor: Vec<_> = v1.iter().zip(&v2).map(|(&x, &y)| x | y).collect();
209    let vand: Vec<_> = v1.iter().zip(&v2).map(|(&x, &y)| x & y).collect();
210    let mut fv1 = StaticFlagVec::<9>::default();
211    for (i, &x) in v1.iter().enumerate() {
212        fv1.set(i, x);
213    }
214    assert_eq!(&fv1.iter().collect::<Vec<_>>()[..100], &v1);
215    let mut fv2 = fv1.clone();
216    for (i, &x) in v2.iter().enumerate() {
217        assert_eq!(fv2.get(i), v1[i]);
218        fv2.set(i, x);
219    }
220    assert_eq!(&fv2.iter().collect::<Vec<_>>()[..100], &v2);
221    let mut fvor = fv1.clone();
222    for (i, &x) in v2.iter().enumerate() {
223        fvor.or_assign(i, x);
224    }
225    assert_eq!(&fvor.iter().collect::<Vec<_>>()[..100], &vor);
226    let mut fvand = fvor;
227    fvand.clone_from(&fv1);
228    for (i, &x) in v2.iter().enumerate() {
229        fvand.and_assign(i, x);
230    }
231    assert_eq!(&fvand.iter().collect::<Vec<_>>()[..100], &vand);
232}
233
234impl<N: FlagLength, I: Into<usize>> Debug for FlagVec<N, I> {
235    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
236        f.debug_list().entries(self.iter()).finish()
237    }
238}