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
109pub struct Range<'a, T> {
110    range: btree_map::Range<'a, u8, T>,
111}
112
113impl<'a, T> Iterator for Range<'a, T> {
114    type Item = (u8, &'a T);
115
116    fn next(&mut self) -> Option<Self::Item> {
117        self.range.next().map(|(&k, v)| (k, v))
118    }
119}
120
121impl<T> Default for ArrayMap<T> {
122    fn default() -> Self {
123        Self {
124            bits: Default::default(),
125            map: Default::default(),
126        }
127    }
128}
129
130impl<T> Extend<(u8, T)> for ArrayMap<T> {
131    fn extend<I: IntoIterator<Item = (u8, T)>>(&mut self, iter: I) {
132        iter.into_iter().for_each(|(key, value)| {
133            self.insert(key, value);
134        });
135    }
136}
137
138impl<T> FromIterator<(u8, T)> for ArrayMap<T> {
139    fn from_iter<I: IntoIterator<Item = (u8, T)>>(iter: I) -> Self {
140        let mut map = Self::default();
141        map.extend(iter);
142        map
143    }
144}
145
146impl<T, const N: usize> From<[(u8, T); N]> for ArrayMap<T> {
147    fn from(value: [(u8, T); N]) -> Self {
148        value.into_iter().collect()
149    }
150}
151
152#[derive(
153    ToOutput,
154    InlineOutput,
155    Tagged,
156    ListHashes,
157    Topological,
158    Parse,
159    ParseInline,
160    Clone,
161    Size,
162    MaybeHasNiche,
163    Default,
164)]
165pub struct ArraySet {
166    bits: Bits,
167}
168
169assert_impl!(
170    impl<E> Inline<E> for ArraySet where E: Clone {}
171);
172
173impl ArraySet {
174    pub fn insert(&mut self, key: u8) -> bool {
175        if !self.bits[key as usize] {
176            self.bits.set(key as usize, true);
177            true
178        } else {
179            false
180        }
181    }
182
183    pub fn contains(&self, key: u8) -> bool {
184        self.bits[key as usize]
185    }
186}
187
188impl Extend<u8> for ArraySet {
189    fn extend<T: IntoIterator<Item = u8>>(&mut self, iter: T) {
190        iter.into_iter().for_each(|key| {
191            self.insert(key);
192        });
193    }
194}
195
196impl FromIterator<u8> for ArraySet {
197    fn from_iter<T: IntoIterator<Item = u8>>(iter: T) -> Self {
198        let mut set = Self::default();
199        set.extend(iter);
200        set
201    }
202}
203
204impl<const N: usize> From<[u8; N]> for ArraySet {
205    fn from(value: [u8; N]) -> Self {
206        value.into_iter().collect()
207    }
208}
209
210impl Equivalent<ArrayMap<()>> for ArraySet {
211    fn into_equivalent(self) -> ArrayMap<()> {
212        ArrayMap {
213            bits: self.bits,
214            map: self
215                .bits
216                .iter_ones()
217                .map(|one| (u8::try_from(one).expect("overflow"), ()))
218                .collect(),
219        }
220    }
221
222    fn from_equivalent(ArrayMap { bits, .. }: ArrayMap<()>) -> Self {
223        Self { bits }
224    }
225}
226
227#[test]
228fn reparse() -> object_rainbow::Result<()> {
229    use object_rainbow::ParseSlice;
230    let mut map = ArrayMap::<u8>::new();
231    map.insert(12, 34);
232    map = map.reparse()?;
233    assert_eq!(*map.get(12).unwrap(), 34);
234    Ok(())
235}
object_rainbow_array_map/lib.rs

object_rainbow_array_map/
lib.rs
