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