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    Inline, InlineOutput, ListHashes, MaybeHasNiche, Parse, ParseAsInline, ParseInline, ParseInput,
10    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> where T: Inline<E> {}
41);
42
43impl<T> ArrayMap<T> {
44    pub fn get_mut(&mut self, key: u8) -> Option<&mut T> {
45        self.map.get_mut(&key)
46    }
47
48    pub fn get(&self, key: u8) -> Option<&T> {
49        self.map.get(&key)
50    }
51
52    pub fn insert(&mut self, key: u8, value: T) -> Option<T> {
53        self.bits.set(key as usize, true);
54        self.map.insert(key, value)
55    }
56
57    pub fn is_empty(&self) -> bool {
58        self.map.is_empty()
59    }
60
61    pub fn len(&self) -> usize {
62        self.map.len()
63    }
64
65    pub fn contains(&self, key: u8) -> bool {
66        self.bits[key as usize]
67    }
68
69    pub fn remove(&mut self, key: u8) -> Option<T> {
70        self.bits.set(key as usize, false);
71        self.map.remove(&key)
72    }
73
74    pub fn range<R: RangeBounds<u8>>(&self, range: R) -> Range<'_, T> {
75        Range {
76            range: self.map.range(range),
77        }
78    }
79
80    pub fn pop_first(&mut self) -> Option<(u8, T)> {
81        self.map
82            .pop_first()
83            .inspect(|&(key, _)| self.bits.set(key as usize, false))
84    }
85
86    pub const fn new() -> Self {
87        Self {
88            bits: BitArray {
89                _ord: PhantomData,
90                data: [0; _],
91            },
92            map: BTreeMap::new(),
93        }
94    }
95}
96
97pub struct Range<'a, T> {
98    range: btree_map::Range<'a, u8, T>,
99}
100
101impl<'a, T> Iterator for Range<'a, T> {
102    type Item = (u8, &'a T);
103
104    fn next(&mut self) -> Option<Self::Item> {
105        self.range.next().map(|(&k, v)| (k, v))
106    }
107}
108
109impl<T> Default for ArrayMap<T> {
110    fn default() -> Self {
111        Self {
112            bits: Default::default(),
113            map: Default::default(),
114        }
115    }
116}
117
118impl<T> Extend<(u8, T)> for ArrayMap<T> {
119    fn extend<I: IntoIterator<Item = (u8, T)>>(&mut self, iter: I) {
120        iter.into_iter().for_each(|(key, value)| {
121            self.insert(key, value);
122        });
123    }
124}
125
126impl<T> FromIterator<(u8, T)> for ArrayMap<T> {
127    fn from_iter<I: IntoIterator<Item = (u8, T)>>(iter: I) -> Self {
128        let mut map = Self::default();
129        map.extend(iter);
130        map
131    }
132}
133
134impl<T, const N: usize> From<[(u8, T); N]> for ArrayMap<T> {
135    fn from(value: [(u8, T); N]) -> Self {
136        value.into_iter().collect()
137    }
138}
139
140#[derive(
141    ToOutput,
142    InlineOutput,
143    Tagged,
144    ListHashes,
145    Topological,
146    Parse,
147    ParseInline,
148    Clone,
149    Size,
150    MaybeHasNiche,
151    Default,
152)]
153pub struct ArraySet {
154    bits: Bits,
155}
156
157assert_impl!(
158    impl<E> Inline<E> for ArraySet {}
159);
160
161impl ArraySet {
162    pub fn insert(&mut self, key: u8) -> bool {
163        if !self.bits[key as usize] {
164            self.bits.set(key as usize, true);
165            true
166        } else {
167            false
168        }
169    }
170
171    pub fn contains(&self, key: u8) -> bool {
172        self.bits[key as usize]
173    }
174}
175
176impl Extend<u8> for ArraySet {
177    fn extend<T: IntoIterator<Item = u8>>(&mut self, iter: T) {
178        iter.into_iter().for_each(|key| {
179            self.insert(key);
180        });
181    }
182}
183
184impl FromIterator<u8> for ArraySet {
185    fn from_iter<T: IntoIterator<Item = u8>>(iter: T) -> Self {
186        let mut set = Self::default();
187        set.extend(iter);
188        set
189    }
190}
191
192impl<const N: usize> From<[u8; N]> for ArraySet {
193    fn from(value: [u8; N]) -> Self {
194        value.into_iter().collect()
195    }
196}