Merge pull request #215 from CosmWasm/prefixer-sub-prefix

Prefixer sub prefix
CosmWasm · Jan 19, 2021 · 7cb45e8 · 7cb45e8
2 parents 4c6e509 + 7d1ea98
commit 7cb45e8
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 14 deletions.
diff --git a/packages/storage-plus/src/indexed_map.rs b/packages/storage-plus/src/indexed_map.rs
@@ -126,6 +126,11 @@ where
     pub fn prefix(&self, p: K::Prefix) -> Prefix<T> {
         Prefix::new(self.pk_namespace, &p.prefix())
     }
+
+    // use sub_prefix to scan -> range
+    pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix<T> {
+        Prefix::new(self.pk_namespace, &p.prefix())
+    }
 }
 
 // short-cut for simple keys, rather than .prefix(()).range(...)

diff --git a/packages/storage-plus/src/indexes.rs b/packages/storage-plus/src/indexes.rs
@@ -189,6 +189,10 @@ where
         Prefix::new_de_fn(self.idx_namespace, &p.prefix(), deserialize_unique_kv)
     }
 
+    pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix<T> {
+        Prefix::new_de_fn(self.idx_namespace, &p.prefix(), deserialize_unique_kv)
+    }
+
     /// returns all items that match this secondary index, always by pk Ascending
     pub fn item(&self, store: &dyn Storage, idx: K) -> StdResult<Option<KV<T>>> {
         let data = self

diff --git a/packages/storage-plus/src/keys.rs b/packages/storage-plus/src/keys.rs
@@ -7,9 +7,10 @@ use crate::Endian;
 // pub trait PrimaryKey<'a>: Copy {
 pub trait PrimaryKey<'a>: Clone {
     type Prefix: Prefixer<'a>;
+    type SubPrefix: Prefixer<'a>;
 
     /// returns a slice of key steps, which can be optionally combined
-    fn key<'b>(&'b self) -> Vec<&'b [u8]>;
+    fn key(&self) -> Vec<&[u8]>;
 
     fn joined_key(&self) -> Vec<u8> {
         let keys = self.key();
@@ -24,8 +25,9 @@ pub trait PrimaryKey<'a>: Clone {
 
 impl<'a> PrimaryKey<'a> for &'a [u8] {
     type Prefix = ();
+    type SubPrefix = ();
 
-    fn key<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn key(&self) -> Vec<&[u8]> {
         // this is simple, we don't add more prefixes
         vec![self]
     }
@@ -38,8 +40,9 @@ impl<'a> PrimaryKey<'a> for &'a [u8] {
 // Provide a string version of this to raw encode strings
 impl<'a> PrimaryKey<'a> for &'a str {
     type Prefix = ();
+    type SubPrefix = ();
 
-    fn key<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn key(&self) -> Vec<&[u8]> {
         // this is simple, we don't add more prefixes
         vec![self.as_bytes()]
     }
@@ -52,8 +55,9 @@ impl<'a> PrimaryKey<'a> for &'a str {
 // use generics for combining there - so we can use &[u8], PkOwned, or IntKey
 impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a>> PrimaryKey<'a> for (T, U) {
     type Prefix = T;
+    type SubPrefix = ();
 
-    fn key<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn key(&self) -> Vec<&[u8]> {
         let mut keys = self.0.key();
         keys.extend(&self.1.key());
         keys
@@ -72,6 +76,7 @@ impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a> + Prefixer<'a>, V:
     PrimaryKey<'a> for (T, U, V)
 {
     type Prefix = (T, U);
+    type SubPrefix = T;
 
     fn key(&self) -> Vec<&[u8]> {
         let mut keys = self.0.key();
@@ -97,31 +102,31 @@ impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a> + Prefixer<'a>, V:
 // pub trait Prefixer<'a>: Copy {
 pub trait Prefixer<'a> {
     /// returns 0 or more namespaces that should length-prefixed and concatenated for range searches
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]>;
+    fn prefix(&self) -> Vec<&[u8]>;
 }
 
 impl<'a> Prefixer<'a> for () {
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn prefix(&self) -> Vec<&[u8]> {
         vec![]
     }
 }
 
 impl<'a> Prefixer<'a> for &'a [u8] {
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn prefix(&self) -> Vec<&[u8]> {
         vec![self]
     }
 }
 
 impl<'a, T: Prefixer<'a>, U: Prefixer<'a>> Prefixer<'a> for (T, U) {
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn prefix(&self) -> Vec<&[u8]> {
         let mut res = self.0.prefix();
         res.extend(self.1.prefix().into_iter());
         res
     }
 }
 
 impl<'a, T: Prefixer<'a>, U: Prefixer<'a>, V: Prefixer<'a>> Prefixer<'a> for (T, U, V) {
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn prefix(&self) -> Vec<&[u8]> {
         let mut res = self.0.prefix();
         res.extend(self.1.prefix().into_iter());
         res.extend(self.2.prefix().into_iter());
@@ -131,7 +136,7 @@ impl<'a, T: Prefixer<'a>, U: Prefixer<'a>, V: Prefixer<'a>> Prefixer<'a> for (T,
 
 // Provide a string version of this to raw encode strings
 impl<'a> Prefixer<'a> for &'a str {
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn prefix(&self) -> Vec<&[u8]> {
         vec![self.as_bytes()]
     }
 }
@@ -151,8 +156,9 @@ pub struct PkOwned(pub Vec<u8>);
 
 impl<'a> PrimaryKey<'a> for PkOwned {
     type Prefix = ();
+    type SubPrefix = ();
 
-    fn key<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn key(&self) -> Vec<&[u8]> {
         vec![&self.0]
     }
 
@@ -162,16 +168,17 @@ impl<'a> PrimaryKey<'a> for PkOwned {
 }
 
 impl<'a> Prefixer<'a> for PkOwned {
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn prefix(&self) -> Vec<&[u8]> {
         vec![&self.0]
     }
 }
 
 // this auto-implements PrimaryKey for all the IntKey types (and more!)
 impl<'a, T: AsRef<PkOwned> + From<PkOwned> + Clone> PrimaryKey<'a> for T {
     type Prefix = ();
+    type SubPrefix = ();
 
-    fn key<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn key(&self) -> Vec<&[u8]> {
         self.as_ref().key()
     }
 
@@ -182,7 +189,7 @@ impl<'a, T: AsRef<PkOwned> + From<PkOwned> + Clone> PrimaryKey<'a> for T {
 
 // this auto-implements Prefixer for all the IntKey types (and more!)
 impl<'a, T: AsRef<PkOwned>> Prefixer<'a> for T {
-    fn prefix<'b>(&'b self) -> Vec<&'b [u8]> {
+    fn prefix(&self) -> Vec<&[u8]> {
         self.as_ref().prefix()
     }
 }

diff --git a/packages/storage-plus/src/map.rs b/packages/storage-plus/src/map.rs
@@ -42,6 +42,11 @@ where
         Prefix::new(self.namespace, &p.prefix())
     }
 
+    #[cfg(feature = "iterator")]
+    pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix<T> {
+        Prefix::new(self.namespace, &p.prefix())
+    }
+
     pub fn save(&self, store: &mut dyn Storage, k: K, data: &T) -> StdResult<()> {
         self.key(k).save(store, data)
     }
@@ -104,6 +109,8 @@ mod test {
     use serde::{Deserialize, Serialize};
     use std::ops::Deref;
 
+    #[cfg(feature = "iterator")]
+    use crate::iter_helpers::to_length_prefixed;
     use crate::U8Key;
     use cosmwasm_std::testing::MockStorage;
     #[cfg(feature = "iterator")]
@@ -290,6 +297,9 @@ mod test {
         TRIPLE
             .save(&mut store, (b"owner", 9u8.into(), "recipient2"), &3000)
             .unwrap();
+        TRIPLE
+            .save(&mut store, (b"owner", 10u8.into(), "recipient3"), &3000)
+            .unwrap();
         TRIPLE
             .save(&mut store, (b"owner2", 9u8.into(), "recipient"), &5000)
             .unwrap();
@@ -308,6 +318,32 @@ mod test {
                 (b"recipient2".to_vec(), 3000)
             ]
         );
+
+        // let's iterate over a sub prefix
+        let all: StdResult<Vec<_>> = TRIPLE
+            .sub_prefix(b"owner")
+            .range(&store, None, None, Order::Ascending)
+            .collect();
+        let all = all.unwrap();
+        assert_eq!(3, all.len());
+        // FIXME: range() works, but remaining keys are still encoded
+        assert_eq!(
+            all,
+            vec![
+                (
+                    [to_length_prefixed(b"\x09"), b"recipient".to_vec()].concat(),
+                    1000
+                ),
+                (
+                    [to_length_prefixed(b"\x09"), b"recipient2".to_vec()].concat(),
+                    3000
+                ),
+                (
+                    [to_length_prefixed(b"\x0a"), b"recipient3".to_vec()].concat(),
+                    3000
+                )
+            ]
+        );
     }
 
     #[test]