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, ParseAsInline, ParseInline, ParseInput, Tagged, ToOutput,
10    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}