cw_storage_plus/
indexed_snapshot.rs

1// this module requires iterator to be useful at all
2#![cfg(feature = "iterator")]
3
4use cosmwasm_std::{StdError, StdResult, Storage};
5use serde::de::DeserializeOwned;
6use serde::Serialize;
7
8use crate::de::KeyDeserialize;
9use crate::iter_helpers::deserialize_kv;
10use crate::keys::{Prefixer, PrimaryKey};
11use crate::namespace::Namespace;
12use crate::prefix::{namespaced_prefix_range, Prefix};
13use crate::snapshot::{ChangeSet, SnapshotMap};
14use crate::PrefixBound;
15use crate::{Bound, IndexList, Map, Path, Strategy};
16
17/// `IndexedSnapshotMap` works like a `SnapshotMap` but has a secondary index
18pub struct IndexedSnapshotMap<K, T, I> {
19    pk_namespace: Namespace,
20    primary: SnapshotMap<K, T>,
21    /// This is meant to be read directly to get the proper types, like:
22    /// map.idx.owner.items(...)
23    pub idx: I,
24}
25
26impl<K, T, I> IndexedSnapshotMap<K, T, I> {
27    /// Examples:
28    ///
29    /// ```rust
30    /// use cw_storage_plus::{IndexedSnapshotMap, Strategy, UniqueIndex};
31    ///
32    /// #[derive(PartialEq, Debug, Clone)]
33    /// struct Data {
34    ///     pub name: String,
35    ///     pub age: u32,
36    /// }
37    ///
38    /// let indexes = UniqueIndex::new(|d: &Data| d.age, "data__age");
39    ///
40    /// IndexedSnapshotMap::<&[u8], Data, UniqueIndex<u32, Data, ()>>::new(
41    ///     "data",
42    ///     "checkpoints",
43    ///     "changelog",
44    ///     Strategy::EveryBlock,
45    ///     indexes,
46    /// );
47    /// ```
48    pub fn new(
49        pk_namespace: impl Into<Namespace>,
50        checkpoints: impl Into<Namespace>,
51        changelog: impl Into<Namespace>,
52        strategy: Strategy,
53        indexes: I,
54    ) -> Self {
55        let pk_namespace = pk_namespace.into();
56        IndexedSnapshotMap {
57            pk_namespace: pk_namespace.clone(),
58            primary: SnapshotMap::new_dyn(pk_namespace, checkpoints, changelog, strategy),
59            idx: indexes,
60        }
61    }
62
63    pub fn changelog(&self) -> &Map<(K, u64), ChangeSet<T>> {
64        self.primary.changelog()
65    }
66}
67
68impl<'a, K, T, I> IndexedSnapshotMap<K, T, I>
69where
70    T: Serialize + DeserializeOwned + Clone,
71    K: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
72    I: IndexList<T>,
73{
74    pub fn add_checkpoint(&self, store: &mut dyn Storage, height: u64) -> StdResult<()> {
75        self.primary.add_checkpoint(store, height)
76    }
77
78    pub fn remove_checkpoint(&self, store: &mut dyn Storage, height: u64) -> StdResult<()> {
79        self.primary.remove_checkpoint(store, height)
80    }
81
82    pub fn may_load_at_height(
83        &self,
84        store: &dyn Storage,
85        k: K,
86        height: u64,
87    ) -> StdResult<Option<T>> {
88        self.primary.may_load_at_height(store, k, height)
89    }
90
91    pub fn assert_checkpointed(&self, store: &dyn Storage, height: u64) -> StdResult<()> {
92        self.primary.assert_checkpointed(store, height)
93    }
94
95    pub fn key(&self, k: K) -> Path<T> {
96        self.primary.key(k)
97    }
98}
99
100impl<'a, K, T, I> IndexedSnapshotMap<K, T, I>
101where
102    K: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
103    T: Serialize + DeserializeOwned + Clone,
104    I: IndexList<T>,
105{
106    /// save will serialize the model and store, returns an error on serialization issues.
107    /// this must load the old value to update the indexes properly
108    /// if you loaded the old value earlier in the same function, use replace to avoid needless db reads
109    pub fn save(&self, store: &mut dyn Storage, key: K, data: &T, height: u64) -> StdResult<()> {
110        let old_data = self.may_load(store, key.clone())?;
111        self.replace(store, key, Some(data), old_data.as_ref(), height)
112    }
113
114    pub fn remove(&self, store: &mut dyn Storage, key: K, height: u64) -> StdResult<()> {
115        let old_data = self.may_load(store, key.clone())?;
116        self.replace(store, key, None, old_data.as_ref(), height)
117    }
118
119    /// replace writes data to key. old_data must be the current stored value (from a previous load)
120    /// and is used to properly update the index. This is used by save, replace, and update
121    /// and can be called directly if you want to optimize
122    pub fn replace(
123        &self,
124        store: &mut dyn Storage,
125        key: K,
126        data: Option<&T>,
127        old_data: Option<&T>,
128        height: u64,
129    ) -> StdResult<()> {
130        // this is the key *relative* to the primary map namespace
131        let pk = key.joined_key();
132        if let Some(old) = old_data {
133            for index in self.idx.get_indexes() {
134                index.remove(store, &pk, old)?;
135            }
136        }
137        if let Some(updated) = data {
138            for index in self.idx.get_indexes() {
139                index.save(store, &pk, updated)?;
140            }
141            self.primary.save(store, key, updated, height)?;
142        } else {
143            self.primary.remove(store, key, height)?;
144        }
145        Ok(())
146    }
147
148    /// Loads the data, perform the specified action, and store the result
149    /// in the database. This is shorthand for some common sequences, which may be useful.
150    ///
151    /// If the data exists, `action(Some(value))` is called. Otherwise `action(None)` is called.
152    pub fn update<A, E>(
153        &self,
154        store: &mut dyn Storage,
155        key: K,
156        height: u64,
157        action: A,
158    ) -> Result<T, E>
159    where
160        A: FnOnce(Option<T>) -> Result<T, E>,
161        E: From<StdError>,
162    {
163        let input = self.may_load(store, key.clone())?;
164        let old_val = input.clone();
165        let output = action(input)?;
166        self.replace(store, key, Some(&output), old_val.as_ref(), height)?;
167        Ok(output)
168    }
169
170    // Everything else, that doesn't touch indexers, is just pass-through from self.core,
171    // thus can be used from while iterating over indexes
172
173    /// load will return an error if no data is set at the given key, or on parse error
174    pub fn load(&self, store: &dyn Storage, key: K) -> StdResult<T> {
175        self.primary.load(store, key)
176    }
177
178    /// may_load will parse the data stored at the key if present, returns Ok(None) if no data there.
179    /// returns an error on issues parsing
180    pub fn may_load(&self, store: &dyn Storage, key: K) -> StdResult<Option<T>> {
181        self.primary.may_load(store, key)
182    }
183
184    // use no_prefix to scan -> range
185    pub fn no_prefix_raw(&self) -> Prefix<Vec<u8>, T, K> {
186        Prefix::new(self.pk_namespace.as_slice(), &[])
187    }
188}
189
190// short-cut for simple keys, rather than .prefix(()).range_raw(...)
191impl<'a, K, T, I> IndexedSnapshotMap<K, T, I>
192where
193    K: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
194    T: Serialize + DeserializeOwned + Clone,
195    I: IndexList<T>,
196{
197    // I would prefer not to copy code from Prefix, but no other way
198    // with lifetimes (create Prefix inside function and return ref = no no)
199    pub fn range_raw<'c>(
200        &self,
201        store: &'c dyn Storage,
202        min: Option<Bound<'a, K>>,
203        max: Option<Bound<'a, K>>,
204        order: cosmwasm_std::Order,
205    ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::Record<T>>> + 'c>
206    where
207        T: 'c,
208    {
209        self.no_prefix_raw().range_raw(store, min, max, order)
210    }
211
212    pub fn keys_raw<'c>(
213        &self,
214        store: &'c dyn Storage,
215        min: Option<Bound<'a, K>>,
216        max: Option<Bound<'a, K>>,
217        order: cosmwasm_std::Order,
218    ) -> Box<dyn Iterator<Item = Vec<u8>> + 'c> {
219        self.no_prefix_raw().keys_raw(store, min, max, order)
220    }
221}
222
223#[cfg(feature = "iterator")]
224impl<'a, K, T, I> IndexedSnapshotMap<K, T, I>
225where
226    T: Serialize + DeserializeOwned + Clone,
227    K: PrimaryKey<'a>,
228    I: IndexList<T>,
229{
230    pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix<K::SuperSuffix, T, K::SuperSuffix> {
231        Prefix::new(self.pk_namespace.as_slice(), &p.prefix())
232    }
233
234    pub fn prefix(&self, p: K::Prefix) -> Prefix<K::Suffix, T, K::Suffix> {
235        Prefix::new(self.pk_namespace.as_slice(), &p.prefix())
236    }
237}
238
239#[cfg(feature = "iterator")]
240impl<'a, K, T, I> IndexedSnapshotMap<K, T, I>
241where
242    T: Serialize + DeserializeOwned + Clone,
243    K: PrimaryKey<'a> + KeyDeserialize,
244    I: IndexList<T>,
245{
246    /// While `range` over a `prefix` fixes the prefix to one element and iterates over the
247    /// remaining, `prefix_range` accepts bounds for the lowest and highest elements of the
248    /// `Prefix` itself, and iterates over those (inclusively or exclusively, depending on
249    /// `PrefixBound`).
250    /// There are some issues that distinguish these two, and blindly casting to `Vec<u8>` doesn't
251    /// solve them.
252    pub fn prefix_range<'c>(
253        &self,
254        store: &'c dyn Storage,
255        min: Option<PrefixBound<'a, K::Prefix>>,
256        max: Option<PrefixBound<'a, K::Prefix>>,
257        order: cosmwasm_std::Order,
258    ) -> Box<dyn Iterator<Item = StdResult<(K::Output, T)>> + 'c>
259    where
260        T: 'c,
261        'a: 'c,
262        K: 'c,
263        K::Output: 'static,
264    {
265        let mapped = namespaced_prefix_range(store, self.pk_namespace.as_slice(), min, max, order)
266            .map(deserialize_kv::<K, T>);
267        Box::new(mapped)
268    }
269
270    pub fn range<'c>(
271        &self,
272        store: &'c dyn Storage,
273        min: Option<Bound<'a, K>>,
274        max: Option<Bound<'a, K>>,
275        order: cosmwasm_std::Order,
276    ) -> Box<dyn Iterator<Item = StdResult<(K::Output, T)>> + 'c>
277    where
278        T: 'c,
279        K::Output: 'static,
280    {
281        self.no_prefix().range(store, min, max, order)
282    }
283
284    pub fn keys<'c>(
285        &self,
286        store: &'c dyn Storage,
287        min: Option<Bound<'a, K>>,
288        max: Option<Bound<'a, K>>,
289        order: cosmwasm_std::Order,
290    ) -> Box<dyn Iterator<Item = StdResult<K::Output>> + 'c>
291    where
292        T: 'c,
293        K::Output: 'static,
294    {
295        self.no_prefix().keys(store, min, max, order)
296    }
297
298    fn no_prefix(&self) -> Prefix<K, T, K> {
299        Prefix::new(self.pk_namespace.as_slice(), &[])
300    }
301}
302
303#[cfg(test)]
304mod test {
305    use super::*;
306
307    use crate::indexes::test::{index_string_tuple, index_tuple};
308    use crate::{Index, MultiIndex, UniqueIndex};
309    use cosmwasm_std::testing::MockStorage;
310    use cosmwasm_std::{MemoryStorage, Order};
311    use serde::{Deserialize, Serialize};
312
313    #[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
314    struct Data {
315        pub name: String,
316        pub last_name: String,
317        pub age: u32,
318    }
319
320    struct DataIndexes<'a> {
321        // Last type parameters are for signaling pk deserialization
322        pub name: MultiIndex<'a, Vec<u8>, Data, String>,
323        pub age: UniqueIndex<'a, u32, Data, String>,
324        pub name_lastname: UniqueIndex<'a, (Vec<u8>, Vec<u8>), Data, String>,
325    }
326
327    // Future Note: this can likely be macro-derived
328    impl<'a> IndexList<Data> for DataIndexes<'a> {
329        fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<Data>> + '_> {
330            let v: Vec<&dyn Index<Data>> = vec![&self.name, &self.age, &self.name_lastname];
331            Box::new(v.into_iter())
332        }
333    }
334
335    // For composite multi index tests
336    struct DataCompositeMultiIndex<'a> {
337        // Last type parameter is for signaling pk deserialization
338        pub name_age: MultiIndex<'a, (Vec<u8>, u32), Data, String>,
339    }
340
341    // Future Note: this can likely be macro-derived
342    impl<'a> IndexList<Data> for DataCompositeMultiIndex<'a> {
343        fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<Data>> + '_> {
344            let v: Vec<&dyn Index<Data>> = vec![&self.name_age];
345            Box::new(v.into_iter())
346        }
347    }
348
349    // Can we make it easier to define this? (less wordy generic)
350    fn build_snapshot_map<'a>() -> IndexedSnapshotMap<&'a str, Data, DataIndexes<'a>> {
351        let indexes = DataIndexes {
352            name: MultiIndex::new(|_pk, d| d.name.as_bytes().to_vec(), "data", "data__name"),
353            age: UniqueIndex::new(|d| d.age, "data__age"),
354            name_lastname: UniqueIndex::new(
355                |d| index_string_tuple(&d.name, &d.last_name),
356                "data__name_lastname",
357            ),
358        };
359        IndexedSnapshotMap::new(
360            "data",
361            "checkpoints",
362            "changelog",
363            Strategy::EveryBlock,
364            indexes,
365        )
366    }
367
368    fn save_data<'a>(
369        store: &mut MockStorage,
370        map: &IndexedSnapshotMap<&'a str, Data, DataIndexes<'a>>,
371    ) -> (Vec<&'a str>, Vec<Data>) {
372        let mut pks = vec![];
373        let mut datas = vec![];
374        let mut height = 0;
375        let data = Data {
376            name: "Maria".to_string(),
377            last_name: "Doe".to_string(),
378            age: 42,
379        };
380        let pk = "1";
381        map.save(store, pk, &data, height).unwrap();
382        height += 1;
383        pks.push(pk);
384        datas.push(data);
385
386        // same name (multi-index), different last name, different age => ok
387        let data = Data {
388            name: "Maria".to_string(),
389            last_name: "Williams".to_string(),
390            age: 23,
391        };
392        let pk = "2";
393        map.save(store, pk, &data, height).unwrap();
394        height += 1;
395        pks.push(pk);
396        datas.push(data);
397
398        // different name, different last name, different age => ok
399        let data = Data {
400            name: "John".to_string(),
401            last_name: "Wayne".to_string(),
402            age: 32,
403        };
404        let pk = "3";
405        map.save(store, pk, &data, height).unwrap();
406        height += 1;
407        pks.push(pk);
408        datas.push(data);
409
410        let data = Data {
411            name: "Maria Luisa".to_string(),
412            last_name: "Rodriguez".to_string(),
413            age: 12,
414        };
415        let pk = "4";
416        map.save(store, pk, &data, height).unwrap();
417        pks.push(pk);
418        datas.push(data);
419
420        (pks, datas)
421    }
422
423    #[test]
424    fn store_and_load_by_index() {
425        let mut store = MockStorage::new();
426        let map = build_snapshot_map();
427
428        // save data
429        let (pks, datas) = save_data(&mut store, &map);
430        let pk = pks[0];
431        let data = &datas[0];
432
433        // load it properly
434        let loaded = map.load(&store, pk).unwrap();
435        assert_eq!(*data, loaded);
436
437        let count = map
438            .idx
439            .name
440            .prefix(b"Maria".to_vec())
441            .range_raw(&store, None, None, Order::Ascending)
442            .count();
443        assert_eq!(2, count);
444
445        // load it by secondary index
446        let marias: Vec<_> = map
447            .idx
448            .name
449            .prefix(b"Maria".to_vec())
450            .range_raw(&store, None, None, Order::Ascending)
451            .collect::<StdResult<_>>()
452            .unwrap();
453        assert_eq!(2, marias.len());
454        let (k, v) = &marias[0];
455        assert_eq!(pk.as_bytes(), k);
456        assert_eq!(data, v);
457
458        // other index doesn't match (1 byte after)
459        let count = map
460            .idx
461            .name
462            .prefix(b"Marib".to_vec())
463            .range_raw(&store, None, None, Order::Ascending)
464            .count();
465        assert_eq!(0, count);
466
467        // other index doesn't match (1 byte before)
468        let count = map
469            .idx
470            .name
471            .prefix(b"Mari`".to_vec())
472            .range_raw(&store, None, None, Order::Ascending)
473            .count();
474        assert_eq!(0, count);
475
476        // other index doesn't match (longer)
477        let count = map
478            .idx
479            .name
480            .prefix(b"Maria5".to_vec())
481            .range_raw(&store, None, None, Order::Ascending)
482            .count();
483        assert_eq!(0, count);
484
485        // match on proper age
486        let proper = 42u32;
487        let aged = map.idx.age.item(&store, proper).unwrap().unwrap();
488        assert_eq!(pk.as_bytes(), aged.0);
489        assert_eq!(*data, aged.1);
490
491        // no match on wrong age
492        let too_old = 43u32;
493        let aged = map.idx.age.item(&store, too_old).unwrap();
494        assert_eq!(None, aged);
495    }
496
497    #[test]
498    fn range_raw_simple_key_by_multi_index() {
499        let mut store = MockStorage::new();
500        let map = build_snapshot_map();
501        let mut height = 1;
502
503        // save data
504        let data1 = Data {
505            name: "Maria".to_string(),
506            last_name: "".to_string(),
507            age: 42,
508        };
509        let pk = "5627";
510        map.save(&mut store, pk, &data1, height).unwrap();
511        height += 1;
512
513        let data2 = Data {
514            name: "Juan".to_string(),
515            last_name: "Perez".to_string(),
516            age: 13,
517        };
518        let pk = "5628";
519        map.save(&mut store, pk, &data2, height).unwrap();
520        height += 1;
521
522        let data3 = Data {
523            name: "Maria".to_string(),
524            last_name: "Williams".to_string(),
525            age: 24,
526        };
527        let pk = "5629";
528        map.save(&mut store, pk, &data3, height).unwrap();
529        height += 1;
530
531        let data4 = Data {
532            name: "Maria Luisa".to_string(),
533            last_name: "Bemberg".to_string(),
534            age: 12,
535        };
536        let pk = "5630";
537        map.save(&mut store, pk, &data4, height).unwrap();
538
539        let marias: Vec<_> = map
540            .idx
541            .name
542            .prefix(b"Maria".to_vec())
543            .range_raw(&store, None, None, Order::Descending)
544            .collect::<StdResult<_>>()
545            .unwrap();
546        let count = marias.len();
547        assert_eq!(2, count);
548
549        // Sorted by (descending) pk
550        assert_eq!(marias[0].0, b"5629");
551        assert_eq!(marias[1].0, b"5627");
552        // Data is correct
553        assert_eq!(marias[0].1, data3);
554        assert_eq!(marias[1].1, data1);
555    }
556
557    #[test]
558    fn range_simple_key_by_multi_index() {
559        let mut store = MockStorage::new();
560        let map = build_snapshot_map();
561        let mut height = 1;
562
563        // save data
564        let data1 = Data {
565            name: "Maria".to_string(),
566            last_name: "".to_string(),
567            age: 42,
568        };
569        let pk = "5627";
570        map.save(&mut store, pk, &data1, height).unwrap();
571        height += 1;
572
573        let data2 = Data {
574            name: "Juan".to_string(),
575            last_name: "Perez".to_string(),
576            age: 13,
577        };
578        let pk = "5628";
579        map.save(&mut store, pk, &data2, height).unwrap();
580        height += 1;
581
582        let data3 = Data {
583            name: "Maria".to_string(),
584            last_name: "Williams".to_string(),
585            age: 24,
586        };
587        let pk = "5629";
588        map.save(&mut store, pk, &data3, height).unwrap();
589        height += 1;
590
591        let data4 = Data {
592            name: "Maria Luisa".to_string(),
593            last_name: "Bemberg".to_string(),
594            age: 12,
595        };
596        let pk = "5630";
597        map.save(&mut store, pk, &data4, height).unwrap();
598
599        let marias: Vec<_> = map
600            .idx
601            .name
602            .prefix(b"Maria".to_vec())
603            .range(&store, None, None, Order::Descending)
604            .collect::<StdResult<_>>()
605            .unwrap();
606        let count = marias.len();
607        assert_eq!(2, count);
608
609        // Sorted by (descending) pk
610        assert_eq!(marias[0].0, "5629");
611        assert_eq!(marias[1].0, "5627");
612        // Data is correct
613        assert_eq!(marias[0].1, data3);
614        assert_eq!(marias[1].1, data1);
615    }
616
617    #[test]
618    fn changelog_range_works() {
619        let mut store = MockStorage::new();
620        let map = build_snapshot_map();
621        let mut height = 1;
622
623        // simple data for testing
624        // EVERY.remove(&mut store, "B", 4).unwrap();
625        let data1 = Data {
626            name: "Maria".to_string(),
627            last_name: "".to_string(),
628            age: 42,
629        };
630        let pk_a = "A";
631        map.save(&mut store, pk_a, &data1, height).unwrap();
632        height += 1;
633
634        let data2 = Data {
635            name: "Juan".to_string(),
636            last_name: "Perez".to_string(),
637            age: 13,
638        };
639        let pk_b = "B";
640        map.save(&mut store, pk_b, &data2, height).unwrap();
641        height += 1;
642
643        let data3 = Data {
644            name: "Maria".to_string(),
645            last_name: "Williams".to_string(),
646            age: 24,
647        };
648        map.update(&mut store, pk_a, height, |_| -> StdResult<Data> {
649            Ok(data3)
650        })
651        .unwrap();
652
653        height += 1;
654        map.remove(&mut store, pk_b, height).unwrap();
655
656        let changes: Vec<_> = map
657            .changelog()
658            .range(&store, None, None, Order::Ascending)
659            .collect::<StdResult<_>>()
660            .unwrap();
661        let count = changes.len();
662        assert_eq!(4, count);
663
664        // sorted by ascending key, height
665        assert_eq!(
666            changes,
667            vec![
668                (("A".into(), 1), ChangeSet { old: None }),
669                (("A".into(), 3), ChangeSet { old: Some(data1) }),
670                (("B".into(), 2), ChangeSet { old: None }),
671                (("B".into(), 4), ChangeSet { old: Some(data2) })
672            ]
673        );
674    }
675
676    #[test]
677    fn range_raw_composite_key_by_multi_index() {
678        let mut store = MockStorage::new();
679        let mut height = 2;
680
681        let indexes = DataCompositeMultiIndex {
682            name_age: MultiIndex::new(
683                |_pk, d| index_tuple(&d.name, d.age),
684                "data",
685                "data__name_age",
686            ),
687        };
688        let map =
689            IndexedSnapshotMap::new("data", "checks", "changes", Strategy::EveryBlock, indexes);
690
691        // save data
692        let data1 = Data {
693            name: "Maria".to_string(),
694            last_name: "".to_string(),
695            age: 42,
696        };
697        let pk1 = "5627";
698        map.save(&mut store, pk1, &data1, height).unwrap();
699        height += 1;
700
701        let data2 = Data {
702            name: "Juan".to_string(),
703            last_name: "Perez".to_string(),
704            age: 13,
705        };
706        let pk2 = "5628";
707        map.save(&mut store, pk2, &data2, height).unwrap();
708        height += 1;
709
710        let data3 = Data {
711            name: "Maria".to_string(),
712            last_name: "Young".to_string(),
713            age: 24,
714        };
715        let pk3 = "5629";
716        map.save(&mut store, pk3, &data3, height).unwrap();
717        height += 1;
718
719        let data4 = Data {
720            name: "Maria Luisa".to_string(),
721            last_name: "Bemberg".to_string(),
722            age: 43,
723        };
724        let pk4 = "5630";
725        map.save(&mut store, pk4, &data4, height).unwrap();
726
727        let marias: Vec<_> = map
728            .idx
729            .name_age
730            .sub_prefix(b"Maria".to_vec())
731            .range_raw(&store, None, None, Order::Descending)
732            .collect::<StdResult<_>>()
733            .unwrap();
734        let count = marias.len();
735        assert_eq!(2, count);
736
737        // Pks (sorted by age descending)
738        assert_eq!(pk1.as_bytes(), marias[0].0);
739        assert_eq!(pk3.as_bytes(), marias[1].0);
740
741        // Data
742        assert_eq!(data1, marias[0].1);
743        assert_eq!(data3, marias[1].1);
744    }
745
746    #[test]
747    fn range_composite_key_by_multi_index() {
748        let mut store = MockStorage::new();
749        let mut height = 2;
750
751        let indexes = DataCompositeMultiIndex {
752            name_age: MultiIndex::new(
753                |_pk, d| index_tuple(&d.name, d.age),
754                "data",
755                "data__name_age",
756            ),
757        };
758        let map =
759            IndexedSnapshotMap::new("data", "checks", "changes", Strategy::EveryBlock, indexes);
760
761        // save data
762        let data1 = Data {
763            name: "Maria".to_string(),
764            last_name: "".to_string(),
765            age: 42,
766        };
767        let pk1 = "5627";
768        map.save(&mut store, pk1, &data1, height).unwrap();
769        height += 1;
770
771        let data2 = Data {
772            name: "Juan".to_string(),
773            last_name: "Perez".to_string(),
774            age: 13,
775        };
776        let pk2 = "5628";
777        map.save(&mut store, pk2, &data2, height).unwrap();
778        height += 1;
779
780        let data3 = Data {
781            name: "Maria".to_string(),
782            last_name: "Young".to_string(),
783            age: 24,
784        };
785        let pk3 = "5629";
786        map.save(&mut store, pk3, &data3, height).unwrap();
787        height += 1;
788
789        let data4 = Data {
790            name: "Maria Luisa".to_string(),
791            last_name: "Bemberg".to_string(),
792            age: 43,
793        };
794        let pk4 = "5630";
795        map.save(&mut store, pk4, &data4, height).unwrap();
796
797        let marias: Vec<_> = map
798            .idx
799            .name_age
800            .sub_prefix(b"Maria".to_vec())
801            .range(&store, None, None, Order::Descending)
802            .collect::<StdResult<_>>()
803            .unwrap();
804        let count = marias.len();
805        assert_eq!(2, count);
806
807        // Pks (sorted by age descending)
808        assert_eq!(pk1.to_string(), marias[0].0);
809        assert_eq!(pk3.to_string(), marias[1].0);
810
811        // Data
812        assert_eq!(data1, marias[0].1);
813        assert_eq!(data3, marias[1].1);
814    }
815
816    #[test]
817    fn unique_index_enforced() {
818        let mut store = MockStorage::new();
819        let map = build_snapshot_map();
820        let mut height = 3;
821
822        // save data
823        let (pks, datas) = save_data(&mut store, &map);
824
825        // different name, different last name, same age => error
826        let data5 = Data {
827            name: "Marta".to_string(),
828            last_name: "Laurens".to_string(),
829            age: 42,
830        };
831        let pk5 = "4";
832
833        // enforce this returns some error
834        map.save(&mut store, pk5, &data5, height).unwrap_err();
835        height += 1;
836
837        // query by unique key
838        // match on proper age
839        let age42 = 42u32;
840        let (k, v) = map.idx.age.item(&store, age42).unwrap().unwrap();
841        assert_eq!(k, pks[0].as_bytes());
842        assert_eq!(v.name, datas[0].name);
843        assert_eq!(v.age, datas[0].age);
844
845        // match on other age
846        let age23 = 23u32;
847        let (k, v) = map.idx.age.item(&store, age23).unwrap().unwrap();
848        assert_eq!(k, pks[1].as_bytes());
849        assert_eq!(v.name, datas[1].name);
850        assert_eq!(v.age, datas[1].age);
851
852        // if we delete the first one, we can add the blocked one
853        map.remove(&mut store, pks[0], height).unwrap();
854        height += 1;
855        map.save(&mut store, pk5, &data5, height).unwrap();
856        // now 42 is the new owner
857        let (k, v) = map.idx.age.item(&store, age42).unwrap().unwrap();
858        assert_eq!(k, pk5.as_bytes());
859        assert_eq!(v.name, data5.name);
860        assert_eq!(v.age, data5.age);
861    }
862
863    #[test]
864    fn unique_index_enforced_composite_key() {
865        let mut store = MockStorage::new();
866        let map = build_snapshot_map();
867        let height = 4;
868
869        // save data
870        save_data(&mut store, &map);
871
872        // same name, same lastname => error
873        let data5 = Data {
874            name: "Maria".to_string(),
875            last_name: "Doe".to_string(),
876            age: 24,
877        };
878        let pk5 = "5";
879        // enforce this returns some error
880        map.save(&mut store, pk5, &data5, height).unwrap_err();
881    }
882
883    #[test]
884    fn remove_and_update_reflected_on_indexes() {
885        let mut store = MockStorage::new();
886        let map = build_snapshot_map();
887        let mut height = 5;
888
889        let name_count = |map: &IndexedSnapshotMap<&str, Data, DataIndexes>,
890                          store: &MemoryStorage,
891                          name: &str|
892         -> usize {
893            map.idx
894                .name
895                .prefix(name.as_bytes().to_vec())
896                .keys_raw(store, None, None, Order::Ascending)
897                .count()
898        };
899
900        // save data
901        let (pks, _) = save_data(&mut store, &map);
902
903        // find 2 Marias, 1 John, and no Mary
904        assert_eq!(name_count(&map, &store, "Maria"), 2);
905        assert_eq!(name_count(&map, &store, "John"), 1);
906        assert_eq!(name_count(&map, &store, "Maria Luisa"), 1);
907        assert_eq!(name_count(&map, &store, "Mary"), 0);
908
909        // remove maria 2
910        map.remove(&mut store, pks[1], height).unwrap();
911        height += 1;
912
913        // change john to mary
914        map.update(&mut store, pks[2], height, |d| -> StdResult<_> {
915            let mut x = d.unwrap();
916            assert_eq!(&x.name, "John");
917            x.name = "Mary".to_string();
918            Ok(x)
919        })
920        .unwrap();
921
922        // find 1 maria, 1 maria luisa, no john, and 1 mary
923        assert_eq!(name_count(&map, &store, "Maria"), 1);
924        assert_eq!(name_count(&map, &store, "Maria Luisa"), 1);
925        assert_eq!(name_count(&map, &store, "John"), 0);
926        assert_eq!(name_count(&map, &store, "Mary"), 1);
927    }
928
929    #[test]
930    fn range_raw_simple_key_by_unique_index() {
931        let mut store = MockStorage::new();
932        let map = build_snapshot_map();
933
934        // save data
935        let (pks, datas) = save_data(&mut store, &map);
936
937        let res: StdResult<Vec<_>> = map
938            .idx
939            .age
940            .range_raw(&store, None, None, Order::Ascending)
941            .collect();
942        let ages = res.unwrap();
943
944        let count = ages.len();
945        assert_eq!(4, count);
946
947        // The pks, sorted by age ascending
948        assert_eq!(pks[3].as_bytes(), ages[0].0);
949        assert_eq!(pks[1].as_bytes(), ages[1].0);
950        assert_eq!(pks[2].as_bytes(), ages[2].0);
951        assert_eq!(pks[0].as_bytes(), ages[3].0);
952
953        // The associated data
954        assert_eq!(datas[3], ages[0].1);
955        assert_eq!(datas[1], ages[1].1);
956        assert_eq!(datas[2], ages[2].1);
957        assert_eq!(datas[0], ages[3].1);
958    }
959
960    #[test]
961    fn range_simple_key_by_unique_index() {
962        let mut store = MockStorage::new();
963        let map = build_snapshot_map();
964
965        // save data
966        let (pks, datas) = save_data(&mut store, &map);
967
968        let res: StdResult<Vec<_>> = map
969            .idx
970            .age
971            .range(&store, None, None, Order::Ascending)
972            .collect();
973        let ages = res.unwrap();
974
975        let count = ages.len();
976        assert_eq!(4, count);
977
978        // The pks, sorted by age ascending
979        assert_eq!(pks[3], ages[0].0);
980        assert_eq!(pks[1], ages[1].0);
981        assert_eq!(pks[2], ages[2].0);
982        assert_eq!(pks[0], ages[3].0);
983
984        // The associated data
985        assert_eq!(datas[3], ages[0].1);
986        assert_eq!(datas[1], ages[1].1);
987        assert_eq!(datas[2], ages[2].1);
988        assert_eq!(datas[0], ages[3].1);
989    }
990
991    #[test]
992    fn range_raw_composite_key_by_unique_index() {
993        let mut store = MockStorage::new();
994        let map = build_snapshot_map();
995
996        // save data
997        let (pks, datas) = save_data(&mut store, &map);
998
999        let res: StdResult<Vec<_>> = map
1000            .idx
1001            .name_lastname
1002            .prefix(b"Maria".to_vec())
1003            .range_raw(&store, None, None, Order::Ascending)
1004            .collect();
1005        let marias = res.unwrap();
1006
1007        // Only two people are called "Maria"
1008        let count = marias.len();
1009        assert_eq!(2, count);
1010
1011        // The pks
1012        assert_eq!(pks[0].as_bytes(), marias[0].0);
1013        assert_eq!(pks[1].as_bytes(), marias[1].0);
1014
1015        // The associated data
1016        assert_eq!(datas[0], marias[0].1);
1017        assert_eq!(datas[1], marias[1].1);
1018    }
1019
1020    #[test]
1021    fn range_composite_key_by_unique_index() {
1022        let mut store = MockStorage::new();
1023        let map = build_snapshot_map();
1024
1025        // save data
1026        let (pks, datas) = save_data(&mut store, &map);
1027
1028        let res: StdResult<Vec<_>> = map
1029            .idx
1030            .name_lastname
1031            .prefix(b"Maria".to_vec())
1032            .range(&store, None, None, Order::Ascending)
1033            .collect();
1034        let marias = res.unwrap();
1035
1036        // Only two people are called "Maria"
1037        let count = marias.len();
1038        assert_eq!(2, count);
1039
1040        // The pks
1041        assert_eq!(pks[0], marias[0].0);
1042        assert_eq!(pks[1], marias[1].0);
1043
1044        // The associated data
1045        assert_eq!(datas[0], marias[0].1);
1046        assert_eq!(datas[1], marias[1].1);
1047    }
1048
1049    #[test]
1050    #[cfg(feature = "iterator")]
1051    fn range_simple_string_key() {
1052        let mut store = MockStorage::new();
1053        let map = build_snapshot_map();
1054
1055        // save data
1056        let (pks, datas) = save_data(&mut store, &map);
1057
1058        // let's try to iterate!
1059        let all: StdResult<Vec<_>> = map.range(&store, None, None, Order::Ascending).collect();
1060        let all = all.unwrap();
1061        assert_eq!(
1062            all,
1063            pks.clone()
1064                .into_iter()
1065                .map(str::to_string)
1066                .zip(datas.clone().into_iter())
1067                .collect::<Vec<_>>()
1068        );
1069
1070        // let's try to iterate over a range
1071        let all: StdResult<Vec<_>> = map
1072            .range(&store, Some(Bound::inclusive("3")), None, Order::Ascending)
1073            .collect();
1074        let all = all.unwrap();
1075        assert_eq!(
1076            all,
1077            pks.into_iter()
1078                .map(str::to_string)
1079                .zip(datas.into_iter())
1080                .rev()
1081                .take(2)
1082                .rev()
1083                .collect::<Vec<_>>()
1084        );
1085    }
1086
1087    #[test]
1088    #[cfg(feature = "iterator")]
1089    fn prefix_simple_string_key() {
1090        let mut store = MockStorage::new();
1091        let map = build_snapshot_map();
1092
1093        // save data
1094        let (pks, datas) = save_data(&mut store, &map);
1095
1096        // Let's prefix and iterate.
1097        // This is similar to calling range() directly, but added here for completeness / prefix
1098        // type checks
1099        let all: StdResult<Vec<_>> = map
1100            .prefix(())
1101            .range(&store, None, None, Order::Ascending)
1102            .collect();
1103        let all = all.unwrap();
1104        assert_eq!(
1105            all,
1106            pks.clone()
1107                .into_iter()
1108                .map(str::to_string)
1109                .zip(datas.into_iter())
1110                .collect::<Vec<_>>()
1111        );
1112    }
1113
1114    #[test]
1115    #[cfg(feature = "iterator")]
1116    fn sub_prefix_simple_string_key() {
1117        let mut store = MockStorage::new();
1118        let map = build_snapshot_map();
1119
1120        // save data
1121        let (pks, datas) = save_data(&mut store, &map);
1122
1123        // Let's sub-prefix and iterate.
1124        // This is similar to calling range() directly, but added here for completeness / sub_prefix
1125        // type checks
1126        let all: StdResult<Vec<_>> = map
1127            .sub_prefix(())
1128            .range(&store, None, None, Order::Ascending)
1129            .collect();
1130        let all = all.unwrap();
1131        assert_eq!(
1132            all,
1133            pks.clone()
1134                .into_iter()
1135                .map(str::to_string)
1136                .zip(datas.into_iter())
1137                .collect::<Vec<_>>()
1138        );
1139    }
1140
1141    #[test]
1142    #[cfg(feature = "iterator")]
1143    fn prefix_range_simple_key() {
1144        let mut store = MockStorage::new();
1145
1146        let indexes = DataCompositeMultiIndex {
1147            name_age: MultiIndex::new(
1148                |_pk, d| index_tuple(&d.name, d.age),
1149                "data",
1150                "data__name_age",
1151            ),
1152        };
1153        let map =
1154            IndexedSnapshotMap::new("data", "checks", "changes", Strategy::EveryBlock, indexes);
1155
1156        // save data
1157        let data1 = Data {
1158            name: "Maria".to_string(),
1159            last_name: "".to_string(),
1160            age: 42,
1161        };
1162        let pk1: (&str, &str) = ("1", "5627");
1163        map.save(&mut store, pk1, &data1, 1).unwrap();
1164
1165        let data2 = Data {
1166            name: "Juan".to_string(),
1167            last_name: "Perez".to_string(),
1168            age: 13,
1169        };
1170        let pk2: (&str, &str) = ("2", "5628");
1171        map.save(&mut store, pk2, &data2, 1).unwrap();
1172
1173        let data3 = Data {
1174            name: "Maria".to_string(),
1175            last_name: "Young".to_string(),
1176            age: 24,
1177        };
1178        let pk3: (&str, &str) = ("2", "5629");
1179        map.save(&mut store, pk3, &data3, 1).unwrap();
1180
1181        let data4 = Data {
1182            name: "Maria Luisa".to_string(),
1183            last_name: "Bemberg".to_string(),
1184            age: 43,
1185        };
1186        let pk4: (&str, &str) = ("3", "5630");
1187        map.save(&mut store, pk4, &data4, 1).unwrap();
1188
1189        // let's prefix-range and iterate
1190        let result: StdResult<Vec<_>> = map
1191            .prefix_range(
1192                &store,
1193                Some(PrefixBound::inclusive("2")),
1194                None,
1195                Order::Ascending,
1196            )
1197            .collect();
1198        let result = result.unwrap();
1199        assert_eq!(
1200            result,
1201            [
1202                (("2".to_string(), "5628".to_string()), data2.clone()),
1203                (("2".to_string(), "5629".to_string()), data3.clone()),
1204                (("3".to_string(), "5630".to_string()), data4)
1205            ]
1206        );
1207
1208        // let's try to iterate over a more restrictive prefix-range!
1209        let result: StdResult<Vec<_>> = map
1210            .prefix_range(
1211                &store,
1212                Some(PrefixBound::inclusive("2")),
1213                Some(PrefixBound::exclusive("3")),
1214                Order::Ascending,
1215            )
1216            .collect();
1217        let result = result.unwrap();
1218        assert_eq!(
1219            result,
1220            [
1221                (("2".to_string(), "5628".to_string()), data2),
1222                (("2".to_string(), "5629".to_string()), data3),
1223            ]
1224        );
1225    }
1226}