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, Size, Tagged, ToOutput, Topological, assert_impl,
11};
12
13type Bits = BitArray<[u8; 32]>;
14
15#[derive(ToOutput, Tagged, ListHashes, Topological, ParseAsInline, Clone)]
16pub struct ArrayMap<T> {
17    bits: Bits,
18    map: BTreeMap<u8, T>,
19}
20
21impl<T: InlineOutput> InlineOutput for ArrayMap<T> {}
22
23impl<T: ParseInline<I>, I: ParseInput> ParseInline<I> for ArrayMap<T> {
24    fn parse_inline(input: &mut I) -> object_rainbow::Result<Self> {
25        let bits = input.parse_inline::<Bits>()?;
26        let map = bits
27            .iter_ones()
28            .map(|one| {
29                Ok((
30                    u8::try_from(one).expect("overflow"),
31                    input.parse_inline::<T>()?,
32                ))
33            })
34            .collect::<object_rainbow::Result<_>>()?;
35        Ok(Self { bits, map })
36    }
37}
38
39assert_impl!(
40    impl<T, E> Inline<E> for ArrayMap<T>
41    where
42        T: Inline<E>,
43        E: Clone,
44    {
45    }
46);
47
48impl<T> ArrayMap<T> {
49    pub fn get_mut(&mut self, key: u8) -> Option<&mut T> {
50        self.map.get_mut(&key)
51    }
52
53    pub fn get(&self, key: u8) -> Option<&T> {
54        self.map.get(&key)
55    }
56
57    pub fn insert(&mut self, key: u8, value: T) -> Option<T> {
58        self.bits.set(key as usize, true);
59        self.map.insert(key, value)
60    }
61
62    pub fn is_empty(&self) -> bool {
63        self.map.is_empty()
64    }
65
66    pub fn len(&self) -> usize {
67        self.map.len()
68    }
69
70    pub fn contains(&self, key: u8) -> bool {
71        self.bits[key as usize]
72    }
73
74    pub fn remove(&mut self, key: u8) -> Option<T> {
75        self.bits.set(key as usize, false);
76        self.map.remove(&key)
77    }
78
79    pub fn range<R: RangeBounds<u8>>(&self, range: R) -> Range<'_, T> {
80        Range {
81            range: self.map.range(range),
82        }
83    }
84
85    pub fn pop_first(&mut self) -> Option<(u8, T)> {
86        self.map
87            .pop_first()
88            .inspect(|&(key, _)| self.bits.set(key as usize, false))
89    }
90
91    pub const fn new() -> Self {
92        Self {
93            bits: BitArray {
94                _ord: PhantomData,
95                data: [0; _],
96            },
97            map: BTreeMap::new(),
98        }
99    }
100}
101
102pub struct Range<'a, T> {
103    range: btree_map::Range<'a, u8, T>,
104}
105
106impl<'a, T> Iterator for Range<'a, T> {
107    type Item = (u8, &'a T);
108
109    fn next(&mut self) -> Option<Self::Item> {
110        self.range.next().map(|(&k, v)| (k, v))
111    }
112}
113
114impl<T> Default for ArrayMap<T> {
115    fn default() -> Self {
116        Self {
117            bits: Default::default(),
118            map: Default::default(),
119        }
120    }
121}
122
123impl<T> Extend<(u8, T)> for ArrayMap<T> {
124    fn extend<I: IntoIterator<Item = (u8, T)>>(&mut self, iter: I) {
125        iter.into_iter().for_each(|(key, value)| {
126            self.insert(key, value);
127        });
128    }
129}
130
131impl<T> FromIterator<(u8, T)> for ArrayMap<T> {
132    fn from_iter<I: IntoIterator<Item = (u8, T)>>(iter: I) -> Self {
133        let mut map = Self::default();
134        map.extend(iter);
135        map
136    }
137}
138
139impl<T, const N: usize> From<[(u8, T); N]> for ArrayMap<T> {
140    fn from(value: [(u8, T); N]) -> Self {
141        value.into_iter().collect()
142    }
143}
144
145#[derive(
146    ToOutput,
147    InlineOutput,
148    Tagged,
149    ListHashes,
150    Topological,
151    Parse,
152    ParseInline,
153    Clone,
154    Size,
155    MaybeHasNiche,
156    Default,
157)]
158pub struct ArraySet {
159    bits: Bits,
160}
161
162assert_impl!(
163    impl<E> Inline<E> for ArraySet where E: Clone {}
164);
165
166impl ArraySet {
167    pub fn insert(&mut self, key: u8) -> bool {
168        if !self.bits[key as usize] {
169            self.bits.set(key as usize, true);
170            true
171        } else {
172            false
173        }
174    }
175
176    pub fn contains(&self, key: u8) -> bool {
177        self.bits[key as usize]
178    }
179}
180
181impl Extend<u8> for ArraySet {
182    fn extend<T: IntoIterator<Item = u8>>(&mut self, iter: T) {
183        iter.into_iter().for_each(|key| {
184            self.insert(key);
185        });
186    }
187}
188
189impl FromIterator<u8> for ArraySet {
190    fn from_iter<T: IntoIterator<Item = u8>>(iter: T) -> Self {
191        let mut set = Self::default();
192        set.extend(iter);
193        set
194    }
195}
196
197impl<const N: usize> From<[u8; N]> for ArraySet {
198    fn from(value: [u8; N]) -> Self {
199        value.into_iter().collect()
200    }
201}
202
203impl Equivalent<ArrayMap<()>> for ArraySet {
204    fn into_equivalent(self) -> ArrayMap<()> {
205        ArrayMap {
206            bits: self.bits,
207            map: self
208                .bits
209                .iter_ones()
210                .map(|one| (u8::try_from(one).expect("overflow"), ()))
211                .collect(),
212        }
213    }
214
215    fn from_equivalent(ArrayMap { bits, .. }: ArrayMap<()>) -> Self {
216        Self { bits }
217    }
218}