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 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 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
68impl 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")]
99type 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 assert_eq!(store.get(b"foo"), Some(b"bar".to_vec()));
150 assert_eq!(store.range(None, None, Order::Ascending).count(), 1);
151
152 store.set(b"ant", b"hill");
154 store.set(b"ze", b"bra");
155
156 store.set(b"bye", b"bye");
158 store.remove(b"bye");
159
160 {
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 {
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 {
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 {
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 {
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 {
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 {
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 {
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 {
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 {
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 {
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 {
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 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 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 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}