cosmwasm_vm/testing/
storage.rs

1use std::collections::BTreeMap;
2#[cfg(feature = "iterator")]
3use std::collections::HashMap;
4#[cfg(feature = "iterator")]
5use std::ops::{Bound, RangeBounds};
6
7#[cfg(feature = "iterator")]
8use cosmwasm_std::{Order, Record};
9
10#[cfg(feature = "iterator")]
11use crate::BackendError;
12use crate::{BackendResult, GasInfo, Storage};
13
14#[cfg(feature = "iterator")]
15const GAS_COST_LAST_ITERATION: u64 = 37;
16
17#[cfg(feature = "iterator")]
18const GAS_COST_RANGE: u64 = 11;
19
20#[cfg(feature = "iterator")]
21#[derive(Default, Debug)]
22struct Iter {
23    data: Vec<Record>,
24    position: usize,
25}
26
27#[derive(Default, Debug)]
28pub struct MockStorage {
29    data: BTreeMap<Vec<u8>, Vec<u8>>,
30    #[cfg(feature = "iterator")]
31    iterators: HashMap<u32, Iter>,
32}
33
34impl MockStorage {
35    pub fn new() -> Self {
36        MockStorage::default()
37    }
38
39    #[cfg(feature = "iterator")]
40    pub fn all(&mut self, iterator_id: u32) -> BackendResult<Vec<Record>> {
41        let mut out: Vec<Record> = Vec::new();
42        let mut total = GasInfo::free();
43        loop {
44            let (result, info) = self.next(iterator_id);
45            total += info;
46            match result {
47                Err(err) => return (Err(err), total),
48                Ok(ok) => {
49                    if let Some(v) = ok {
50                        out.push(v);
51                    } else {
52                        break;
53                    }
54                }
55            }
56        }
57        (Ok(out), total)
58    }
59}
60
61impl Storage for MockStorage {
62    fn get(&self, key: &[u8]) -> BackendResult<Option<Vec<u8>>> {
63        let gas_info = GasInfo::with_externally_used(key.len() as u64);
64        (Ok(self.data.get(key).cloned()), gas_info)
65    }
66
67    #[cfg(feature = "iterator")]
68    fn scan(
69        &mut self,
70        start: Option<&[u8]>,
71        end: Option<&[u8]>,
72        order: Order,
73    ) -> BackendResult<u32> {
74        let gas_info = GasInfo::with_externally_used(GAS_COST_RANGE);
75        let bounds = range_bounds(start, end);
76
77        let values: Vec<Record> = match (bounds.start_bound(), bounds.end_bound()) {
78            // BTreeMap.range panics if range is start > end.
79            // However, this cases represent just empty range and we treat it as such.
80            (Bound::Included(start), Bound::Excluded(end)) if start > end => Vec::new(),
81            _ => match order {
82                Order::Ascending => self.data.range(bounds).map(clone_item).collect(),
83                Order::Descending => self.data.range(bounds).rev().map(clone_item).collect(),
84            },
85        };
86
87        let last_id: u32 = self
88            .iterators
89            .len()
90            .try_into()
91            .expect("Found more iterator IDs than supported");
92        let new_id = last_id + 1;
93        let iter = Iter {
94            data: values,
95            position: 0,
96        };
97        self.iterators.insert(new_id, iter);
98
99        (Ok(new_id), gas_info)
100    }
101
102    #[cfg(feature = "iterator")]
103    fn next(&mut self, iterator_id: u32) -> BackendResult<Option<Record>> {
104        let iterator = match self.iterators.get_mut(&iterator_id) {
105            Some(i) => i,
106            None => {
107                return (
108                    Err(BackendError::iterator_does_not_exist(iterator_id)),
109                    GasInfo::free(),
110                )
111            }
112        };
113
114        let (value, gas_info): (Option<Record>, GasInfo) =
115            if iterator.data.len() > iterator.position {
116                let item = iterator.data[iterator.position].clone();
117                iterator.position += 1;
118                let gas_cost = (item.0.len() + item.1.len()) as u64;
119                (Some(item), GasInfo::with_cost(gas_cost))
120            } else {
121                (None, GasInfo::with_externally_used(GAS_COST_LAST_ITERATION))
122            };
123
124        (Ok(value), gas_info)
125    }
126
127    fn set(&mut self, key: &[u8], value: &[u8]) -> BackendResult<()> {
128        self.data.insert(key.to_vec(), value.to_vec());
129        let gas_info = GasInfo::with_externally_used((key.len() + value.len()) as u64);
130        (Ok(()), gas_info)
131    }
132
133    fn remove(&mut self, key: &[u8]) -> BackendResult<()> {
134        self.data.remove(key);
135        let gas_info = GasInfo::with_externally_used(key.len() as u64);
136        (Ok(()), gas_info)
137    }
138}
139
140#[cfg(feature = "iterator")]
141fn range_bounds(start: Option<&[u8]>, end: Option<&[u8]>) -> impl RangeBounds<Vec<u8>> {
142    (
143        start.map_or(Bound::Unbounded, |x| Bound::Included(x.to_vec())),
144        end.map_or(Bound::Unbounded, |x| Bound::Excluded(x.to_vec())),
145    )
146}
147
148#[cfg(feature = "iterator")]
149/// The BTreeMap specific key-value pair reference type, as returned by BTreeMap<Vec<u8>, Vec<u8>>::range.
150/// This is internal as it can change any time if the map implementation is swapped out.
151type BTreeMapRecordRef<'a> = (&'a Vec<u8>, &'a Vec<u8>);
152
153#[cfg(feature = "iterator")]
154fn clone_item(item_ref: BTreeMapRecordRef) -> Record {
155    let (key, value) = item_ref;
156    (key.clone(), value.clone())
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn get_and_set() {
165        let mut store = MockStorage::new();
166        assert_eq!(None, store.get(b"foo").0.unwrap());
167        store.set(b"foo", b"bar").0.unwrap();
168        assert_eq!(Some(b"bar".to_vec()), store.get(b"foo").0.unwrap());
169        assert_eq!(None, store.get(b"food").0.unwrap());
170    }
171
172    #[test]
173    fn delete() {
174        let mut store = MockStorage::new();
175        store.set(b"foo", b"bar").0.unwrap();
176        store.set(b"food", b"bank").0.unwrap();
177        store.remove(b"foo").0.unwrap();
178
179        assert_eq!(None, store.get(b"foo").0.unwrap());
180        assert_eq!(Some(b"bank".to_vec()), store.get(b"food").0.unwrap());
181    }
182
183    #[test]
184    #[cfg(feature = "iterator")]
185    fn iterator() {
186        let mut store = MockStorage::new();
187        store.set(b"foo", b"bar").0.expect("error setting value");
188
189        // ensure we had previously set "foo" = "bar"
190        assert_eq!(store.get(b"foo").0.unwrap(), Some(b"bar".to_vec()));
191        let iter_id = store.scan(None, None, Order::Ascending).0.unwrap();
192        assert_eq!(store.all(iter_id).0.unwrap().len(), 1);
193
194        // setup - add some data, and delete part of it as well
195        store.set(b"ant", b"hill").0.expect("error setting value");
196        store.set(b"ze", b"bra").0.expect("error setting value");
197
198        // noise that should be ignored
199        store.set(b"bye", b"bye").0.expect("error setting value");
200        store.remove(b"bye").0.expect("error removing key");
201
202        // unbounded
203        {
204            let iter_id = store.scan(None, None, Order::Ascending).0.unwrap();
205            let elements = store.all(iter_id).0.unwrap();
206            assert_eq!(
207                elements,
208                vec![
209                    (b"ant".to_vec(), b"hill".to_vec()),
210                    (b"foo".to_vec(), b"bar".to_vec()),
211                    (b"ze".to_vec(), b"bra".to_vec()),
212                ]
213            );
214        }
215
216        // unbounded (descending)
217        {
218            let iter_id = store.scan(None, None, Order::Descending).0.unwrap();
219            let elements = store.all(iter_id).0.unwrap();
220            assert_eq!(
221                elements,
222                vec![
223                    (b"ze".to_vec(), b"bra".to_vec()),
224                    (b"foo".to_vec(), b"bar".to_vec()),
225                    (b"ant".to_vec(), b"hill".to_vec()),
226                ]
227            );
228        }
229
230        // bounded
231        {
232            let iter_id = store
233                .scan(Some(b"f"), Some(b"n"), Order::Ascending)
234                .0
235                .unwrap();
236            let elements = store.all(iter_id).0.unwrap();
237            assert_eq!(elements, vec![(b"foo".to_vec(), b"bar".to_vec())]);
238        }
239
240        // bounded (descending)
241        {
242            let iter_id = store
243                .scan(Some(b"air"), Some(b"loop"), Order::Descending)
244                .0
245                .unwrap();
246            let elements = store.all(iter_id).0.unwrap();
247            assert_eq!(
248                elements,
249                vec![
250                    (b"foo".to_vec(), b"bar".to_vec()),
251                    (b"ant".to_vec(), b"hill".to_vec()),
252                ]
253            );
254        }
255
256        // bounded empty [a, a)
257        {
258            let iter_id = store
259                .scan(Some(b"foo"), Some(b"foo"), Order::Ascending)
260                .0
261                .unwrap();
262            let elements = store.all(iter_id).0.unwrap();
263            assert_eq!(elements, vec![]);
264        }
265
266        // bounded empty [a, a) (descending)
267        {
268            let iter_id = store
269                .scan(Some(b"foo"), Some(b"foo"), Order::Descending)
270                .0
271                .unwrap();
272            let elements = store.all(iter_id).0.unwrap();
273            assert_eq!(elements, vec![]);
274        }
275
276        // bounded empty [a, b) with b < a
277        {
278            let iter_id = store
279                .scan(Some(b"z"), Some(b"a"), Order::Ascending)
280                .0
281                .unwrap();
282            let elements = store.all(iter_id).0.unwrap();
283            assert_eq!(elements, vec![]);
284        }
285
286        // bounded empty [a, b) with b < a (descending)
287        {
288            let iter_id = store
289                .scan(Some(b"z"), Some(b"a"), Order::Descending)
290                .0
291                .unwrap();
292            let elements = store.all(iter_id).0.unwrap();
293            assert_eq!(elements, vec![]);
294        }
295
296        // right unbounded
297        {
298            let iter_id = store.scan(Some(b"f"), None, Order::Ascending).0.unwrap();
299            let elements = store.all(iter_id).0.unwrap();
300            assert_eq!(
301                elements,
302                vec![
303                    (b"foo".to_vec(), b"bar".to_vec()),
304                    (b"ze".to_vec(), b"bra".to_vec()),
305                ]
306            );
307        }
308
309        // right unbounded (descending)
310        {
311            let iter_id = store.scan(Some(b"f"), None, Order::Descending).0.unwrap();
312            let elements = store.all(iter_id).0.unwrap();
313            assert_eq!(
314                elements,
315                vec![
316                    (b"ze".to_vec(), b"bra".to_vec()),
317                    (b"foo".to_vec(), b"bar".to_vec()),
318                ]
319            );
320        }
321
322        // left unbounded
323        {
324            let iter_id = store.scan(None, Some(b"f"), Order::Ascending).0.unwrap();
325            let elements = store.all(iter_id).0.unwrap();
326            assert_eq!(elements, vec![(b"ant".to_vec(), b"hill".to_vec()),]);
327        }
328
329        // left unbounded (descending)
330        {
331            let iter_id = store.scan(None, Some(b"no"), Order::Descending).0.unwrap();
332            let elements = store.all(iter_id).0.unwrap();
333            assert_eq!(
334                elements,
335                vec![
336                    (b"foo".to_vec(), b"bar".to_vec()),
337                    (b"ant".to_vec(), b"hill".to_vec()),
338                ]
339            );
340        }
341    }
342}
cosmwasm_vm/testing/storage.rs

cosmwasm_vm/testing/
storage.rs
