default_vec2/
flag_vec.rs

1use crate::default_vec::DefaultVec;
2use alloc::vec::Vec;
3use core::fmt::{Debug, Formatter};
4use core::marker::PhantomData;
5
6type Elt = u32;
7
8pub trait FlagLength: Clone + Eq {
9    fn len(&self) -> u32;
10
11    #[inline]
12    fn base_mask(&self) -> Elt {
13        (1 << self.len()) - 1
14    }
15
16    #[inline]
17    fn chunk_size(&self) -> u32 {
18        Elt::BITS / self.len()
19    }
20
21    #[inline]
22    fn split(&self, x: usize) -> (usize, u32) {
23        let chunk_size = self.chunk_size();
24        let offset = (x % chunk_size as usize) as u32 * self.len();
25        (x / chunk_size as usize, offset)
26    }
27}
28
29#[derive(Copy, Clone, Eq, PartialEq)]
30pub struct Static<const N: u32>;
31
32impl<const N: u32> FlagLength for Static<N> {
33    #[inline]
34    fn len(&self) -> u32 {
35        N
36    }
37}
38
39impl FlagLength for u32 {
40    #[inline]
41    fn len(&self) -> u32 {
42        *self
43    }
44}
45
46/// A data structure an indexed collection of `N` bit bit-flag like objects
47/// Behaves like [`DefaultVec<u32>`](DefaultVec) but only considers the last `N` bits
48///
49/// Use [`StaticFlagVec`] when `N` is known at compile time, or [`DynamicFlagVec`] if it isn't
50pub struct FlagVec<N: FlagLength, I = usize>(N, DefaultVec<Elt>, PhantomData<I>);
51
52/// [`FlagVec`] where `N` is known at compile time
53pub type StaticFlagVec<const N: u32, I = usize> = FlagVec<Static<N>, I>;
54
55impl<const N: u32, I> Default for StaticFlagVec<N, I> {
56    fn default() -> Self {
57        FlagVec(Static, DefaultVec::default(), PhantomData)
58    }
59}
60
61/// [`DynamicFlagVec`] where `N` is not known at compile time
62pub type DynamicFlagVec<I = usize> = FlagVec<u32, I>;
63
64impl DynamicFlagVec {
65    pub fn new(flag_bit_len: u32) -> Self {
66        FlagVec(flag_bit_len, DefaultVec::default(), PhantomData)
67    }
68}
69
70impl<N: FlagLength, I> Clone for FlagVec<N, I> {
71    fn clone(&self) -> Self {
72        FlagVec(self.0.clone(), self.1.clone(), PhantomData)
73    }
74
75    fn clone_from(&mut self, source: &Self) {
76        self.0.clone_from(&source.0)
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    let v1 = vec![
192        145, 114, 177, 130, 57, 228, 108, 147, 117, 119, 102, 143, 216, 215, 2, 215, 191, 217, 96,
193        157, 200, 82, 220, 211, 66, 183, 16, 173, 174, 246, 232, 248, 174, 40, 33, 169, 12, 191,
194        171, 24, 32, 196, 104, 101, 216, 155, 132, 91, 32, 95, 122, 149, 64, 56, 218, 129, 41, 26,
195        63, 87, 77, 120, 101, 213, 26, 141, 166, 167, 70, 16, 136, 159, 157, 144, 94, 52, 121, 188,
196        219, 72, 75, 74, 223, 148, 160, 13, 126, 89, 148, 149, 24, 221, 120, 204, 148, 167, 179,
197        120, 93, 126,
198    ];
199    let v2 = vec![
200        177, 234, 192, 132, 135, 63, 26, 136, 95, 252, 137, 0, 200, 214, 60, 61, 93, 45, 218, 173,
201        81, 162, 23, 188, 97, 228, 26, 159, 199, 238, 128, 59, 19, 143, 15, 133, 52, 182, 208, 56,
202        54, 99, 83, 47, 80, 208, 8, 64, 142, 205, 248, 11, 71, 252, 101, 32, 219, 117, 160, 120,
203        217, 111, 3, 69, 215, 49, 122, 147, 147, 7, 199, 157, 69, 73, 159, 250, 241, 136, 85, 101,
204        14, 32, 118, 64, 123, 154, 236, 7, 239, 206, 159, 12, 170, 84, 101, 136, 54, 138, 182, 247,
205    ];
206    assert_eq!(v1.len(), v2.len());
207    let vor: Vec<_> = v1.iter().zip(&v2).map(|(&x, &y)| x | y).collect();
208    let vand: Vec<_> = v1.iter().zip(&v2).map(|(&x, &y)| x & y).collect();
209    let mut fv1 = StaticFlagVec::<9>::default();
210    for (i, &x) in v1.iter().enumerate() {
211        fv1.set(i, x);
212    }
213    assert_eq!(&fv1.iter().collect::<Vec<_>>()[..100], &v1);
214    let mut fv2 = fv1.clone();
215    for (i, &x) in v2.iter().enumerate() {
216        assert_eq!(fv2.get(i), v1[i]);
217        fv2.set(i, x);
218    }
219    assert_eq!(&fv2.iter().collect::<Vec<_>>()[..100], &v2);
220    let mut fvor = fv1.clone();
221    for (i, &x) in v2.iter().enumerate() {
222        fvor.or_assign(i, x);
223    }
224    assert_eq!(&fvor.iter().collect::<Vec<_>>()[..100], &vor);
225    let mut fvand = fv1.clone();
226    for (i, &x) in v2.iter().enumerate() {
227        fvand.and_assign(i, x);
228    }
229    assert_eq!(&fvand.iter().collect::<Vec<_>>()[..100], &vand);
230}
231
232impl<N: FlagLength, I: Into<usize>> Debug for FlagVec<N, I> {
233    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
234        f.debug_list().entries(self.iter()).finish()
235    }
236}
default_vec2/flag_vec.rs

default_vec2/
flag_vec.rs
