UniqueIndex range() returns KV<T> iterator

CosmWasm · Dec 23, 2020 · c58ef15 · c58ef15
1 parent 6de36d0
commit c58ef15
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 36 deletions.
diff --git a/packages/storage-plus/src/indexed_map.rs b/packages/storage-plus/src/indexed_map.rs
@@ -484,23 +484,17 @@ mod test {
         let count = ages.len();
         assert_eq!(4, count);
 
-        // The (index) keys are the (unique, encoded) ages, in ascending order
-        assert_eq!(12u32.to_be_bytes(), ages[0].0.as_slice());
-        assert_eq!(13u32.to_be_bytes(), ages[1].0.as_slice());
-        assert_eq!(24u32.to_be_bytes(), ages[2].0.as_slice());
-        assert_eq!(42u32.to_be_bytes(), ages[3].0.as_slice());
-
-        // The pks are in the (UniqueRef) values
-        assert_eq!(b"5630".to_vec(), ages[0].1.pk);
-        assert_eq!(b"5628".to_vec(), ages[1].1.pk);
-        assert_eq!(b"5629".to_vec(), ages[2].1.pk);
-        assert_eq!(b"5627".to_vec(), ages[3].1.pk);
-
-        // The associated data is in the (UniqueRef) values
-        assert_eq!(data4, ages[0].1.value);
-        assert_eq!(data2, ages[1].1.value);
-        assert_eq!(data3, ages[2].1.value);
-        assert_eq!(data1, ages[3].1.value);
+        // The pks
+        assert_eq!(b"5630".to_vec(), ages[0].0);
+        assert_eq!(b"5628".to_vec(), ages[1].0);
+        assert_eq!(b"5629".to_vec(), ages[2].0);
+        assert_eq!(b"5627".to_vec(), ages[3].0);
+
+        // The associated data
+        assert_eq!(data4, ages[0].1);
+        assert_eq!(data2, ages[1].1);
+        assert_eq!(data3, ages[2].1);
+        assert_eq!(data1, ages[3].1);
     }
 
     #[test]
@@ -553,16 +547,12 @@ mod test {
         let count = marias.len();
         assert_eq!(2, count);
 
-        // The (index) keys are the (encoded) last names, in ascending order
-        assert_eq!(b"", marias[0].0.as_slice());
-        assert_eq!(b"Young", marias[1].0.as_slice());
+        // The pks
+        assert_eq!(b"5627".to_vec(), marias[0].0);
+        assert_eq!(b"5629".to_vec(), marias[1].0);
 
-        // The pks are in the (UniqueRef) values
-        assert_eq!(b"5627".to_vec(), marias[0].1.pk);
-        assert_eq!(b"5629".to_vec(), marias[1].1.pk);
-
-        // The associated data is in the (UniqueRef) values
-        assert_eq!(data1, marias[0].1.value);
-        assert_eq!(data3, marias[1].1.value);
+        // The associated data
+        assert_eq!(data1, marias[0].1);
+        assert_eq!(data3, marias[1].1);
     }
 }
diff --git a/packages/storage-plus/src/indexes.rs b/packages/storage-plus/src/indexes.rs
@@ -4,12 +4,12 @@
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 
-use cosmwasm_std::{Binary, Order, StdError, StdResult, Storage, KV};
+use cosmwasm_std::{from_slice, Binary, Order, StdError, StdResult, Storage, KV};
 
 use crate::keys::EmptyPrefix;
 use crate::map::Map;
 use crate::prefix::range_with_prefix;
-use crate::{Bound, PkOwned, Prefix, PrimaryKey};
+use crate::{Bound, PkOwned, Prefix, Prefixer, PrimaryKey};
 
 /// MARKER is stored in the multi-index as value, but we only look at the key (which is pk)
 const MARKER: u32 = 1;
@@ -123,15 +123,16 @@ where
 }
 
 #[derive(Deserialize, Serialize)]
-pub struct UniqueRef<T> {
+pub(crate) struct UniqueRef<T> {
     // note, we collapse the pk - combining everything under the namespace - even if it is composite
-    pub pk: Binary,
-    pub value: T,
+    pk: Binary,
+    value: T,
 }
 
 pub struct UniqueIndex<'a, K, T> {
     index: fn(&T) -> K,
     idx_map: Map<'a, K, UniqueRef<T>>,
+    idx_namespace: &'a [u8],
 }
 
 impl<'a, K, T> UniqueIndex<'a, K, T> {
@@ -140,6 +141,7 @@ impl<'a, K, T> UniqueIndex<'a, K, T> {
         UniqueIndex {
             index: idx_fn,
             idx_map: Map::new(idx_namespace),
+            idx_namespace: idx_namespace.as_bytes(),
         }
     }
 }
@@ -172,14 +174,19 @@ where
     }
 }
 
+pub(crate) fn deserialize_unique_kv<T: DeserializeOwned>(kv: KV) -> StdResult<KV<T>> {
+    let (_, v) = kv;
+    let t = from_slice::<UniqueRef<T>>(&v)?;
+    Ok((t.pk.into(), t.value))
+}
+
 impl<'a, K, T> UniqueIndex<'a, K, T>
 where
     T: Serialize + DeserializeOwned + Clone,
     K: PrimaryKey<'a>,
 {
-    pub fn prefix(&self, p: K::Prefix) -> Prefix<UniqueRef<T>> {
-        // Prefix::<T>::new(self.idx_namespace, &p.prefix())
-        self.idx_map.prefix(p)
+    pub fn prefix(&self, p: K::Prefix) -> 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
@@ -207,7 +214,7 @@ where
         min: Option<Bound>,
         max: Option<Bound>,
         order: cosmwasm_std::Order,
-    ) -> Box<dyn Iterator<Item = StdResult<KV<UniqueRef<T>>>> + 'c>
+    ) -> Box<dyn Iterator<Item = StdResult<KV<T>>> + 'c>
     where
         T: 'c,
     {