Skip to main content

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