Add range bounds to MultiIndex

CosmWasm · Oct 16, 2020 · f99c26c · f99c26c
1 parent 508fbdc
commit f99c26c
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 18 deletions.
diff --git a/packages/storage-plus/src/indexed_map.rs b/packages/storage-plus/src/indexed_map.rs
@@ -140,7 +140,7 @@ mod test {
 
     use crate::indexes::{index_int, index_string, MultiIndex, UniqueIndex};
     use cosmwasm_std::testing::MockStorage;
-    use cosmwasm_std::MemoryStorage;
+    use cosmwasm_std::{MemoryStorage, Order};
     use serde::{Deserialize, Serialize};
 
     #[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
@@ -188,32 +188,54 @@ mod test {
         let loaded = map.load(&store, pk).unwrap();
         assert_eq!(data, loaded);
 
-        let count = map.idx.name.items(&store, &index_string("Maria")).count();
+        let count = map
+            .idx
+            .name
+            .all_items(&store, &index_string("Maria"))
+            .unwrap()
+            .len();
         assert_eq!(1, count);
 
         // TODO: we load by wrong keys - get full storage key!
 
         // load it by secondary index (we must know how to compute this)
         // let marias: StdResult<Vec<_>> = map
-        let marias: Vec<StdResult<_>> =
-            map.idx.name.items(&store, &index_string("Maria")).collect();
+        let marias = map
+            .idx
+            .name
+            .all_items(&store, &index_string("Maria"))
+            .unwrap();
         assert_eq!(1, marias.len());
-        assert!(marias[0].is_ok());
-        let (k, v) = marias[0].as_ref().unwrap();
+        let (k, v) = &marias[0];
         assert_eq!(pk, k.as_slice());
         assert_eq!(&data, v);
 
         // other index doesn't match (1 byte after)
-        let marias = map.idx.name.items(&store, &index_string("Marib")).count();
-        assert_eq!(0, marias);
+        let count = map
+            .idx
+            .name
+            .all_items(&store, &index_string("Marib"))
+            .unwrap()
+            .len();
+        assert_eq!(0, count);
 
         // other index doesn't match (1 byte before)
-        let marias = map.idx.name.items(&store, &index_string("Mari`")).count();
-        assert_eq!(0, marias);
+        let count = map
+            .idx
+            .name
+            .all_items(&store, &index_string("Mari`"))
+            .unwrap()
+            .len();
+        assert_eq!(0, count);
 
         // other index doesn't match (longer)
-        let marias = map.idx.name.items(&store, &index_string("Maria5")).count();
-        assert_eq!(0, marias);
+        let count = map
+            .idx
+            .name
+            .all_items(&store, &index_string("Maria5"))
+            .unwrap()
+            .len();
+        assert_eq!(0, count);
 
         // match on proper age
         let proper = index_int(42);
@@ -291,7 +313,12 @@ mod test {
             |map: &IndexedMap<&[u8], Data, MemoryStorage, DataIndexes<MemoryStorage>>,
              store: &MemoryStorage,
              name: &str|
-             -> usize { map.idx.name.pks(store, &index_string(name)).count() };
+             -> usize {
+                map.idx
+                    .name
+                    .pks(store, &index_string(name), None, None, Order::Ascending)
+                    .count()
+            };
 
         // set up some data
         let data1 = Data {

diff --git a/packages/storage-plus/src/indexes.rs b/packages/storage-plus/src/indexes.rs
@@ -9,7 +9,7 @@ use cosmwasm_std::{Binary, Order, StdError, StdResult, Storage, KV};
 
 use crate::map::Map;
 use crate::prefix::range_with_prefix;
-use crate::Endian;
+use crate::{Bound, Endian};
 
 /// MARKER is stored in the multi-index as value, but we only look at the key (which is pk)
 const MARKER: u32 = 1;
@@ -84,10 +84,16 @@ where
     S: Storage,
     T: Serialize + DeserializeOwned + Clone,
 {
-    pub fn pks<'c>(&self, store: &'c S, idx: &[u8]) -> Box<dyn Iterator<Item = Vec<u8>> + 'c> {
+    pub fn pks<'c>(
+        &self,
+        store: &'c S,
+        idx: &[u8],
+        min: Option<Bound>,
+        max: Option<Bound>,
+        order: Order,
+    ) -> Box<dyn Iterator<Item = Vec<u8>> + 'c> {
         let prefix = self.idx_map.prefix(idx);
-        let mapped =
-            range_with_prefix(store, &prefix, None, None, Order::Ascending).map(|(k, _)| k);
+        let mapped = range_with_prefix(store, &prefix, min, max, order).map(|(k, _)| k);
         Box::new(mapped)
     }
 
@@ -96,13 +102,32 @@ where
         &'c self,
         store: &'c S,
         idx: &[u8],
+        min: Option<Bound>,
+        max: Option<Bound>,
+        order: Order,
     ) -> Box<dyn Iterator<Item = StdResult<KV<T>>> + 'c> {
-        let mapped = self.pks(store, idx).map(move |pk| {
+        let mapped = self.pks(store, idx, min, max, order).map(move |pk| {
             let v = self.pk_map.load(store, &pk)?;
             Ok((pk, v))
         });
         Box::new(mapped)
     }
+
+    #[cfg(test)]
+    pub fn count<'c>(&self, store: &'c S, idx: &[u8]) -> usize {
+        self.pks(store, idx, None, None, Order::Ascending).count()
+    }
+
+    #[cfg(test)]
+    pub fn all_pks<'c>(&self, store: &'c S, idx: &[u8]) -> Vec<Vec<u8>> {
+        self.pks(store, idx, None, None, Order::Ascending).collect()
+    }
+
+    #[cfg(test)]
+    pub fn all_items<'c>(&self, store: &'c S, idx: &[u8]) -> StdResult<Vec<KV<T>>> {
+        self.items(store, idx, None, None, Order::Ascending)
+            .collect()
+    }
 }
 
 #[derive(Deserialize, Serialize)]