Skip to main content

object_rainbow_array_map/
lib.rs

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