clone_cw_multi_test/wasm_emulation/storage/
mock_storage.rs

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

clone_cw_multi_test/wasm_emulation/storage/
mock_storage.rs
