Skip to main content

object_rainbow_array_map/
lib.rs

1use 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}