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
46pub struct FlagVec<N: FlagLength, I = usize>(N, DefaultVec<Elt>, PhantomData<I>);
51
52pub 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
61pub 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 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 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 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 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 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 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 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}