Primary key / index key helpers

packages/storage-plus/src/indexed_map.rs

@@ -9,6 +9,7 @@ use crate::indexes::Index;
 use crate::keys::{EmptyPrefix, Prefixer, PrimaryKey};
 use crate::map::Map;
 use crate::prefix::{Bound, Prefix};
+use crate::Path;
 
 pub trait IndexList<T> {
     fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<T>> + '_>;
@@ -43,6 +44,10 @@ where
             idx: indexes,
         }
     }
+
+    pub fn key(&self, k: K) -> Path<T> {
+        self.primary.key(k)
+    }
 }
 
 impl<'a, K, T, I> IndexedMap<'a, K, T, I>
@@ -337,6 +342,54 @@ mod test {
             .count();
         assert_eq!(0, count);
 
+        // index_key() over MultiIndex works (empty pk)
+        // In a MultiIndex, an index key is composed by the index and the primary key.
+        // Primary key may be empty (so that to iterate over all elements that match just the index)
+        let key = (PkOwned(b"Maria".to_vec()), PkOwned(b"".to_vec()));
+        // Use the index_key() helper to build the (raw) index key
+        let key = map.idx.name.index_key(key);
+        // Iterate using a bound over the raw key
+        let count = map
+            .idx
+            .name
+            .range(&store, Some(Bound::inclusive(key)), None, Order::Ascending)
+            .count();
+        // gets all the "Maria"s
+        assert_eq!(3, count);
+
+        // index_key() over MultiIndex works (non-empty pk)
+        // Build key including a non-empty pk
+        let key = (PkOwned(b"Maria".to_vec()), PkOwned(b"1".to_vec()));
+        // Use the index_key() helper to build the (raw) index key
+        let key = map.idx.name.index_key(key);
+        // Iterate using a (exclusive) bound over the raw key.
+        // (Useful for pagination / continuation contexts).
+        let count = map
+            .idx
+            .name
+            .range(&store, Some(Bound::exclusive(key)), None, Order::Ascending)
+            .count();
+        // gets all the "Maria"s except for the first one
+        assert_eq!(2, count);
+
+        // index_key() over UniqueIndex works.
+        let age_key = U32Key::from(23);
+        // Use the index_key() helper to build the (raw) index key
+        let age_key = map.idx.age.index_key(age_key);
+        // Iterate using a (inclusive) bound over the raw key.
+        let count = map
+            .idx
+            .age
+            .range(
+                &store,
+                Some(Bound::inclusive(age_key)),
+                None,
+                Order::Ascending,
+            )
+            .count();
+        // gets all the greater than or equal to 23 years old people
+        assert_eq!(3, count);
+
         // match on proper age
         let proper = U32Key::new(42);
         let aged = map.idx.age.item(&store, proper).unwrap().unwrap();

packages/storage-plus/src/indexed_snapshot.rs

@@ -8,7 +8,7 @@ use serde::Serialize;
 use crate::keys::{EmptyPrefix, Prefixer, PrimaryKey};
 use crate::prefix::{Bound, Prefix};
 use crate::snapshot::SnapshotMap;
-use crate::{IndexList, Strategy};
+use crate::{IndexList, Path, Strategy};
 
 /// IndexedSnapshotMap works like a SnapshotMap but has a secondary index
 pub struct IndexedSnapshotMap<'a, K, T, I>
@@ -71,6 +71,10 @@ where
     pub fn assert_checkpointed(&self, store: &dyn Storage, height: u64) -> StdResult<()> {
         self.primary.assert_checkpointed(store, height)
     }
+
+    pub fn key(&self, k: K) -> Path<T> {
+        self.primary.key(k)
+    }
 }
 
 impl<'a, K, T, I> IndexedSnapshotMap<'a, K, T, I>

packages/storage-plus/src/indexes.rs

@@ -138,6 +138,10 @@ where
         )
     }
 
+    pub fn index_key(&self, k: K) -> Vec<u8> {
+        k.joined_key()
+    }
+
     // FIXME?: Move to Prefix<T> for ergonomics
     pub fn pks<'c>(
         &self,
@@ -267,6 +271,10 @@ where
     T: Serialize + DeserializeOwned + Clone,
     K: PrimaryKey<'a>,
 {
+    pub fn index_key(&self, k: K) -> Vec<u8> {
+        k.joined_key()
+    }
+
     pub fn prefix(&self, p: K::Prefix) -> Prefix<T> {
         Prefix::with_deserialization_function(self.idx_namespace, &p.prefix(), &[], |_, _, kv| {
             deserialize_unique_kv(kv)