Skip to main content

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        self.1.clone_from(&source.1);
78    }
79}
80
81impl<N: FlagLength, I> PartialEq<Self> for FlagVec<N, I> {
82    fn eq(&self, other: &Self) -> bool {
83        self.1 == other.1
84    }
85}
86
87impl<N: FlagLength, I> Eq for FlagVec<N, I> {}
88
89impl<N: FlagLength, I: Into<usize>> FlagVec<N, I> {
90    /// Equivalent to `self.set(x, self.get(x) | v)`
91    ///
92    /// ```
93    /// use default_vec2::StaticFlagVec;
94    /// let mut s: StaticFlagVec<4> = StaticFlagVec::default();
95    /// s.or_assign(0, 3);
96    /// assert_eq!(s.get(0), 3);
97    /// s.or_assign(0, 9);
98    /// assert_eq!(s.get(0), 11);
99    /// ```
100    pub fn or_assign(&mut self, x: I, v: Elt) {
101        let (chunk_idx, offset) = self.0.split(x.into());
102        let mask = (self.0.base_mask() & v) << offset;
103        let chunk = self.1.get_mut(chunk_idx);
104        *chunk |= mask;
105    }
106
107    /// Equivalent to `self.set(x, self.get(x) & v)`
108    ///
109    /// ```
110    /// use default_vec2::StaticFlagVec;
111    /// let mut s: StaticFlagVec<4> = StaticFlagVec::default();
112    /// s.or_assign(0, 11);
113    /// s.and_assign(0, 5);
114    /// assert_eq!(s.get(0), 1);
115    /// ```
116    pub fn and_assign(&mut self, x: I, v: u32) {
117        let (chunk_idx, offset) = self.0.split(x.into());
118        let mask = (self.0.base_mask() & !v) << offset;
119        let chunk = self.1.get_mut(chunk_idx);
120        *chunk &= !mask;
121    }
122
123    /// Sets the element and index `x` to the last `N` bits of `v`
124    ///
125    /// ```
126    /// use default_vec2::StaticFlagVec;
127    /// let mut s: StaticFlagVec<4> = StaticFlagVec::default();
128    /// s.set(0, 3);
129    /// assert_eq!(s.get(0), 3);
130    /// s.set(0, 9);
131    /// assert_eq!(s.get(0), 9);
132    /// s.set(0, 18);
133    /// assert_eq!(s.get(0), 2);
134    /// ```
135    pub fn set(&mut self, x: I, v: u32) {
136        let (chunk_idx, offset) = self.0.split(x.into());
137        let mask_or = (self.0.base_mask() & v) << offset;
138        let mask_and = (self.0.base_mask() & !v) << offset;
139        let chunk = self.1.get_mut(chunk_idx);
140        *chunk |= mask_or;
141        *chunk &= !mask_and;
142    }
143
144    /// Returns the element and index `x`
145    pub fn get(&self, x: I) -> Elt {
146        let (chunk_idx, offset) = self.0.split(x.into());
147        let chunk = self.1.get(chunk_idx);
148        (chunk >> offset) & self.0.base_mask()
149    }
150
151    /// Same as `get` but already reserves space for `x`
152    pub fn get_reserve(&mut self, x: I) -> Elt {
153        let (chunk_idx, offset) = self.0.split(x.into());
154        let chunk = self.1.get_mut(chunk_idx);
155        (*chunk >> offset) & self.0.base_mask()
156    }
157
158    /// Removes all elements from the set
159    pub fn clear(&mut self) {
160        self.1.clear()
161    }
162
163    pub fn capacity(&self) -> usize {
164        self.1.capacity() * self.0.chunk_size() as usize
165    }
166
167    /// Iterate over the elements of `self`
168    /// ```
169    /// use default_vec2::StaticFlagVec;
170    /// let mut s: StaticFlagVec<10> = StaticFlagVec::default();
171    /// s.set(0, 42);
172    /// s.set(2, 999);
173    /// s.set(3, 365);
174    /// let res: Vec<_> = s.iter().collect();
175    /// assert_eq!(&res[..4], [42, 0, 999, 365])
176    /// ```
177    pub fn iter(&self) -> impl Iterator<Item = Elt> + '_ {
178        self.1.iter().flat_map(move |x| {
179            let mut x = *x;
180            (0..self.0.chunk_size()).map(move |_| {
181                let res = x & self.0.base_mask();
182                x >>= self.0.len();
183                res
184            })
185        })
186    }
187}
188
189#[test]
190fn test() {
191    use alloc::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}