jk_cosmwasm_std/
storage.rs

1use std::collections::BTreeMap;
2use std::fmt;
3#[cfg(feature = "iterator")]
4use std::iter;
5#[cfg(feature = "iterator")]
6use std::ops::{Bound, RangeBounds};
7
8#[cfg(feature = "iterator")]
9use crate::iterator::{Order, Pair};
10use crate::traits::Storage;
11
12#[derive(Default)]
13pub struct MemoryStorage {
14    data: BTreeMap<Vec<u8>, Vec<u8>>,
15}
16
17impl MemoryStorage {
18    pub fn new() -> Self {
19        MemoryStorage::default()
20    }
21}
22
23impl Storage for MemoryStorage {
24    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
25        self.data.get(key).cloned()
26    }
27
28    fn set(&mut self, key: &[u8], value: &[u8]) {
29        if value.is_empty() {
30            panic!("TL;DR: Value must not be empty in Storage::set but in most cases you can use Storage::remove instead. Long story: Getting empty values from storage is not well supported at the moment. Some of our internal interfaces cannot differentiate between a non-existent key and an empty value. Right now, you cannot rely on the behaviour of empty values. To protect you from trouble later on, we stop here. Sorry for the inconvenience! We highly welcome you to contribute to CosmWasm, making this more solid one way or the other.");
31        }
32
33        self.data.insert(key.to_vec(), value.to_vec());
34    }
35
36    fn remove(&mut self, key: &[u8]) {
37        self.data.remove(key);
38    }
39
40    #[cfg(feature = "iterator")]
41    /// range allows iteration over a set of keys, either forwards or backwards
42    /// uses standard rust range notation, and eg db.range(b"foo"..b"bar") also works reverse
43    fn range<'a>(
44        &'a self,
45        start: Option<&[u8]>,
46        end: Option<&[u8]>,
47        order: Order,
48    ) -> Box<dyn Iterator<Item = Pair> + 'a> {
49        let bounds = range_bounds(start, end);
50
51        // BTreeMap.range panics if range is start > end.
52        // However, this cases represent just empty range and we treat it as such.
53        match (bounds.start_bound(), bounds.end_bound()) {
54            (Bound::Included(start), Bound::Excluded(end)) if start > end => {
55                return Box::new(iter::empty());
56            }
57            _ => {}
58        }
59
60        let iter = self.data.range(bounds);
61        match order {
62            Order::Ascending => Box::new(iter.map(clone_item)),
63            Order::Descending => Box::new(iter.rev().map(clone_item)),
64        }
65    }
66}
67
68/// This debug implementation is made for inspecting storages in unit testing.
69/// It is made for human readability only and the output can change at any time.
70impl fmt::Debug for MemoryStorage {
71    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
72        write!(f, "MemoryStorage ({} entries)", self.data.len())?;
73        f.write_str(" {\n")?;
74        for (key, value) in &self.data {
75            f.write_str("  0x")?;
76            for byte in key {
77                write!(f, "{:02x}", byte)?;
78            }
79            f.write_str(": 0x")?;
80            for byte in value {
81                write!(f, "{:02x}", byte)?;
82            }
83            f.write_str("\n")?;
84        }
85        f.write_str("}")?;
86        Ok(())
87    }
88}
89
90#[cfg(feature = "iterator")]
91fn range_bounds(start: Option<&[u8]>, end: Option<&[u8]>) -> impl RangeBounds<Vec<u8>> {
92    (
93        start.map_or(Bound::Unbounded, |x| Bound::Included(x.to_vec())),
94        end.map_or(Bound::Unbounded, |x| Bound::Excluded(x.to_vec())),
95    )
96}
97
98#[cfg(feature = "iterator")]
99/// The BTreeMap specific key-value pair reference type, as returned by BTreeMap<Vec<u8>, T>::range.
100/// This is internal as it can change any time if the map implementation is swapped out.
101type BTreeMapPairRef<'a, T = Vec<u8>> = (&'a Vec<u8>, &'a T);
102
103#[cfg(feature = "iterator")]
104fn clone_item<T: Clone>(item_ref: BTreeMapPairRef<T>) -> Pair<T> {
105    let (key, value) = item_ref;
106    (key.clone(), value.clone())
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn get_and_set() {
115        let mut store = MemoryStorage::new();
116        assert_eq!(store.get(b"foo"), None);
117        store.set(b"foo", b"bar");
118        assert_eq!(store.get(b"foo"), Some(b"bar".to_vec()));
119        assert_eq!(store.get(b"food"), None);
120    }
121
122    #[test]
123    #[should_panic(
124        expected = "Getting empty values from storage is not well supported at the moment."
125    )]
126    fn set_panics_for_empty() {
127        let mut store = MemoryStorage::new();
128        store.set(b"foo", b"");
129    }
130
131    #[test]
132    fn delete() {
133        let mut store = MemoryStorage::new();
134        store.set(b"foo", b"bar");
135        store.set(b"food", b"bank");
136        store.remove(b"foo");
137
138        assert_eq!(store.get(b"foo"), None);
139        assert_eq!(store.get(b"food"), Some(b"bank".to_vec()));
140    }
141
142    #[test]
143    #[cfg(feature = "iterator")]
144    fn iterator() {
145        let mut store = MemoryStorage::new();
146        store.set(b"foo", b"bar");
147
148        // ensure we had previously set "foo" = "bar"
149        assert_eq!(store.get(b"foo"), Some(b"bar".to_vec()));
150        assert_eq!(store.range(None, None, Order::Ascending).count(), 1);
151
152        // setup - add some data, and delete part of it as well
153        store.set(b"ant", b"hill");
154        store.set(b"ze", b"bra");
155
156        // noise that should be ignored
157        store.set(b"bye", b"bye");
158        store.remove(b"bye");
159
160        // unbounded
161        {
162            let iter = store.range(None, None, Order::Ascending);
163            let elements: Vec<Pair> = iter.collect();
164            assert_eq!(
165                elements,
166                vec![
167                    (b"ant".to_vec(), b"hill".to_vec()),
168                    (b"foo".to_vec(), b"bar".to_vec()),
169                    (b"ze".to_vec(), b"bra".to_vec()),
170                ]
171            );
172        }
173
174        // unbounded (descending)
175        {
176            let iter = store.range(None, None, Order::Descending);
177            let elements: Vec<Pair> = iter.collect();
178            assert_eq!(
179                elements,
180                vec![
181                    (b"ze".to_vec(), b"bra".to_vec()),
182                    (b"foo".to_vec(), b"bar".to_vec()),
183                    (b"ant".to_vec(), b"hill".to_vec()),
184                ]
185            );
186        }
187
188        // bounded
189        {
190            let iter = store.range(Some(b"f"), Some(b"n"), Order::Ascending);
191            let elements: Vec<Pair> = iter.collect();
192            assert_eq!(elements, vec![(b"foo".to_vec(), b"bar".to_vec())]);
193        }
194
195        // bounded (descending)
196        {
197            let iter = store.range(Some(b"air"), Some(b"loop"), Order::Descending);
198            let elements: Vec<Pair> = iter.collect();
199            assert_eq!(
200                elements,
201                vec![
202                    (b"foo".to_vec(), b"bar".to_vec()),
203                    (b"ant".to_vec(), b"hill".to_vec()),
204                ]
205            );
206        }
207
208        // bounded empty [a, a)
209        {
210            let iter = store.range(Some(b"foo"), Some(b"foo"), Order::Ascending);
211            let elements: Vec<Pair> = iter.collect();
212            assert_eq!(elements, vec![]);
213        }
214
215        // bounded empty [a, a) (descending)
216        {
217            let iter = store.range(Some(b"foo"), Some(b"foo"), Order::Descending);
218            let elements: Vec<Pair> = iter.collect();
219            assert_eq!(elements, vec![]);
220        }
221
222        // bounded empty [a, b) with b < a
223        {
224            let iter = store.range(Some(b"z"), Some(b"a"), Order::Ascending);
225            let elements: Vec<Pair> = iter.collect();
226            assert_eq!(elements, vec![]);
227        }
228
229        // bounded empty [a, b) with b < a (descending)
230        {
231            let iter = store.range(Some(b"z"), Some(b"a"), Order::Descending);
232            let elements: Vec<Pair> = iter.collect();
233            assert_eq!(elements, vec![]);
234        }
235
236        // right unbounded
237        {
238            let iter = store.range(Some(b"f"), None, Order::Ascending);
239            let elements: Vec<Pair> = iter.collect();
240            assert_eq!(
241                elements,
242                vec![
243                    (b"foo".to_vec(), b"bar".to_vec()),
244                    (b"ze".to_vec(), b"bra".to_vec()),
245                ]
246            );
247        }
248
249        // right unbounded (descending)
250        {
251            let iter = store.range(Some(b"f"), None, Order::Descending);
252            let elements: Vec<Pair> = iter.collect();
253            assert_eq!(
254                elements,
255                vec![
256                    (b"ze".to_vec(), b"bra".to_vec()),
257                    (b"foo".to_vec(), b"bar".to_vec()),
258                ]
259            );
260        }
261
262        // left unbounded
263        {
264            let iter = store.range(None, Some(b"f"), Order::Ascending);
265            let elements: Vec<Pair> = iter.collect();
266            assert_eq!(elements, vec![(b"ant".to_vec(), b"hill".to_vec()),]);
267        }
268
269        // left unbounded (descending)
270        {
271            let iter = store.range(None, Some(b"no"), Order::Descending);
272            let elements: Vec<Pair> = iter.collect();
273            assert_eq!(
274                elements,
275                vec![
276                    (b"foo".to_vec(), b"bar".to_vec()),
277                    (b"ant".to_vec(), b"hill".to_vec()),
278                ]
279            );
280        }
281    }
282
283    #[test]
284    fn memory_storage_implements_debug() {
285        let store = MemoryStorage::new();
286        assert_eq!(
287            format!("{:?}", store),
288            "MemoryStorage (0 entries) {\n\
289            }"
290        );
291
292        // With one element
293        let mut store = MemoryStorage::new();
294        store.set(&[0x00, 0xAB, 0xDD], &[0xFF, 0xD5]);
295        assert_eq!(
296            format!("{:?}", store),
297            "MemoryStorage (1 entries) {\n\
298            \x20\x200x00abdd: 0xffd5\n\
299            }"
300        );
301
302        // Sorted by key
303        let mut store = MemoryStorage::new();
304        store.set(&[0x00, 0xAB, 0xDD], &[0xFF, 0xD5]);
305        store.set(&[0x00, 0xAB, 0xEE], &[0xFF, 0xD5]);
306        store.set(&[0x00, 0xAB, 0xCC], &[0xFF, 0xD5]);
307        assert_eq!(
308            format!("{:?}", store),
309            "MemoryStorage (3 entries) {\n\
310            \x20\x200x00abcc: 0xffd5\n\
311            \x20\x200x00abdd: 0xffd5\n\
312            \x20\x200x00abee: 0xffd5\n\
313            }"
314        );
315
316        // Different lengths
317        let mut store = MemoryStorage::new();
318        store.set(&[0xAA], &[0x11]);
319        store.set(&[0xAA, 0xBB], &[0x11, 0x22]);
320        store.set(&[0xAA, 0xBB, 0xCC], &[0x11, 0x22, 0x33]);
321        store.set(&[0xAA, 0xBB, 0xCC, 0xDD], &[0x11, 0x22, 0x33, 0x44]);
322        assert_eq!(
323            format!("{:?}", store),
324            "MemoryStorage (4 entries) {\n\
325            \x20\x200xaa: 0x11\n\
326            \x20\x200xaabb: 0x1122\n\
327            \x20\x200xaabbcc: 0x112233\n\
328            \x20\x200xaabbccdd: 0x11223344\n\
329            }"
330        );
331    }
332}