Multi index generic key

contracts/cw721-base/src/contract.rs

@@ -15,7 +15,7 @@ use crate::msg::{HandleMsg, InitMsg, MintMsg, MinterResponse, QueryMsg};
 use crate::state::{
     increment_tokens, num_tokens, tokens, Approval, TokenInfo, CONTRACT_INFO, MINTER, OPERATORS,
 };
-use cw_storage_plus::Bound;
+use cw_storage_plus::{Bound, PkOwned};
 
 // version info for migration info
 const CONTRACT_NAME: &str = "crates.io:cw721-base";
@@ -497,14 +497,22 @@ fn query_tokens(
     let start = start_after.map(Bound::exclusive);
 
     let owner_raw = deps.api.canonical_address(&owner)?;
-    let tokens: Result<Vec<String>, _> = tokens()
+    let res: Result<Vec<_>, _> = tokens()
         .idx
         .owner
-        .pks(deps.storage, &owner_raw, start, None, Order::Ascending)
+        .pks(
+            deps.storage,
+            PkOwned(owner_raw.into()),
+            start,
+            None,
+            Order::Ascending,
+        )
         .take(limit)
-        .map(String::from_utf8)
         .collect();
-    let tokens = tokens.map_err(StdError::invalid_utf8)?;
+    let pks = res?;
+
+    let res: Result<Vec<_>, _> = pks.iter().map(|v| String::from_utf8(v.to_vec())).collect();
+    let tokens = res.map_err(StdError::invalid_utf8)?;
     Ok(TokensResponse { tokens })
 }
 

contracts/cw721-base/src/state.rs

@@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize};
 
 use cosmwasm_std::{CanonicalAddr, StdResult, Storage};
 use cw721::{ContractInfoResponse, Expiration};
