object_rainbow_array_map/
lib.rs1use std::{
2 collections::{BTreeMap, btree_map},
3 marker::PhantomData,
4 ops::RangeBounds,
5};
6
7use bitvec::array::BitArray;
8use object_rainbow::{
9 Equivalent, Inline, InlineOutput, ListHashes, MaybeHasNiche, Parse, ParseAsInline, ParseInline,
10 ParseInput, RainbowIterator, Size, Tagged, ToOutput, Topological, assert_impl,
11};
12
13type Bits = BitArray<[u8; 32]>;
14
15#[derive(Tagged, ListHashes, Topological, ParseAsInline, Clone)]
16pub struct ArrayMap<T> {
17 bits: Bits,
18 map: BTreeMap<u8, T>,
19}
20
21impl<T: InlineOutput> ToOutput for ArrayMap<T> {
22 fn to_output(&self, output: &mut dyn object_rainbow::Output) {
23 self.bits.to_output(output);
24 self.map.values().iter_to_output(output);
25 }
26}
27
28impl<T: InlineOutput> InlineOutput for ArrayMap<T> {}
29
30impl<T: ParseInline<I>, I: ParseInput> ParseInline<I> for ArrayMap<T> {
31 fn parse_inline(input: &mut I) -> object_rainbow::Result<Self> {
32 let bits = input.parse_inline::<Bits>()?;
33 let map = bits
34 .iter_ones()
35 .map(|one| {
36 Ok((
37 u8::try_from(one).expect("overflow"),
38 input.parse_inline::<T>()?,
39 ))
40 })
41 .collect::<object_rainbow::Result<_>>()?;
42 Ok(Self { bits, map })
43 }
44}
45
46assert_impl!(
47 impl<T, E> Inline<E> for ArrayMap<T>
48 where
49 T: Inline<E>,
50 E: Clone,
51 {
52 }
53);
54
55impl<T> ArrayMap<T> {
56 pub fn get_mut(&mut self, key: u8) -> Option<&mut T> {
57 self.map.get_mut(&key)
58 }
59
60 pub fn get(&self, key: u8) -> Option<&T> {
61 self.map.get(&key)
62 }
63
64 pub fn insert(&mut self, key: u8, value: T) -> Option<T> {
65 self.bits.set(key as usize, true);
66 self.map.insert(key, value)
67 }
68
69 pub fn is_empty(&self) -> bool {
70 self.map.is_empty()
71 }
72
73 pub fn len(&self) -> usize {
74 self.map.len()
75 }
76
77 pub fn contains(&self, key: u8) -> bool {
78 self.bits[key as usize]
79 }
80
81 pub fn remove(&mut self, key: u8) -> Option<T> {
82 self.bits.set(key as usize, false);
83 self.map.remove(&key)
84 }
85
86 pub fn range<R: RangeBounds<u8>>(&self, range: R) -> Range<'_, T> {
87 Range {
88 range: self.map.range(range),
89 }
90 }
91
92 pub fn pop_first(&mut self) -> Option<(u8, T)> {
93 self.map
94 .pop_first()
95 .inspect(|&(key, _)| self.bits.set(key as usize, false))
96 }
97
98 pub const fn new() -> Self {
99 Self {
100 bits: BitArray {
101 _ord: PhantomData,
102 data: [0; _],
103 },
104 map: BTreeMap::new(),
105 }
106 }
107
108 pub fn iter(&self) -> Entries<'_, T> {
109 Entries {
110 inner: self.map.iter(),
111 }
112 }
113
114 pub fn iter_mut(&mut self) -> EntriesMut<'_, T> {
115 EntriesMut {
116 inner: self.map.iter_mut(),
117 }
118 }
119}
120
121pub struct Range<'a, T> {
122 range: btree_map::Range<'a, u8, T>,
123}
124
125impl<'a, T> Iterator for Range<'a, T> {
126 type Item = (u8, &'a T);
127
128 fn next(&mut self) -> Option<Self::Item> {
129 self.range.next().map(|(&k, v)| (k, v))
130 }
131}
132
133impl<T> Default for ArrayMap<T> {
134 fn default() -> Self {
135 Self {
136 bits: Default::default(),
137 map: Default::default(),
138 }
139 }
140}
141
142impl<T> Extend<(u8, T)> for ArrayMap<T> {
143 fn extend<I: IntoIterator<Item = (u8, T)>>(&mut self, iter: I) {
144 iter.into_iter().for_each(|(key, value)| {
145 self.insert(key, value);
146 });
147 }
148}
149
150impl<T> FromIterator<(u8, T)> for ArrayMap<T> {
151 fn from_iter<I: IntoIterator<Item = (u8, T)>>(iter: I) -> Self {
152 let mut map = Self::default();
153 map.extend(iter);
154 map
155 }
156}
157
158impl<T, const N: usize> From<[(u8, T); N]> for ArrayMap<T> {
159 fn from(value: [(u8, T); N]) -> Self {
160 value.into_iter().collect()
161 }
162}
163
164pub struct Entries<'a, T> {
165 inner: std::collections::btree_map::Iter<'a, u8, T>,
166}
167
168impl<'a, T> Iterator for Entries<'a, T> {
169 type Item = (u8, &'a T);
170
171 fn next(&mut self) -> Option<Self::Item> {
172 let (key, value) = self.inner.next()?;
173 Some((*key, value))
174 }
175
176 fn size_hint(&self) -> (usize, Option<usize>) {
177 self.inner.size_hint()
178 }
179}
180
181impl<'a, T> IntoIterator for &'a ArrayMap<T> {
182 type Item = (u8, &'a T);
183
184 type IntoIter = Entries<'a, T>;
185
186 fn into_iter(self) -> Self::IntoIter {
187 self.iter()
188 }
189}
190
191pub struct EntriesMut<'a, T> {
192 inner: std::collections::btree_map::IterMut<'a, u8, T>,
193}
194
195impl<'a, T> Iterator for EntriesMut<'a, T> {
196 type Item = (u8, &'a mut T);
197
198 fn next(&mut self) -> Option<Self::Item> {
199 let (key, value) = self.inner.next()?;
200 Some((*key, value))
201 }
202
203 fn size_hint(&self) -> (usize, Option<usize>) {
204 self.inner.size_hint()
205 }
206}
207
208impl<'a, T> IntoIterator for &'a mut ArrayMap<T> {
209 type Item = (u8, &'a mut T);
210
211 type IntoIter = EntriesMut<'a, T>;
212
213 fn into_iter(self) -> Self::IntoIter {
214 self.iter_mut()
215 }
216}
217
218#[derive(
219 ToOutput,
220 InlineOutput,
221 Tagged,
222 ListHashes,
223 Topological,
224 Parse,
225 ParseInline,
226 Clone,
227 Size,
228 MaybeHasNiche,
229 Default,
230)]
231pub struct ArraySet {
232 bits: Bits,
233}
234
235assert_impl!(
236 impl<E> Inline<E> for ArraySet where E: Clone {}
237);
238
239impl ArraySet {
240 pub fn insert(&mut self, key: u8) -> bool {
241 if !self.bits[key as usize] {
242 self.bits.set(key as usize, true);
243 true
244 } else {
245 false
246 }
247 }
248
249 pub fn contains(&self, key: u8) -> bool {
250 self.bits[key as usize]
251 }
252}
253
254impl Extend<u8> for ArraySet {
255 fn extend<T: IntoIterator<Item = u8>>(&mut self, iter: T) {
256 iter.into_iter().for_each(|key| {
257 self.insert(key);
258 });
259 }
260}
261
262impl FromIterator<u8> for ArraySet {
263 fn from_iter<T: IntoIterator<Item = u8>>(iter: T) -> Self {
264 let mut set = Self::default();
265 set.extend(iter);
266 set
267 }
268}
269
270impl<const N: usize> From<[u8; N]> for ArraySet {
271 fn from(value: [u8; N]) -> Self {
272 value.into_iter().collect()
273 }
274}
275
276impl Equivalent<ArrayMap<()>> for ArraySet {
277 fn into_equivalent(self) -> ArrayMap<()> {
278 ArrayMap {
279 bits: self.bits,
280 map: self
281 .bits
282 .iter_ones()
283 .map(|one| (u8::try_from(one).expect("overflow"), ()))
284 .collect(),
285 }
286 }
287
288 fn from_equivalent(ArrayMap { bits, .. }: ArrayMap<()>) -> Self {
289 Self { bits }
290 }
291}
292
293#[test]
294fn reparse() -> object_rainbow::Result<()> {
295 use object_rainbow::ParseSlice;
296 let mut map = ArrayMap::<u8>::new();
297 map.insert(12, 34);
298 map = map.reparse()?;
299 assert_eq!(*map.get(12).unwrap(), 34);
300 Ok(())
301}