agsol_common/
max_len_btree.rs

1use super::{MaxSerializedLen, CONTENTS_FULL};
2
3use borsh::{BorshDeserialize, BorshSerialize};
4
5use std::cmp::Ordering;
6use std::collections::BTreeMap;
7use std::convert::TryFrom;
8use std::hash::Hash;
9
10#[repr(C)]
11#[derive(BorshDeserialize, BorshSerialize, Clone, Debug)]
12pub struct MaxLenBTreeMap<K, V, const N: usize>
13where
14    K: MaxSerializedLen + Clone + Ord + Hash,
15    V: MaxSerializedLen + Clone,
16{
17    contents: BTreeMap<K, V>,
18}
19
20impl<K, V, const N: usize> MaxSerializedLen for MaxLenBTreeMap<K, V, N>
21where
22    K: MaxSerializedLen + Clone + Ord + Hash,
23    V: MaxSerializedLen + Clone,
24{
25    const MAX_SERIALIZED_LEN: usize = 4 + N * (K::MAX_SERIALIZED_LEN + V::MAX_SERIALIZED_LEN);
26}
27
28impl<K, V, const N: usize> MaxLenBTreeMap<K, V, N>
29where
30    K: MaxSerializedLen + Clone + Ord + Hash,
31    V: MaxSerializedLen + Clone,
32{
33    pub fn new() -> Self {
34        Self {
35            contents: BTreeMap::new(),
36        }
37    }
38
39    pub fn len(&self) -> usize {
40        self.contents.len()
41    }
42
43    pub fn is_empty(&self) -> bool {
44        self.contents.len() == 0
45    }
46
47    pub fn is_full(&self) -> bool {
48        self.contents.len() == N
49    }
50
51    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, &'static str> {
52        if self.contents.keys().any(|k| k.cmp(&key) == Ordering::Equal) || !self.is_full() {
53            Ok(self.contents.insert(key, value))
54        } else {
55            Err(CONTENTS_FULL)
56        }
57    }
58
59    pub fn remove(&mut self, key: &K) -> Option<V> {
60        self.contents.remove(key)
61    }
62
63    pub fn get(&self, key: &K) -> Option<&V> {
64        self.contents.get(key)
65    }
66
67    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
68        self.contents.get_mut(key)
69    }
70
71    pub fn contains_key(&self, key: &K) -> bool {
72        self.contents.contains_key(key)
73    }
74
75    pub fn clear(&mut self) {
76        self.contents.clear();
77    }
78
79    pub fn contents(&self) -> &BTreeMap<K, V> {
80        &self.contents
81    }
82}
83
84impl<K, V, const N: usize> TryFrom<BTreeMap<K, V>> for MaxLenBTreeMap<K, V, N>
85where
86    K: MaxSerializedLen + Clone + Ord + Hash,
87    V: MaxSerializedLen + Clone,
88{
89    type Error = &'static str;
90
91    fn try_from(btree: BTreeMap<K, V>) -> Result<Self, Self::Error> {
92        if btree.len() > N {
93            return Err(CONTENTS_FULL);
94        }
95        Ok(Self { contents: btree })
96    }
97}
98
99impl<K, V, const N: usize> Default for MaxLenBTreeMap<K, V, N>
100where
101    K: MaxSerializedLen + Clone + Ord + Hash,
102    V: MaxSerializedLen + Clone,
103{
104    fn default() -> Self {
105        Self::new()
106    }
107}
108
109#[cfg(test)]
110mod test_max_len_btreemap {
111    use super::*;
112
113    type TestBTree = MaxLenBTreeMap<u8, u32, 5>;
114    type BaseBTree = BTreeMap<u8, u32>;
115
116    #[test]
117    fn valid_conversions() {
118        let mut btree = BaseBTree::new();
119
120        for i in 0..5 {
121            btree.insert(i as u8, i);
122        }
123
124        let max_len_btree = TestBTree::try_from(btree.clone()).unwrap();
125        assert_eq!(btree, max_len_btree.contents);
126    }
127
128    #[test]
129    fn invalid_conversions() {
130        let mut btree = BaseBTree::new();
131
132        for i in 0..6 {
133            btree.insert(i as u8, i);
134        }
135
136        assert!(TestBTree::try_from(btree).is_err());
137    }
138
139    #[test]
140    fn max_len_btreemap_serialized_len() {
141        let mut test_btree: TestBTree = TestBTree::new();
142        assert!(test_btree.try_to_vec().unwrap().len() <= TestBTree::MAX_SERIALIZED_LEN);
143
144        for i in 0..4 {
145            assert!(test_btree.insert(i as u8, i).is_ok());
146        }
147        // re-insert with the same key
148        assert_eq!(test_btree.insert(1_u8, 2), Ok(Some(1)));
149
150        assert!(test_btree.try_to_vec().unwrap().len() <= TestBTree::MAX_SERIALIZED_LEN);
151        assert!(test_btree.insert(83_u8, 81237).is_ok());
152        assert_eq!(
153            test_btree.try_to_vec().unwrap().len(),
154            TestBTree::MAX_SERIALIZED_LEN
155        );
156
157        assert_eq!(test_btree.insert(85_u8, 81237), Err(CONTENTS_FULL));
158        // re-insert into full map
159        assert_eq!(test_btree.insert(3_u8, 4), Ok(Some(3)));
160    }
161}
agsol_common/max_len_btree.rs

agsol_common/
max_len_btree.rs