-use cw_storage_plus::{Index, IndexList, IndexedMap, Item, Map, MultiIndex};
+use cw_storage_plus::{Index, IndexList, IndexedMap, Item, Map, MultiIndex, PkOwned};
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 pub struct TokenInfo {
@@ -46,7 +46,8 @@ pub fn increment_tokens(storage: &mut dyn Storage) -> StdResult<u64> {
 }
 
 pub struct TokenIndexes<'a> {
-    pub owner: MultiIndex<'a, TokenInfo>,
+    // pk goes to second tuple element
+    pub owner: MultiIndex<'a, (PkOwned, PkOwned), TokenInfo>,
 }
 
 impl<'a> IndexList<TokenInfo> for TokenIndexes<'a> {
@@ -58,7 +59,11 @@ impl<'a> IndexList<TokenInfo> for TokenIndexes<'a> {
 
 pub fn tokens<'a>() -> IndexedMap<'a, &'a str, TokenInfo, TokenIndexes<'a>> {
     let indexes = TokenIndexes {
-        owner: MultiIndex::new(|d| d.owner.to_vec(), "tokens", "tokens__owner"),
+        owner: MultiIndex::new(
+            |d, k| (PkOwned(d.owner.to_vec()), PkOwned(k)),
+            "tokens",
+            "tokens__owner",
+        ),
     };
     IndexedMap::new("tokens", indexes)
 }

packages/storage-plus/src/indexed_map.rs

@@ -161,7 +161,7 @@ where
 mod test {
     use super::*;
 
-    use crate::indexes::{index_string, index_string_tuple, MultiIndex, UniqueIndex};
+    use crate::indexes::{index_string_tuple, index_triple, MultiIndex, UniqueIndex};
     use crate::{PkOwned, U32Key};
     use cosmwasm_std::testing::MockStorage;
     use cosmwasm_std::{MemoryStorage, Order};
@@ -175,7 +175,8 @@ mod test {
     }
 
     struct DataIndexes<'a> {
-        pub name: MultiIndex<'a, Data>,
+        // Second arg is for storing pk
+        pub name: MultiIndex<'a, (PkOwned, PkOwned), Data>,
         pub age: UniqueIndex<'a, U32Key, Data>,
         pub name_lastname: UniqueIndex<'a, (PkOwned, PkOwned), Data>,
     }
@@ -188,10 +189,28 @@ mod test {
         }
     }
 
+    // For composite multi index tests
+    struct DataCompositeMultiIndex<'a> {
+        // Third arg needed for storing pk
+        pub name_age: MultiIndex<'a, (PkOwned, U32Key, PkOwned), Data>,
+    }
+
+    // Future Note: this can likely be macro-derived
+    impl<'a> IndexList<Data> for DataCompositeMultiIndex<'a> {
+        fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<Data>> + '_> {
+            let v: Vec<&dyn Index<Data>> = vec![&self.name_age];
+            Box::new(v.into_iter())
+        }
+    }
+
     // Can we make it easier to define this? (less wordy generic)
     fn build_map<'a>() -> IndexedMap<'a, &'a [u8], Data, DataIndexes<'a>> {
         let indexes = DataIndexes {
-            name: MultiIndex::new(|d| index_string(&d.name), "data", "data__name"),
+            name: MultiIndex::new(
+                |d, k| (PkOwned(d.name.as_bytes().to_vec()), PkOwned(k)),
+                "data",
+                "data__name",
+            ),
             age: UniqueIndex::new(|d| U32Key::new(d.age), "data__age"),
             name_lastname: UniqueIndex::new(
                 |d| index_string_tuple(&d.name, &d.last_name),
@@ -269,19 +288,22 @@ mod test {
         let count = map
             .idx
             .name
-            .all_items(&store, &index_string("Maria"))
-            .unwrap()
+            .prefix(PkOwned(b"Maria".to_vec()))
+            .range(&store, None, None, Order::Ascending)
+            .collect::<Vec<_>>()
             .len();
         assert_eq!(2, 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 = map
+        // let marias: Vec<_>> = map
+        let marias: Vec<_> = map
             .idx
             .name
-            .all_items(&store, &index_string("Maria"))
+            .prefix(PkOwned(b"Maria".to_vec()))
+            .range(&store, None, None, Order::Ascending)
+            .collect::<StdResult<_>>()
             .unwrap();
         assert_eq!(2, marias.len());
         let (k, v) = &marias[0];
@@ -292,26 +314,29 @@ mod test {
         let count = map
             .idx
             .name
-            .all_items(&store, &index_string("Marib"))
-            .unwrap()
+            .prefix(PkOwned(b"Marib".to_vec()))
+            .range(&store, None, None, Order::Ascending)
+            .collect::<Vec<_>>()
             .len();
         assert_eq!(0, count);
 
         // other index doesn't match (1 byte before)
         let count = map
             .idx
             .name
-            .all_items(&store, &index_string("Mari`"))
-            .unwrap()
+            .prefix(PkOwned(b"Mari`".to_vec()))
+            .range(&store, None, None, Order::Ascending)
+            .collect::<Vec<_>>()
             .len();
         assert_eq!(0, count);
 
         // other index doesn't match (longer)
         let count = map
             .idx
             .name
-            .all_items(&store, &index_string("Maria5"))
-            .unwrap()
+            .prefix(PkOwned(b"Maria5".to_vec()))
+            .range(&store, None, None, Order::Ascending)
+            .collect::<Vec<_>>()
             .len();
         assert_eq!(0, count);
 
@@ -327,6 +352,127 @@ mod test {
         assert_eq!(None, aged);
     }
 
+    #[test]
+    fn range_simple_key_by_multi_index() {
+        let mut store = MockStorage::new();
+        let map = build_map();
+
+        // save data
+        let data1 = Data {
+            name: "Maria".to_string(),
+            last_name: "".to_string(),
+            age: 42,
+        };
+        let pk: &[u8] = b"5627";
+        map.save(&mut store, pk, &data1).unwrap();
+
+        let data2 = Data {
+            name: "Juan".to_string(),
+            last_name: "Perez".to_string(),
+            age: 13,
+        };
+        let pk: &[u8] = b"5628";
+        map.save(&mut store, pk, &data2).unwrap();
+
+        let data3 = Data {
+            name: "Maria".to_string(),
+            last_name: "Williams".to_string(),
+            age: 24,
+        };
+        let pk: &[u8] = b"5629";
+        map.save(&mut store, pk, &data3).unwrap();
+
+        let data4 = Data {
+            name: "Maria Luisa".to_string(),
+            last_name: "Bemberg".to_string(),
+            age: 12,
+        };
+        let pk: &[u8] = b"5630";
+        map.save(&mut store, pk, &data4).unwrap();
+
+        let marias: Vec<_> = map
+            .idx
+            .name
+            .prefix(PkOwned(b"Maria".to_vec()))
+            .range(&store, None, None, Order::Descending)
+            .collect::<StdResult<_>>()
+            .unwrap();
+        let count = marias.len();
+        assert_eq!(2, count);
+
+        // Sorted by (descending) pk
+        assert_eq!(marias[0].0, b"5629");
+        assert_eq!(marias[1].0, b"5627");
+        // Data is correct
+        assert_eq!(marias[0].1, data3);
+        assert_eq!(marias[1].1, data1);
+    }
+
+    #[test]
+    fn range_composite_key_by_multi_index() {
+        let mut store = MockStorage::new();
+
+        let indexes = DataCompositeMultiIndex {
+            name_age: MultiIndex::new(
+                |d, k| index_triple(&d.name, d.age, k),
+                "data",
+                "data__name_age",
+            ),
+        };
+        let map = IndexedMap::new("data", indexes);
+
+        // save data
+        let data1 = Data {
+            name: "Maria".to_string(),
+            last_name: "".to_string(),
+            age: 42,
+        };
+        let pk1: &[u8] = b"5627";
+        map.save(&mut store, pk1, &data1).unwrap();
+
+        let data2 = Data {
+            name: "Juan".to_string(),
+            last_name: "Perez".to_string(),
+            age: 13,
+        };
+        let pk2: &[u8] = b"5628";
+        map.save(&mut store, pk2, &data2).unwrap();
+
+        let data3 = Data {
+            name: "Maria".to_string(),
+            last_name: "Young".to_string(),
+            age: 24,
+        };
+        let pk3: &[u8] = b"5629";
+        map.save(&mut store, pk3, &data3).unwrap();
+
+        let data4 = Data {
+            name: "Maria Luisa".to_string(),
+            last_name: "Bemberg".to_string(),
+            age: 43,
+        };
+        let pk4: &[u8] = b"5630";
+        map.save(&mut store, pk4, &data4).unwrap();
+
+        let marias: Vec<_> = map
+            .idx
+            .name_age
+            .sub_prefix(PkOwned(b"Maria".to_vec()))
+            .range(&store, None, None, Order::Descending)
+            .collect::<StdResult<_>>()
+            .unwrap();
+        let count = marias.len();
+        assert_eq!(2, count);
+
+        // Pks (sorted by age descending)
+        assert_eq!(pk1, marias[0].0);
+        assert_eq!(pk3, marias[1].0);
+
+        // Data
+        assert_eq!(data1, marias[0].1);
+        assert_eq!(data3, marias[1].1);
+    }
+
     #[test]
     fn unique_index_enforced() {
         let mut store = MockStorage::new();
@@ -401,7 +547,13 @@ mod test {
          -> usize {
             map.idx
                 .name
-                .pks(store, &index_string(name), None, None, Order::Ascending)
+                .pks(
+                    store,
+                    PkOwned(name.as_bytes().to_vec()),
+                    None,
+                    None,
+                    Order::Ascending,
+                )
                 .count()
         };
 
@@ -436,10 +588,10 @@ mod test {
     #[test]
     fn unique_index_simple_key_range() {
         let mut store = MockStorage::new();
-        let mut map = build_map();
+        let map = build_map();
 
         // save data
-        let (pks, datas) = save_data(&mut store, &mut map);
+        let (pks, datas) = save_data(&mut store, &map);
 
         let res: StdResult<Vec<_>> = map
             .idx
@@ -467,10 +619,10 @@ mod test {
     #[test]
     fn unique_index_composite_key_range() {
         let mut store = MockStorage::new();
-        let mut map = build_map();
+        let map = build_map();
 
         // save data
-        let (pks, datas) = save_data(&mut store, &mut map);
+        let (pks, datas) = save_data(&mut store, &map);
 
         let res: StdResult<Vec<_>> = map
             .idx