secret_toolkit_storage/
keymap.rs

1use std::any::type_name;
2use std::collections::HashMap;
3use std::convert::TryInto;
4use std::marker::PhantomData;
5use std::sync::Mutex;
6
7use serde::Deserialize;
8use serde::{de::DeserializeOwned, Serialize};
9
10use cosmwasm_std::{StdError, StdResult, Storage};
11use cosmwasm_storage::to_length_prefixed;
12
13use secret_toolkit_serialization::{Bincode2, Serde};
14
15use crate::{IterOption, WithIter, WithoutIter};
16
17const INDEXES: &[u8] = b"indexes";
18const MAP_LENGTH: &[u8] = b"length";
19
20const DEFAULT_PAGE_SIZE: u32 = 1;
21
22#[derive(Serialize, Deserialize)]
23struct InternalItem<T, Ser>
24where
25    T: Serialize + DeserializeOwned,
26    Ser: Serde,
27{
28    item_vec: Vec<u8>,
29    // only Some if we enabled iterator
30    index_pos: Option<u32>,
31    item_type: PhantomData<T>,
32    serialization_type: PhantomData<Ser>,
33}
34
35impl<T: Serialize + DeserializeOwned, Ser: Serde> InternalItem<T, Ser> {
36    fn new(index_pos: Option<u32>, item: &T) -> StdResult<Self> {
37        Ok(Self {
38            item_vec: Ser::serialize(item)?,
39            index_pos,
40            item_type: PhantomData,
41            serialization_type: PhantomData,
42        })
43    }
44
45    fn get_item(&self) -> StdResult<T> {
46        Ser::deserialize(&self.item_vec)
47    }
48}
49
50pub struct KeymapBuilder<'a, K, T, Ser = Bincode2, I = WithIter> {
51    /// namespace of the newly constructed Storage
52    namespace: &'a [u8],
53    page_size: u32,
54    key_type: PhantomData<K>,
55    item_type: PhantomData<T>,
56    serialization_type: PhantomData<Ser>,
57    iter_option: PhantomData<I>,
58}
59
60impl<'a, K, T, Ser> KeymapBuilder<'a, K, T, Ser, WithIter>
61where
62    K: Serialize + DeserializeOwned,
63    T: Serialize + DeserializeOwned,
64    Ser: Serde,
65{
66    /// Creates a KeymapBuilder with default features
67    pub const fn new(namespace: &'a [u8]) -> Self {
68        Self {
69            namespace,
70            page_size: DEFAULT_PAGE_SIZE,
71            key_type: PhantomData,
72            item_type: PhantomData,
73            serialization_type: PhantomData,
74            iter_option: PhantomData,
75        }
76    }
77    /// Modifies the number of keys stored in one page of indexing, for the iterator
78    pub const fn with_page_size(&self, indexes_size: u32) -> Self {
79        if indexes_size == 0 {
80            panic!("zero index page size used in keymap")
81        }
82        Self {
83            namespace: self.namespace,
84            page_size: indexes_size,
85            key_type: self.key_type,
86            item_type: self.item_type,
87            serialization_type: self.serialization_type,
88            iter_option: self.iter_option,
89        }
90    }
91    /// Disables the iterator of the keymap, saving at least 4000 gas in each insertion.
92    pub const fn without_iter(&self) -> KeymapBuilder<'a, K, T, Ser, WithoutIter> {
93        KeymapBuilder {
94            namespace: self.namespace,
95            page_size: self.page_size,
96            key_type: PhantomData,
97            item_type: PhantomData,
98            serialization_type: PhantomData,
99            iter_option: PhantomData,
100        }
101    }
102    /// Returns a keymap with the given configuration
103    pub const fn build(&self) -> Keymap<'a, K, T, Ser, WithIter> {
104        Keymap {
105            namespace: self.namespace,
106            prefix: None,
107            page_size: self.page_size,
108            length: Mutex::new(None),
109            key_type: self.key_type,
110            item_type: self.item_type,
111            iter_option: self.iter_option,
112            serialization_type: self.serialization_type,
113        }
114    }
115}
116
117// This enables writing `.iter().skip(n).rev()`
118impl<'a, K, T, Ser> KeymapBuilder<'a, K, T, Ser, WithoutIter>
119where
120    K: Serialize + DeserializeOwned,
121    T: Serialize + DeserializeOwned,
122    Ser: Serde,
123{
124    pub const fn build(&self) -> Keymap<'a, K, T, Ser, WithoutIter> {
125        Keymap {
126            namespace: self.namespace,
127            prefix: None,
128            page_size: self.page_size,
129            length: Mutex::new(None),
130            key_type: self.key_type,
131            item_type: self.item_type,
132            iter_option: self.iter_option,
133            serialization_type: self.serialization_type,
134        }
135    }
136}
137
138pub struct Keymap<'a, K, T, Ser = Bincode2, I = WithIter>
139where
140    K: Serialize + DeserializeOwned,
141    T: Serialize + DeserializeOwned,
142    Ser: Serde,
143    I: IterOption,
144{
145    /// prefix of the newly constructed Storage
146    namespace: &'a [u8],
147    /// needed if any suffixes were added to the original namespace.
148    prefix: Option<Vec<u8>>,
149    page_size: u32,
150    length: Mutex<Option<u32>>,
151    key_type: PhantomData<K>,
152    item_type: PhantomData<T>,
153    iter_option: PhantomData<I>,
154    serialization_type: PhantomData<Ser>,
155}
156
157impl<'a, K: Serialize + DeserializeOwned, T: Serialize + DeserializeOwned, Ser: Serde>
158    Keymap<'a, K, T, Ser>
159{
160    /// constructor
161    pub const fn new(namespace: &'a [u8]) -> Self {
162        Self {
163            namespace,
164            prefix: None,
165            page_size: DEFAULT_PAGE_SIZE,
166            length: Mutex::new(None),
167            key_type: PhantomData,
168            item_type: PhantomData,
169            serialization_type: PhantomData,
170            iter_option: PhantomData,
171        }
172    }
173
174    /// This is used to produce a new Keymap. This can be used when you want to associate an Keymap to each user
175    /// and you still get to define the Keymap as a static constant
176    pub fn add_suffix(&self, suffix: &[u8]) -> Self {
177        let suffix = to_length_prefixed(suffix);
178        let prefix = self.prefix.as_deref().unwrap_or(self.namespace);
179        let prefix = [prefix, suffix.as_slice()].concat();
180        Self {
181            namespace: self.namespace,
182            prefix: Some(prefix),
183            page_size: self.page_size,
184            length: Mutex::new(None),
185            key_type: self.key_type,
186            item_type: self.item_type,
187            serialization_type: self.serialization_type,
188            iter_option: self.iter_option,
189        }
190    }
191}
192
193impl<K: Serialize + DeserializeOwned, T: Serialize + DeserializeOwned, Ser: Serde>
194    Keymap<'_, K, T, Ser, WithoutIter>
195{
196    /// Serialize key
197    fn serialize_key(&self, key: &K) -> StdResult<Vec<u8>> {
198        Ser::serialize(key)
199    }
200
201    /// user facing get function
202    pub fn get(&self, storage: &dyn Storage, key: &K) -> Option<T> {
203        self.get_from_key(storage, key).ok()
204    }
205
206    /// internal item get function
207    fn get_from_key(&self, storage: &dyn Storage, key: &K) -> StdResult<T> {
208        let key_vec = self.serialize_key(key)?;
209        self.load_impl(storage, &key_vec)
210    }
211
212    /// user facing remove function
213    pub fn remove(&self, storage: &mut dyn Storage, key: &K) -> StdResult<()> {
214        let key_vec = self.serialize_key(key)?;
215        self.remove_impl(storage, &key_vec);
216
217        Ok(())
218    }
219
220    /// user facing insert function
221    pub fn insert(&self, storage: &mut dyn Storage, key: &K, item: &T) -> StdResult<()> {
222        let key_vec = self.serialize_key(key)?;
223        self.save_impl(storage, &key_vec, item)
224    }
225
226    /// user facing method that checks if any item is stored with this key.
227    pub fn contains(&self, storage: &dyn Storage, key: &K) -> bool {
228        match self.serialize_key(key) {
229            Ok(key_vec) => self.contains_impl(storage, &key_vec),
230            Err(_) => false,
231        }
232    }
233}
234
235impl<'a, K: Serialize + DeserializeOwned, T: Serialize + DeserializeOwned, Ser: Serde>
236    Keymap<'a, K, T, Ser, WithIter>
237{
238    /// Serialize key
239    fn serialize_key(&self, key: &K) -> StdResult<Vec<u8>> {
240        Ser::serialize(key)
241    }
242
243    /// Deserialize key
244    fn deserialize_key(&self, key_data: &[u8]) -> StdResult<K> {
245        Ser::deserialize(key_data)
246    }
247
248    fn page_from_position(&self, position: u32) -> u32 {
249        position / self.page_size
250    }
251
252    /// get total number of objects saved
253    pub fn get_len(&self, storage: &dyn Storage) -> StdResult<u32> {
254        let mut may_len = self.length.lock().unwrap();
255        match *may_len {
256            Some(length) => Ok(length),
257            None => {
258                let len_key = [self.as_slice(), MAP_LENGTH].concat();
259                if let Some(len_vec) = storage.get(&len_key) {
260                    let len_bytes = len_vec
261                        .as_slice()
262                        .try_into()
263                        .map_err(|err| StdError::parse_err("u32", err))?;
264                    let len = u32::from_be_bytes(len_bytes);
265                    *may_len = Some(len);
266                    Ok(len)
267                } else {
268                    *may_len = Some(0);
269                    Ok(0)
270                }
271            }
272        }
273    }
274
275    /// checks if the collection has any elements
276    pub fn is_empty(&self, storage: &dyn Storage) -> StdResult<bool> {
277        Ok(self.get_len(storage)? == 0)
278    }
279
280    /// set length of the map
281    fn set_len(&self, storage: &mut dyn Storage, len: u32) -> StdResult<()> {
282        let len_key = [self.as_slice(), MAP_LENGTH].concat();
283        storage.set(&len_key, &len.to_be_bytes());
284
285        let mut may_len = self.length.lock().unwrap();
286        *may_len = Some(len);
287
288        Ok(())
289    }
290
291    /// Used to get the indexes stored in the given page number
292    fn get_indexes(&self, storage: &dyn Storage, page: u32) -> StdResult<Vec<Vec<u8>>> {
293        let indexes_key = [self.as_slice(), INDEXES, page.to_be_bytes().as_slice()].concat();
294        if self.page_size == 1 {
295            let maybe_item_data = storage.get(&indexes_key);
296            match maybe_item_data {
297                Some(item_data) => Ok(vec![item_data]),
298                None => Ok(vec![]),
299            }
300        } else {
301            let maybe_serialized = storage.get(&indexes_key);
302            match maybe_serialized {
303                Some(serialized) => Bincode2::deserialize(&serialized),
304                None => Ok(vec![]),
305            }
306        }
307    }
308
309    /// Set an indexes page
310    fn set_indexes_page(
311        &self,
312        storage: &mut dyn Storage,
313        page: u32,
314        indexes: &Vec<Vec<u8>>,
315    ) -> StdResult<()> {
316        let indexes_key = [self.as_slice(), INDEXES, page.to_be_bytes().as_slice()].concat();
317        if self.page_size == 1 {
318            if let Some(item_data) = indexes.first() {
319                storage.set(&indexes_key, item_data);
320            } else {
321                storage.remove(&indexes_key);
322            }
323        } else {
324            storage.set(&indexes_key, &Bincode2::serialize(indexes)?);
325        }
326        Ok(())
327    }
328
329    /// user facing get function
330    pub fn get(&self, storage: &dyn Storage, key: &K) -> Option<T> {
331        if let Ok(internal_item) = self.get_from_key(storage, key) {
332            internal_item.get_item().ok()
333        } else {
334            None
335        }
336    }
337
338    /// internal item get function
339    fn get_from_key(&self, storage: &dyn Storage, key: &K) -> StdResult<InternalItem<T, Ser>> {
340        let key_vec = self.serialize_key(key)?;
341        self.load_impl(storage, &key_vec)
342    }
343
344    /// user facing remove function
345    pub fn remove(&self, storage: &mut dyn Storage, key: &K) -> StdResult<()> {
346        let key_vec = self.serialize_key(key)?;
347
348        let removed_pos = self.get_from_key(storage, key)?.index_pos.unwrap();
349
350        let page = self.page_from_position(removed_pos);
351
352        let mut len = self.get_len(storage)?;
353        len -= 1;
354        self.set_len(storage, len)?;
355
356        let mut indexes = self.get_indexes(storage, page)?;
357
358        let pos_in_indexes = (removed_pos % self.page_size) as usize;
359
360        if indexes[pos_in_indexes] != key_vec {
361            return Err(StdError::generic_err(
362                "tried to remove from keymap, but key not found in indexes - should never happen",
363            ));
364        }
365
366        // if our object is the last item, then just remove it
367        if len == 0 || len == removed_pos {
368            indexes.pop();
369            self.set_indexes_page(storage, page, &indexes)?;
370            self.remove_impl(storage, &key_vec);
371            return Ok(());
372        }
373
374        // max page should use previous_len - 1 which is exactly the current len
375        let max_page = self.page_from_position(len);
376        if max_page == page {
377            // last page indexes is the same as indexes
378            let last_key = indexes.pop().ok_or_else(|| {
379                StdError::generic_err("last item's key not found - should never happen")
380            })?;
381            // modify last item
382            let mut last_internal_item = self.load_impl(storage, &last_key)?;
383            last_internal_item.index_pos = Some(removed_pos);
384            self.save_impl(storage, &last_key, &last_internal_item)?;
385            // save to indexes
386            indexes[pos_in_indexes] = last_key;
387            self.set_indexes_page(storage, page, &indexes)?;
388        } else {
389            let mut last_page_indexes = self.get_indexes(storage, max_page)?;
390            let last_key = last_page_indexes.pop().ok_or_else(|| {
391                StdError::generic_err("last item's key not found - should never happen")
392            })?;
393            // modify last item
394            let mut last_internal_item = self.load_impl(storage, &last_key)?;
395            last_internal_item.index_pos = Some(removed_pos);
396            self.save_impl(storage, &last_key, &last_internal_item)?;
397            // save indexes
398            indexes[pos_in_indexes] = last_key;
399            self.set_indexes_page(storage, page, &indexes)?;
400            self.set_indexes_page(storage, max_page, &last_page_indexes)?;
401        }
402
403        self.remove_impl(storage, &key_vec);
404
405        Ok(())
406    }
407
408    /// user facing insert function
409    pub fn insert(&self, storage: &mut dyn Storage, key: &K, item: &T) -> StdResult<()> {
410        let key_vec = self.serialize_key(key)?;
411
412        match self.may_load_impl(storage, &key_vec)? {
413            Some(existing_internal_item) => {
414                // if item already exists
415                let new_internal_item = InternalItem::new(existing_internal_item.index_pos, item)?;
416                self.save_impl(storage, &key_vec, &new_internal_item)
417            }
418            None => {
419                // not already saved
420                let pos = self.get_len(storage)?;
421                self.set_len(storage, pos + 1)?;
422                let page = self.page_from_position(pos);
423                // save the item
424                let internal_item = InternalItem::new(Some(pos), item)?;
425                self.save_impl(storage, &key_vec, &internal_item)?;
426                // add index
427                let mut indexes = self.get_indexes(storage, page)?;
428                indexes.push(key_vec);
429                self.set_indexes_page(storage, page, &indexes)
430            }
431        }
432    }
433
434    /// user facing method that checks if any item is stored with this key.
435    pub fn contains(&self, storage: &dyn Storage, key: &K) -> bool {
436        match self.serialize_key(key) {
437            Ok(key_vec) => self.contains_impl(storage, &key_vec),
438            Err(_) => false,
439        }
440    }
441
442    /// paginates (key, item) pairs.
443    pub fn paging(
444        &self,
445        storage: &dyn Storage,
446        start_page: u32,
447        size: u32,
448    ) -> StdResult<Vec<(K, T)>> {
449        let start_pos = start_page * size;
450
451        let max_size = self.get_len(storage)?;
452
453        if max_size == 0 {
454            return Ok(vec![]);
455        }
456
457        if start_pos > max_size {
458            return Err(StdError::not_found("out of bounds"));
459        }
460
461        self.iter(storage)?
462            .skip(start_pos as usize)
463            .take(size as usize)
464            .collect()
465    }
466
467    /// paginates only the keys. More efficient than paginating both items and keys
468    pub fn paging_keys(
469        &self,
470        storage: &dyn Storage,
471        start_page: u32,
472        size: u32,
473    ) -> StdResult<Vec<K>> {
474        let start_pos = start_page * size;
475
476        let max_size = self.get_len(storage)?;
477
478        if max_size == 0 {
479            return Ok(vec![]);
480        }
481
482        if start_pos > max_size {
483            return Err(StdError::not_found("out of bounds"));
484        }
485
486        self.iter_keys(storage)?
487            .skip(start_pos as usize)
488            .take(size as usize)
489            .collect()
490    }
491
492    /// Returns a readonly iterator only for keys. More efficient than iter().
493    pub fn iter_keys(&self, storage: &'a dyn Storage) -> StdResult<KeyIter<K, T, Ser>> {
494        let len = self.get_len(storage)?;
495        let iter = KeyIter::new(self, storage, 0, len);
496        Ok(iter)
497    }
498
499    /// Returns a readonly iterator for (key-item) pairs
500    pub fn iter(&self, storage: &'a dyn Storage) -> StdResult<KeyItemIter<K, T, Ser>> {
501        let len = self.get_len(storage)?;
502        let iter = KeyItemIter::new(self, storage, 0, len);
503        Ok(iter)
504    }
505}
506
507impl<K: Serialize + DeserializeOwned, T: Serialize + DeserializeOwned, Ser: Serde>
508    PrefixedTypedStorage<InternalItem<T, Ser>, Bincode2> for Keymap<'_, K, T, Ser, WithIter>
509{
510    fn as_slice(&self) -> &[u8] {
511        if let Some(prefix) = &self.prefix {
512            prefix
513        } else {
514            self.namespace
515        }
516    }
517}
518
519impl<K: Serialize + DeserializeOwned, T: Serialize + DeserializeOwned, Ser: Serde>
520    PrefixedTypedStorage<T, Ser> for Keymap<'_, K, T, Ser, WithoutIter>
521{
522    fn as_slice(&self) -> &[u8] {
523        if let Some(prefix) = &self.prefix {
524            prefix
525        } else {
526            self.namespace
527        }
528    }
529}
530
531/// An iterator over the keys of the Keymap.
532pub struct KeyIter<'a, K, T, Ser>
533where
534    K: Serialize + DeserializeOwned,
535    T: Serialize + DeserializeOwned,
536    Ser: Serde,
537{
538    keymap: &'a Keymap<'a, K, T, Ser>,
539    storage: &'a dyn Storage,
540    start: u32,
541    end: u32,
542    cache: HashMap<u32, Vec<Vec<u8>>>,
543}
544
545impl<'a, K, T, Ser> KeyIter<'a, K, T, Ser>
546where
547    K: Serialize + DeserializeOwned,
548    T: Serialize + DeserializeOwned,
549    Ser: Serde,
550{
551    /// constructor
552    pub fn new(
553        keymap: &'a Keymap<'a, K, T, Ser>,
554        storage: &'a dyn Storage,
555        start: u32,
556        end: u32,
557    ) -> Self {
558        Self {
559            keymap,
560            storage,
561            start,
562            end,
563            cache: HashMap::new(),
564        }
565    }
566}
567
568impl<K, T, Ser> Iterator for KeyIter<'_, K, T, Ser>
569where
570    K: Serialize + DeserializeOwned,
571    T: Serialize + DeserializeOwned,
572    Ser: Serde,
573{
574    type Item = StdResult<K>;
575
576    fn next(&mut self) -> Option<Self::Item> {
577        if self.start >= self.end {
578            return None;
579        }
580
581        let key;
582        let page = self.keymap.page_from_position(self.start);
583        let indexes_pos = (self.start % self.keymap.page_size) as usize;
584
585        match self.cache.get(&page) {
586            Some(indexes) => {
587                let key_data = &indexes[indexes_pos];
588                key = self.keymap.deserialize_key(key_data);
589            }
590            None => match self.keymap.get_indexes(self.storage, page) {
591                Ok(indexes) => {
592                    let key_data = &indexes[indexes_pos];
593                    key = self.keymap.deserialize_key(key_data);
594                    self.cache.insert(page, indexes);
595                }
596                Err(e) => key = Err(e),
597            },
598        }
599        self.start += 1;
600        Some(key)
601    }
602
603    // This needs to be implemented correctly for `ExactSizeIterator` to work.
604    fn size_hint(&self) -> (usize, Option<usize>) {
605        let len = (self.end - self.start) as usize;
606        (len, Some(len))
607    }
608
609    // I implement `nth` manually because it is used in the standard library whenever
610    // it wants to skip over elements, but the default implementation repeatedly calls next.
611    // because that is very expensive in this case, and the items are just discarded, we wan
612    // do better here.
613    // In practice, this enables cheap paging over the storage by calling:
614    // `.iter().skip(start).take(length).collect()`
615    fn nth(&mut self, n: usize) -> Option<Self::Item> {
616        self.start = self.start.saturating_add(n as u32);
617        self.next()
618    }
619}
620
621impl<K, T, Ser> DoubleEndedIterator for KeyIter<'_, K, T, Ser>
622where
623    K: Serialize + DeserializeOwned,
624    T: Serialize + DeserializeOwned,
625    Ser: Serde,
626{
627    fn next_back(&mut self) -> Option<Self::Item> {
628        if self.start >= self.end {
629            return None;
630        }
631        self.end -= 1;
632
633        let key;
634        let page = self.keymap.page_from_position(self.end);
635        let indexes_pos = (self.end % self.keymap.page_size) as usize;
636
637        match self.cache.get(&page) {
638            Some(indexes) => {
639                let key_data = &indexes[indexes_pos];
640                key = self.keymap.deserialize_key(key_data);
641            }
642            None => match self.keymap.get_indexes(self.storage, page) {
643                Ok(indexes) => {
644                    let key_data = &indexes[indexes_pos];
645                    key = self.keymap.deserialize_key(key_data);
646                    self.cache.insert(page, indexes);
647                }
648                Err(e) => key = Err(e),
649            },
650        }
651        Some(key)
652    }
653
654    // I implement `nth_back` manually because it is used in the standard library whenever
655    // it wants to skip over elements, but the default implementation repeatedly calls next_back.
656    // because that is very expensive in this case, and the items are just discarded, we wan
657    // do better here.
658    // In practice, this enables cheap paging over the storage by calling:
659    // `.iter().skip(start).take(length).collect()`
660    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
661        self.end = self.end.saturating_sub(n as u32);
662        self.next_back()
663    }
664}
665
666// This enables writing `.iter().skip(n).rev()`
667impl<K, T, Ser> ExactSizeIterator for KeyIter<'_, K, T, Ser>
668where
669    K: Serialize + DeserializeOwned,
670    T: Serialize + DeserializeOwned,
671    Ser: Serde,
672{
673}
674
675// ===============================================================================================
676
677/// An iterator over the (key, item) pairs of the Keymap. Less efficient than just iterating over keys.
678pub struct KeyItemIter<'a, K, T, Ser>
679where
680    K: Serialize + DeserializeOwned,
681    T: Serialize + DeserializeOwned,
682    Ser: Serde,
683{
684    keymap: &'a Keymap<'a, K, T, Ser>,
685    storage: &'a dyn Storage,
686    start: u32,
687    end: u32,
688    cache: HashMap<u32, Vec<Vec<u8>>>,
689}
690
691impl<'a, K, T, Ser> KeyItemIter<'a, K, T, Ser>
692where
693    K: Serialize + DeserializeOwned,
694    T: Serialize + DeserializeOwned,
695    Ser: Serde,
696{
697    /// constructor
698    pub fn new(
699        keymap: &'a Keymap<'a, K, T, Ser>,
700        storage: &'a dyn Storage,
701        start: u32,
702        end: u32,
703    ) -> Self {
704        Self {
705            keymap,
706            storage,
707            start,
708            end,
709            cache: HashMap::new(),
710        }
711    }
712}
713
714impl<K, T, Ser> Iterator for KeyItemIter<'_, K, T, Ser>
715where
716    K: Serialize + DeserializeOwned,
717    T: Serialize + DeserializeOwned,
718    Ser: Serde,
719{
720    type Item = StdResult<(K, T)>;
721
722    fn next(&mut self) -> Option<Self::Item> {
723        if self.start >= self.end {
724            return None;
725        }
726
727        let key;
728        let page = self.keymap.page_from_position(self.start);
729        let indexes_pos = (self.start % self.keymap.page_size) as usize;
730
731        match self.cache.get(&page) {
732            Some(indexes) => {
733                let key_data = &indexes[indexes_pos];
734                key = self.keymap.deserialize_key(key_data);
735            }
736            None => match self.keymap.get_indexes(self.storage, page) {
737                Ok(indexes) => {
738                    let key_data = &indexes[indexes_pos];
739                    key = self.keymap.deserialize_key(key_data);
740                    self.cache.insert(page, indexes);
741                }
742                Err(e) => key = Err(e),
743            },
744        }
745        self.start += 1;
746        // turn key into pair
747        let pair = match key {
748            Ok(k) => match self.keymap.get_from_key(self.storage, &k) {
749                Ok(internal_item) => match internal_item.get_item() {
750                    Ok(item) => Ok((k, item)),
751                    Err(e) => Err(e),
752                },
753                Err(e) => Err(e),
754            },
755            Err(e) => Err(e),
756        };
757        Some(pair)
758    }
759
760    // This needs to be implemented correctly for `ExactSizeIterator` to work.
761    fn size_hint(&self) -> (usize, Option<usize>) {
762        let len = (self.end - self.start) as usize;
763        (len, Some(len))
764    }
765
766    // I implement `nth` manually because it is used in the standard library whenever
767    // it wants to skip over elements, but the default implementation repeatedly calls next.
768    // because that is very expensive in this case, and the items are just discarded, we wan
769    // do better here.
770    // In practice, this enables cheap paging over the storage by calling:
771    // `.iter().skip(start).take(length).collect()`
772    fn nth(&mut self, n: usize) -> Option<Self::Item> {
773        self.start = self.start.saturating_add(n as u32);
774        self.next()
775    }
776}
777
778impl<K, T, Ser> DoubleEndedIterator for KeyItemIter<'_, K, T, Ser>
779where
780    K: Serialize + DeserializeOwned,
781    T: Serialize + DeserializeOwned,
782    Ser: Serde,
783{
784    fn next_back(&mut self) -> Option<Self::Item> {
785        if self.start >= self.end {
786            return None;
787        }
788        self.end -= 1;
789
790        let key;
791        let page = self.keymap.page_from_position(self.end);
792        let indexes_pos = (self.end % self.keymap.page_size) as usize;
793
794        match self.cache.get(&page) {
795            Some(indexes) => {
796                let key_data = &indexes[indexes_pos];
797                key = self.keymap.deserialize_key(key_data);
798            }
799            None => match self.keymap.get_indexes(self.storage, page) {
800                Ok(indexes) => {
801                    let key_data = &indexes[indexes_pos];
802                    key = self.keymap.deserialize_key(key_data);
803                    self.cache.insert(page, indexes);
804                }
805                Err(e) => key = Err(e),
806            },
807        }
808        // turn key into pair
809        let pair = match key {
810            Ok(k) => match self.keymap.get_from_key(self.storage, &k) {
811                Ok(internal_item) => match internal_item.get_item() {
812                    Ok(item) => Ok((k, item)),
813                    Err(e) => Err(e),
814                },
815                Err(e) => Err(e),
816            },
817            Err(e) => Err(e),
818        };
819        Some(pair)
820    }
821
822    // I implement `nth_back` manually because it is used in the standard library whenever
823    // it wants to skip over elements, but the default implementation repeatedly calls next_back.
824    // because that is very expensive in this case, and the items are just discarded, we wan
825    // do better here.
826    // In practice, this enables cheap paging over the storage by calling:
827    // `.iter().skip(start).take(length).collect()`
828    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
829        self.end = self.end.saturating_sub(n as u32);
830        self.next_back()
831    }
832}
833
834// This enables writing `.iter().skip(n).rev()`
835impl<K, T, Ser> ExactSizeIterator for KeyItemIter<'_, K, T, Ser>
836where
837    K: Serialize + DeserializeOwned,
838    T: Serialize + DeserializeOwned,
839    Ser: Serde,
840{
841}
842
843trait PrefixedTypedStorage<T: Serialize + DeserializeOwned, Ser: Serde> {
844    fn as_slice(&self) -> &[u8];
845
846    /// Returns bool from retrieving the item with the specified key.
847    ///
848    /// # Arguments
849    ///
850    /// * `storage` - a reference to the storage this item is in
851    /// * `key` - a byte slice representing the key to access the stored item
852    fn contains_impl(&self, storage: &dyn Storage, key: &[u8]) -> bool {
853        let prefixed_key = [self.as_slice(), key].concat();
854        storage.get(&prefixed_key).is_some()
855    }
856
857    /// Returns StdResult<T> from retrieving the item with the specified key.  Returns a
858    /// StdError::NotFound if there is no item with that key
859    ///
860    /// # Arguments
861    ///
862    /// * `storage` - a reference to the storage this item is in
863    /// * `key` - a byte slice representing the key to access the stored item
864    fn load_impl(&self, storage: &dyn Storage, key: &[u8]) -> StdResult<T> {
865        let prefixed_key = [self.as_slice(), key].concat();
866        Ser::deserialize(
867            &storage
868                .get(&prefixed_key)
869                .ok_or_else(|| StdError::not_found(type_name::<T>()))?,
870        )
871    }
872
873    /// Returns StdResult<Option<T>> from retrieving the item with the specified key.  Returns a
874    /// None if there is no item with that key
875    ///
876    /// # Arguments
877    ///
878    /// * `storage` - a reference to the storage this item is in
879    /// * `key` - a byte slice representing the key to access the stored item
880    fn may_load_impl(&self, storage: &dyn Storage, key: &[u8]) -> StdResult<Option<T>> {
881        let prefixed_key = [self.as_slice(), key].concat();
882        match storage.get(&prefixed_key) {
883            Some(value) => Ser::deserialize(&value).map(Some),
884            None => Ok(None),
885        }
886    }
887
888    /// Returns StdResult<()> resulting from saving an item to storage
889    ///
890    /// # Arguments
891    ///
892    /// * `storage` - a mutable reference to the storage this item should go to
893    /// * `key` - a byte slice representing the key to access the stored item
894    /// * `value` - a reference to the item to store
895    fn save_impl(&self, storage: &mut dyn Storage, key: &[u8], value: &T) -> StdResult<()> {
896        let prefixed_key = [self.as_slice(), key].concat();
897        storage.set(&prefixed_key, &Ser::serialize(value)?);
898        Ok(())
899    }
900
901    /// Removes an item from storage
902    ///
903    /// # Arguments
904    ///
905    /// * `storage` - a mutable reference to the storage this item is in
906    /// * `key` - a byte slice representing the key to access the stored item
907    fn remove_impl(&self, storage: &mut dyn Storage, key: &[u8]) {
908        let prefixed_key = [self.as_slice(), key].concat();
909        storage.remove(&prefixed_key);
910    }
911}
912
913#[cfg(test)]
914mod tests {
915    use secret_toolkit_serialization::Json;
916    use serde::{Deserialize, Serialize};
917
918    use cosmwasm_std::testing::MockStorage;
919
920    use super::*;
921
922    #[derive(Serialize, Deserialize, Eq, PartialEq, Debug, Clone)]
923    struct Foo {
924        string: String,
925        number: i32,
926    }
927    #[test]
928    fn test_keymap_perf_insert() -> StdResult<()> {
929        let mut storage = MockStorage::new();
930
931        let total_items: i32 = 1000;
932
933        let keymap: Keymap<Vec<u8>, i32> = Keymap::new(b"test");
934
935        for i in 0..total_items {
936            let key: Vec<u8> = i.to_be_bytes().to_vec();
937            keymap.insert(&mut storage, &key, &i)?;
938        }
939
940        assert_eq!(keymap.get_len(&storage)?, 1000);
941
942        Ok(())
943    }
944
945    #[test]
946    fn test_keymap_perf_insert_remove() -> StdResult<()> {
947        let mut storage = MockStorage::new();
948
949        let total_items = 100;
950
951        let keymap: Keymap<i32, i32> = Keymap::new(b"test");
952
953        for i in 0..total_items {
954            keymap.insert(&mut storage, &i, &i)?;
955        }
956
957        for i in 0..total_items {
958            keymap.remove(&mut storage, &i)?;
959        }
960
961        assert_eq!(keymap.get_len(&storage)?, 0);
962
963        Ok(())
964    }
965
966    #[test]
967    fn test_keymap_paging() -> StdResult<()> {
968        let mut storage = MockStorage::new();
969
970        let page_size: u32 = 5;
971        let total_items: u32 = 50;
972        let keymap: Keymap<Vec<u8>, u32> = Keymap::new(b"test");
973
974        for i in 0..total_items {
975            let key: Vec<u8> = (i as i32).to_be_bytes().to_vec();
976            keymap.insert(&mut storage, &key, &i)?;
977        }
978
979        for i in 0..((total_items / page_size) - 1) {
980            let start_page = i;
981
982            let values = keymap.paging(&storage, start_page, page_size)?;
983
984            for (index, (key_value, value)) in values.iter().enumerate() {
985                let i = page_size * start_page + index as u32;
986                let key: Vec<u8> = (i as i32).to_be_bytes().to_vec();
987                assert_eq!(key_value, &key);
988                assert_eq!(value, &i);
989            }
990        }
991
992        Ok(())
993    }
994
995    #[test]
996    fn test_keymap_paging_overflow() -> StdResult<()> {
997        let mut storage = MockStorage::new();
998
999        let page_size = 50;
1000        let total_items = 10;
1001        let keymap: Keymap<i32, u32> = Keymap::new(b"test");
1002
1003        for i in 0..total_items {
1004            keymap.insert(&mut storage, &(i as i32), &i)?;
1005        }
1006
1007        let values = keymap.paging_keys(&storage, 0, page_size)?;
1008
1009        assert_eq!(values.len(), total_items as usize);
1010
1011        for (index, value) in values.iter().enumerate() {
1012            assert_eq!(value, &(index as i32))
1013        }
1014
1015        Ok(())
1016    }
1017
1018    #[test]
1019    fn test_keymap_insert_multiple() -> StdResult<()> {
1020        let mut storage = MockStorage::new();
1021
1022        let keymap: Keymap<Vec<u8>, Foo> = Keymap::new(b"test");
1023        let foo1 = Foo {
1024            string: "string one".to_string(),
1025            number: 1111,
1026        };
1027        let foo2 = Foo {
1028            string: "string two".to_string(),
1029            number: 1111,
1030        };
1031
1032        keymap.insert(&mut storage, &b"key1".to_vec(), &foo1)?;
1033        keymap.insert(&mut storage, &b"key2".to_vec(), &foo2)?;
1034
1035        let read_foo1 = keymap.get(&storage, &b"key1".to_vec()).unwrap();
1036        let read_foo2 = keymap.get(&storage, &b"key2".to_vec()).unwrap();
1037
1038        assert_eq!(foo1, read_foo1);
1039        assert_eq!(foo2, read_foo2);
1040        Ok(())
1041    }
1042
1043    #[test]
1044    fn test_keymap_contains() -> StdResult<()> {
1045        let mut storage = MockStorage::new();
1046
1047        let keymap: Keymap<Vec<u8>, Foo> = Keymap::new(b"test");
1048        let foo1 = Foo {
1049            string: "string one".to_string(),
1050            number: 1111,
1051        };
1052
1053        keymap.insert(&mut storage, &b"key1".to_vec(), &foo1)?;
1054        let contains_k1 = keymap.contains(&storage, &b"key1".to_vec());
1055
1056        assert!(contains_k1);
1057
1058        Ok(())
1059    }
1060
1061    #[test]
1062    fn test_keymap_iter() -> StdResult<()> {
1063        let mut storage = MockStorage::new();
1064
1065        let keymap: Keymap<Vec<u8>, Foo> = Keymap::new(b"test");
1066        let foo1 = Foo {
1067            string: "string one".to_string(),
1068            number: 1111,
1069        };
1070        let foo2 = Foo {
1071            string: "string two".to_string(),
1072            number: 1111,
1073        };
1074
1075        keymap.insert(&mut storage, &b"key1".to_vec(), &foo1)?;
1076        keymap.insert(&mut storage, &b"key2".to_vec(), &foo2)?;
1077
1078        let mut x = keymap.iter(&storage)?;
1079        let (len, _) = x.size_hint();
1080        assert_eq!(len, 2);
1081
1082        assert_eq!(x.next().unwrap()?, (b"key1".to_vec(), foo1));
1083
1084        assert_eq!(x.next().unwrap()?, (b"key2".to_vec(), foo2));
1085
1086        Ok(())
1087    }
1088
1089    #[test]
1090    fn test_keymap_iter_keys() -> StdResult<()> {
1091        let mut storage = MockStorage::new();
1092
1093        let keymap: Keymap<String, Foo> = Keymap::new(b"test");
1094        let foo1 = Foo {
1095            string: "string one".to_string(),
1096            number: 1111,
1097        };
1098        let foo2 = Foo {
1099            string: "string two".to_string(),
1100            number: 1111,
1101        };
1102
1103        let key1 = "key1".to_string();
1104        let key2 = "key2".to_string();
1105
1106        keymap.insert(&mut storage, &key1, &foo1)?;
1107        keymap.insert(&mut storage, &key2, &foo2)?;
1108
1109        let mut x = keymap.iter_keys(&storage)?;
1110        let (len, _) = x.size_hint();
1111        assert_eq!(len, 2);
1112
1113        assert_eq!(x.next().unwrap()?, key1);
1114
1115        assert_eq!(x.next().unwrap()?, key2);
1116
1117        Ok(())
1118    }
1119
1120    #[test]
1121    fn test_keymap_overwrite() -> StdResult<()> {
1122        let mut storage = MockStorage::new();
1123
1124        let keymap: Keymap<Vec<u8>, Foo> = Keymap::new(b"test");
1125        let foo1 = Foo {
1126            string: "string one".to_string(),
1127            number: 1111,
1128        };
1129        let foo2 = Foo {
1130            string: "string two".to_string(),
1131            number: 2222,
1132        };
1133
1134        keymap.insert(&mut storage, &b"key1".to_vec(), &foo1)?;
1135        keymap.insert(&mut storage, &b"key1".to_vec(), &foo2)?;
1136
1137        let foo3 = keymap.get(&storage, &b"key1".to_vec()).unwrap();
1138
1139        assert_eq!(foo3, foo2);
1140
1141        Ok(())
1142    }
1143
1144    #[test]
1145    fn test_keymap_suffixed_basics() -> StdResult<()> {
1146        let mut storage = MockStorage::new();
1147
1148        let original_keymap: Keymap<String, Foo> = Keymap::new(b"test");
1149        let keymap = original_keymap.add_suffix(b"test_suffix");
1150        let foo1 = Foo {
1151            string: "string one".to_string(),
1152            number: 1111,
1153        };
1154        let foo2 = Foo {
1155            string: "string one".to_string(),
1156            number: 1111,
1157        };
1158        keymap.insert(&mut storage, &"key1".to_string(), &foo1)?;
1159        keymap.insert(&mut storage, &"key2".to_string(), &foo2)?;
1160
1161        let read_foo1 = keymap.get(&storage, &"key1".to_string()).unwrap();
1162        let read_foo2 = keymap.get(&storage, &"key2".to_string()).unwrap();
1163
1164        assert_eq!(original_keymap.get_len(&storage)?, 0);
1165        assert_eq!(foo1, read_foo1);
1166        assert_eq!(foo2, read_foo2);
1167
1168        let alternative_keymap: Keymap<String, Foo> = Keymap::new(b"alternative");
1169        let alt_same_suffix = alternative_keymap.add_suffix(b"test_suffix");
1170
1171        assert!(alt_same_suffix.is_empty(&storage)?);
1172
1173        // show that it loads foo1 before removal
1174        let before_remove_foo1 = keymap.get(&storage, &"key1".to_string());
1175        assert!(before_remove_foo1.is_some());
1176        assert_eq!(foo1, before_remove_foo1.unwrap());
1177        // and returns None after removal
1178        keymap.remove(&mut storage, &"key1".to_string())?;
1179        let removed_foo1 = keymap.get(&storage, &"key1".to_string());
1180        assert!(removed_foo1.is_none());
1181
1182        // show what happens when reading from keys that have not been set yet.
1183        assert!(keymap.get(&storage, &"key3".to_string()).is_none());
1184
1185        Ok(())
1186    }
1187
1188    #[test]
1189    fn test_keymap_length() -> StdResult<()> {
1190        test_keymap_length_with_page_size(1)?;
1191        test_keymap_length_with_page_size(5)?;
1192        test_keymap_length_with_page_size(13)?;
1193        Ok(())
1194    }
1195
1196    fn test_keymap_length_with_page_size(page_size: u32) -> StdResult<()> {
1197        let mut storage = MockStorage::new();
1198
1199        let keymap: Keymap<String, Foo> = KeymapBuilder::new(b"test")
1200            .with_page_size(page_size)
1201            .build();
1202        let foo1 = Foo {
1203            string: "string one".to_string(),
1204            number: 1111,
1205        };
1206        let foo2 = Foo {
1207            string: "string one".to_string(),
1208            number: 1111,
1209        };
1210
1211        assert!(keymap.length.lock().unwrap().eq(&None));
1212        assert_eq!(keymap.get_len(&storage)?, 0);
1213        assert!(keymap.length.lock().unwrap().eq(&Some(0)));
1214
1215        let key1 = "k1".to_string();
1216        let key2 = "k2".to_string();
1217
1218        keymap.insert(&mut storage, &key1, &foo1)?;
1219        assert_eq!(keymap.get_len(&storage)?, 1);
1220        assert!(keymap.length.lock().unwrap().eq(&Some(1)));
1221
1222        // add another item
1223        keymap.insert(&mut storage, &key2, &foo2)?;
1224        assert_eq!(keymap.get_len(&storage)?, 2);
1225        assert!(keymap.length.lock().unwrap().eq(&Some(2)));
1226
1227        // remove item and check length
1228        keymap.remove(&mut storage, &key1)?;
1229        assert_eq!(keymap.get_len(&storage)?, 1);
1230        assert!(keymap.length.lock().unwrap().eq(&Some(1)));
1231
1232        // override item (should not change length)
1233        keymap.insert(&mut storage, &key2, &foo1)?;
1234        assert_eq!(keymap.get_len(&storage)?, 1);
1235        assert!(keymap.length.lock().unwrap().eq(&Some(1)));
1236
1237        // remove item and check length
1238        keymap.remove(&mut storage, &key2)?;
1239        assert_eq!(keymap.get_len(&storage)?, 0);
1240        assert!(keymap.length.lock().unwrap().eq(&Some(0)));
1241
1242        Ok(())
1243    }
1244
1245    #[test]
1246    fn test_keymap_without_iter() -> StdResult<()> {
1247        test_keymap_without_iter_custom_page(1)?;
1248        test_keymap_without_iter_custom_page(2)?;
1249        test_keymap_without_iter_custom_page(3)?;
1250        Ok(())
1251    }
1252
1253    fn test_keymap_without_iter_custom_page(page_size: u32) -> StdResult<()> {
1254        let mut storage = MockStorage::new();
1255
1256        let keymap: Keymap<String, Foo, Json, _> = KeymapBuilder::new(b"test")
1257            .with_page_size(page_size)
1258            .without_iter()
1259            .build();
1260
1261        let foo1 = Foo {
1262            string: "string one".to_string(),
1263            number: 1111,
1264        };
1265        let foo2 = Foo {
1266            string: "string one".to_string(),
1267            number: 1111,
1268        };
1269        keymap.insert(&mut storage, &"key1".to_string(), &foo1)?;
1270        keymap.insert(&mut storage, &"key2".to_string(), &foo2)?;
1271
1272        let read_foo1 = keymap.get(&storage, &"key1".to_string()).unwrap();
1273        let read_foo2 = keymap.get(&storage, &"key2".to_string()).unwrap();
1274
1275        assert_eq!(foo1, read_foo1);
1276        assert_eq!(foo2, read_foo2);
1277        assert!(keymap.contains(&storage, &"key1".to_string()));
1278
1279        keymap.remove(&mut storage, &"key1".to_string())?;
1280
1281        let read_foo1 = keymap.get(&storage, &"key1".to_string());
1282        let read_foo2 = keymap.get(&storage, &"key2".to_string()).unwrap();
1283
1284        assert!(read_foo1.is_none());
1285        assert_eq!(foo2, read_foo2);
1286
1287        Ok(())
1288    }
1289
1290    #[test]
1291    fn test_keymap_custom_paging() -> StdResult<()> {
1292        let mut storage = MockStorage::new();
1293
1294        let page_size: u32 = 5;
1295        let total_items: u32 = 50;
1296        let keymap: Keymap<Vec<u8>, u32> = KeymapBuilder::new(b"test").with_page_size(13).build();
1297
1298        for i in 0..total_items {
1299            let key: Vec<u8> = (i as i32).to_be_bytes().to_vec();
1300            keymap.insert(&mut storage, &key, &i)?;
1301        }
1302
1303        for i in 0..((total_items / page_size) - 1) {
1304            let start_page = i;
1305
1306            let values = keymap.paging(&storage, start_page, page_size)?;
1307
1308            for (index, (key_value, value)) in values.iter().enumerate() {
1309                let i = page_size * start_page + index as u32;
1310                let key: Vec<u8> = (i as i32).to_be_bytes().to_vec();
1311                assert_eq!(key_value, &key);
1312                assert_eq!(value, &i);
1313            }
1314        }
1315
1316        Ok(())
1317    }
1318
1319    #[test]
1320    fn test_keymap_custom_paging_overflow() -> StdResult<()> {
1321        let mut storage = MockStorage::new();
1322
1323        let page_size = 50;
1324        let total_items = 10;
1325        let keymap: Keymap<i32, u32, Json> = KeymapBuilder::new(b"test").with_page_size(3).build();
1326
1327        for i in 0..total_items {
1328            keymap.insert(&mut storage, &(i as i32), &i)?;
1329        }
1330
1331        let values = keymap.paging_keys(&storage, 0, page_size)?;
1332
1333        assert_eq!(values.len(), total_items as usize);
1334
1335        for (index, value) in values.iter().enumerate() {
1336            assert_eq!(value, &(index as i32))
1337        }
1338
1339        Ok(())
1340    }
1341
1342    #[test]
1343    fn test_keymap_custom_page_iter() -> StdResult<()> {
1344        let mut storage = MockStorage::new();
1345
1346        let keymap: Keymap<Vec<u8>, Foo> = KeymapBuilder::new(b"test").with_page_size(2).build();
1347        let foo1 = Foo {
1348            string: "string one".to_string(),
1349            number: 1111,
1350        };
1351        let foo2 = Foo {
1352            string: "string two".to_string(),
1353            number: 1111,
1354        };
1355        let foo3 = Foo {
1356            string: "string three".to_string(),
1357            number: 1111,
1358        };
1359
1360        keymap.insert(&mut storage, &b"key1".to_vec(), &foo1)?;
1361        keymap.insert(&mut storage, &b"key2".to_vec(), &foo2)?;
1362        keymap.insert(&mut storage, &b"key3".to_vec(), &foo3)?;
1363
1364        let mut x = keymap.iter(&storage)?;
1365        let (len, _) = x.size_hint();
1366        assert_eq!(len, 3);
1367
1368        assert_eq!(x.next().unwrap()?, (b"key1".to_vec(), foo1));
1369
1370        assert_eq!(x.next().unwrap()?, (b"key2".to_vec(), foo2));
1371
1372        assert_eq!(x.next().unwrap()?, (b"key3".to_vec(), foo3));
1373
1374        assert_eq!(x.next(), None);
1375
1376        Ok(())
1377    }
1378
1379    #[test]
1380    fn test_keymap_reverse_iter() -> StdResult<()> {
1381        test_keymap_custom_page_reverse_iterator(1)?;
1382        test_keymap_custom_page_reverse_iterator(2)?;
1383        test_keymap_custom_page_reverse_iterator(5)?;
1384        test_keymap_custom_page_reverse_iterator(25)?;
1385        Ok(())
1386    }
1387
1388    fn test_keymap_custom_page_reverse_iterator(page_size: u32) -> StdResult<()> {
1389        let mut storage = MockStorage::new();
1390        let keymap: Keymap<i32, i32> = KeymapBuilder::new(b"test")
1391            .with_page_size(page_size)
1392            .build();
1393        keymap.insert(&mut storage, &1234, &1234)?;
1394        keymap.insert(&mut storage, &2143, &2143)?;
1395        keymap.insert(&mut storage, &3412, &3412)?;
1396        keymap.insert(&mut storage, &4321, &4321)?;
1397
1398        let mut iter = keymap.iter(&storage)?.rev();
1399        assert_eq!(iter.next(), Some(Ok((4321, 4321))));
1400        assert_eq!(iter.next(), Some(Ok((3412, 3412))));
1401        assert_eq!(iter.next(), Some(Ok((2143, 2143))));
1402        assert_eq!(iter.next(), Some(Ok((1234, 1234))));
1403        assert_eq!(iter.next(), None);
1404
1405        // iterate twice to make sure nothing changed
1406        let mut iter = keymap.iter(&storage)?.rev();
1407        assert_eq!(iter.next(), Some(Ok((4321, 4321))));
1408        assert_eq!(iter.next(), Some(Ok((3412, 3412))));
1409        assert_eq!(iter.next(), Some(Ok((2143, 2143))));
1410        assert_eq!(iter.next(), Some(Ok((1234, 1234))));
1411        assert_eq!(iter.next(), None);
1412
1413        // make sure our implementation of `nth_back` doesn't break anything
1414        let mut iter = keymap.iter(&storage)?.rev().skip(2);
1415        assert_eq!(iter.next(), Some(Ok((2143, 2143))));
1416        assert_eq!(iter.next(), Some(Ok((1234, 1234))));
1417        assert_eq!(iter.next(), None);
1418
1419        // make sure our implementation of `ExactSizeIterator` works well
1420        let mut iter = keymap.iter(&storage)?.skip(2).rev();
1421        assert_eq!(iter.next(), Some(Ok((4321, 4321))));
1422        assert_eq!(iter.next(), Some(Ok((3412, 3412))));
1423        assert_eq!(iter.next(), None);
1424
1425        Ok(())
1426    }
1427
1428    #[test]
1429    fn test_serializations() -> StdResult<()> {
1430        test_serializations_with_page_size(1)?;
1431        test_serializations_with_page_size(3)?;
1432        test_serializations_with_page_size(19)?;
1433        Ok(())
1434    }
1435
1436    fn test_serializations_with_page_size(page_size: u32) -> StdResult<()> {
1437        // Check the default behavior is Bincode2
1438        let mut storage = MockStorage::new();
1439
1440        let keymap: Keymap<i32, i32> = KeymapBuilder::new(b"test")
1441            .with_page_size(page_size)
1442            .build();
1443        keymap.insert(&mut storage, &1234, &1234)?;
1444
1445        let page_key = [keymap.as_slice(), INDEXES, &0_u32.to_be_bytes()].concat();
1446        if keymap.page_size == 1 {
1447            let item_data = storage.get(&page_key);
1448            let expected_data = Bincode2::serialize(&1234)?;
1449            assert_eq!(item_data, Some(expected_data));
1450        } else {
1451            let page_bytes = storage.get(&page_key);
1452            let expected_bincode2 = Bincode2::serialize(&vec![Bincode2::serialize(&1234)?])?;
1453            assert_eq!(page_bytes, Some(expected_bincode2));
1454        }
1455
1456        // Check that overriding the serializer with Json works
1457        let mut storage = MockStorage::new();
1458        let json_keymap: Keymap<i32, i32, Json> = KeymapBuilder::new(b"test2")
1459            .with_page_size(page_size)
1460            .build();
1461        json_keymap.insert(&mut storage, &1234, &1234)?;
1462
1463        let key = [json_keymap.as_slice(), INDEXES, &0_u32.to_be_bytes()].concat();
1464        if json_keymap.page_size == 1 {
1465            let item_data = storage.get(&key);
1466            let expected = b"1234".to_vec();
1467            assert_eq!(item_data, Some(expected));
1468        } else {
1469            let bytes = storage.get(&key);
1470            let expected = Bincode2::serialize(&vec![b"1234".to_vec()])?;
1471            assert_eq!(bytes, Some(expected));
1472        }
1473
1474        Ok(())
1475    }
1476
1477    #[test]
1478    fn test_keymap_paging_last_page() -> StdResult<()> {
1479        let mut storage = MockStorage::new();
1480
1481        let total_items: u32 = 20;
1482        let keymap: Keymap<Vec<u8>, u32> = Keymap::new(b"test");
1483
1484        for i in 0..total_items {
1485            let key: Vec<u8> = (i as i32).to_be_bytes().to_vec();
1486            keymap.insert(&mut storage, &key, &i)?;
1487        }
1488
1489        assert_eq!(keymap.paging(&storage, 0, 23)?.len(), 20);
1490        assert_eq!(keymap.paging_keys(&storage, 0, 23)?.len(), 20);
1491        assert_eq!(keymap.paging(&storage, 2, 8)?.len(), 4);
1492        assert_eq!(keymap.paging_keys(&storage, 2, 8)?.len(), 4);
1493        assert_eq!(keymap.paging(&storage, 2, 7)?.len(), 6);
1494        assert_eq!(keymap.paging_keys(&storage, 2, 7)?.len(), 6);
1495
1496        Ok(())
1497    }
1498
1499    #[test]
1500    fn test_add_remove_one() -> StdResult<()> {
1501        let mut storage = MockStorage::new();
1502        let keymap: Keymap<i32, i32> = Keymap::new(b"test");
1503        keymap.insert(&mut storage, &1, &1)?;
1504        assert_eq!(keymap.get_len(&storage)?, 1);
1505        keymap.remove(&mut storage, &1)?;
1506        assert_eq!(keymap.get_len(&storage)?, 0);
1507        assert!(keymap.get(&storage, &1).is_none());
1508        keymap.insert(&mut storage, &1, &1)?;
1509        assert_eq!(keymap.get_len(&storage)?, 1);
1510        Ok(())
1511    }
1512}
secret_toolkit_storage/keymap.rs

secret_toolkit_storage/
keymap.rs
