Add primary key to MultiIndex index fn params

MIGRATING.md

@@ -2,6 +2,30 @@
 
 This guide lists API changes between *cw-plus* major releases.
 
+## v0.13.x -> v0.14.0
+
+### Breaking Issues / PRs
+
+- `MultiIndex` index fn params now include the pk [\#670](https://github.com/CosmWasm/cw-plus/issues/670)
+
+The `idx_fn` param of the `MultiIndex` constructor now has signature `fn(&[u8], &T) -> IK`, where the first param is the
+primary key of the index (in raw format), and the second param is the associated value.
+That allows us to use the pk or parts of it for building the index key.
+
+Migration of existing code is straight-forward. Just add an (unused) `_pk` param to the index function definition:
+
+```diff
+fn build_map<'a>() -> IndexedMap<'a, &'a str, Data, DataIndexes<'a>> {
+   let indexes = DataIndexes {
+-     name: MultiIndex::new(|d| d.name.clone(), "data", "data__name"),
++     name: MultiIndex::new(|_pk, d| d.name.clone(), "data", "data__name"),
+      age: UniqueIndex::new(|d| d.age, "data__age"),
+      name_lastname: UniqueIndex::new(
+```
+
+If you want to leverage this new functionality, take a look at the `pk_based_index()` test / example
+in `src/indexed_map.rs`.
+
 ## v0.11.0 -> v0.12.0
 
 ### Breaking Issues / PRs

packages/storage-macro/tests/index_list.rs

@@ -23,7 +23,7 @@ mod test {
         let _: IndexedMap<u64, TestStruct, TestIndexes> = IndexedMap::new(
             "t",
             TestIndexes {
-                id: MultiIndex::new(|t| t.id2, "t", "t_id2"),
+                id: MultiIndex::new(|_pk, t| t.id2, "t", "t_id2"),
                 addr: UniqueIndex::new(|t| t.addr.clone(), "t_addr"),
             },
         );
@@ -48,7 +48,7 @@ mod test {
         let idm: IndexedMap<u64, TestStruct, TestIndexes> = IndexedMap::new(
             "t",
             TestIndexes {
-                id: MultiIndex::new(|t| t.id2, "t", "t_2"),
+                id: MultiIndex::new(|_pk, t| t.id2, "t", "t_2"),
                 addr: UniqueIndex::new(|t| t.addr.clone(), "t_addr"),
             },
         );

packages/storage-plus/src/indexed_map.rs

@@ -322,7 +322,7 @@ mod test {
     // Can we make it easier to define this? (less wordy generic)
     fn build_map<'a>() -> IndexedMap<'a, &'a str, Data, DataIndexes<'a>> {
         let indexes = DataIndexes {
-            name: MultiIndex::new(|d| d.name.clone(), "data", "data__name"),
+            name: MultiIndex::new(|_pk, d| d.name.clone(), "data", "data__name"),
             age: UniqueIndex::new(|d| d.age, "data__age"),
             name_lastname: UniqueIndex::new(
                 |d| index_string_tuple(&d.name, &d.last_name),
@@ -651,7 +651,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map = IndexedMap::new("data", indexes);
 
@@ -712,7 +716,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map = IndexedMap::new("data", indexes);
 
@@ -1070,7 +1078,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map = IndexedMap::new("data", indexes);
 
@@ -1125,7 +1137,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map = IndexedMap::new("data", indexes);
 
@@ -1177,7 +1193,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map = IndexedMap::new("data", indexes);
 
@@ -1235,7 +1255,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map = IndexedMap::new("data", indexes);
 
@@ -1316,7 +1340,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map = IndexedMap::new("data", indexes);
 
@@ -1478,7 +1506,7 @@ mod test {
         fn composite_key_query() {
             let indexes = Indexes {
                 secondary: MultiIndex::new(
-                    |secondary| *secondary,
+                    |_pk, secondary| *secondary,
                     "test_map",
                     "test_map__secondary",
                 ),
@@ -1529,4 +1557,126 @@ mod test {
             );
         }
     }
+
+    mod pk_multi_index {
+        use super::*;
+        use cosmwasm_std::{Addr, Uint128};
+
+        struct Indexes<'a> {
+            // The last type param must match the `IndexedMap` primary key type below
+            spender: MultiIndex<'a, Addr, Uint128, (Addr, Addr)>,
+        }
+
+        impl<'a> IndexList<Uint128> for Indexes<'a> {
+            fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<Uint128>> + '_> {
+                let v: Vec<&dyn Index<Uint128>> = vec![&self.spender];
+                Box::new(v.into_iter())
+            }
+        }
+
+        #[test]
+        #[cfg(feature = "iterator")]
+        fn pk_based_index() {
+            fn pk_index(pk: &[u8]) -> Addr {
+                let (_owner, spender) = <(Addr, Addr)>::from_slice(pk).unwrap(); // mustn't fail
+                spender
+            }
+
+            let indexes = Indexes {
+                spender: MultiIndex::new(
+                    |pk, _allow| pk_index(pk),
+                    "allowances",
+                    "allowances__spender",
+                ),
+            };
+            let map = IndexedMap::<(&Addr, &Addr), Uint128, Indexes>::new("allowances", indexes);
+            let mut store = MockStorage::new();
+
+            map.save(
+                &mut store,
+                (&Addr::unchecked("owner1"), &Addr::unchecked("spender1")),
+                &Uint128::new(11),
+            )
+            .unwrap();
+            map.save(
+                &mut store,
+                (&Addr::unchecked("owner1"), &Addr::unchecked("spender2")),
+                &Uint128::new(12),
+            )
+            .unwrap();
+            map.save(
+                &mut store,
+                (&Addr::unchecked("owner2"), &Addr::unchecked("spender1")),
+                &Uint128::new(21),
+            )
+            .unwrap();
+
+            // Iterate over the main values
+            let items: Vec<_> = map
+                .range_raw(&store, None, None, Order::Ascending)
+                .collect::<Result<_, _>>()
+                .unwrap();
+
+            // Strip the index from values (for simpler comparison)
+            let items: Vec<_> = items.into_iter().map(|(_, v)| v.u128()).collect();
+
+            assert_eq!(items, vec![11, 12, 21]);
+
+            // Iterate over the indexed values
+            let items: Vec<_> = map
+                .idx
+                .spender
+                .range_raw(&store, None, None, Order::Ascending)
+                .collect::<Result<_, _>>()
+                .unwrap();
+
+            // Strip the index from values (for simpler comparison)
+            let items: Vec<_> = items.into_iter().map(|(_, v)| v.u128()).collect();
+
+            assert_eq!(items, vec![11, 21, 12]);
+
+            // Prefix over the main values
+            let items: Vec<_> = map
+                .prefix(&Addr::unchecked("owner1"))
+                .range_raw(&store, None, None, Order::Ascending)
+                .collect::<Result<_, _>>()
+                .unwrap();
+
+            // Strip the index from values (for simpler comparison)
+            let items: Vec<_> = items.into_iter().map(|(_, v)| v.u128()).collect();
+
+            assert_eq!(items, vec![11, 12]);
+
+            // Prefix over the indexed values
+            let items: Vec<_> = map
+                .idx
+                .spender
+                .prefix(Addr::unchecked("spender1"))
+                .range_raw(&store, None, None, Order::Ascending)
+                .collect::<Result<_, _>>()
+                .unwrap();
+
+            // Strip the index from values (for simpler comparison)
+            let items: Vec<_> = items.into_iter().map(|(_, v)| v.u128()).collect();
+
+            assert_eq!(items, vec![11, 21]);
+
+            // Prefix over the indexed values, and deserialize primary key as well
+            let items: Vec<_> = map
+                .idx
+                .spender
+                .prefix(Addr::unchecked("spender2"))
+                .range(&store, None, None, Order::Ascending)
+                .collect::<Result<_, _>>()
+                .unwrap();
+
+            assert_eq!(
+                items,
+                vec![(
+                    (Addr::unchecked("owner1"), Addr::unchecked("spender2")),
+                    Uint128::new(12)
+                )]
+            );
+        }
+    }
 }

packages/storage-plus/src/indexed_snapshot.rs

@@ -347,7 +347,7 @@ mod test {
     // Can we make it easier to define this? (less wordy generic)
     fn build_snapshot_map<'a>() -> IndexedSnapshotMap<'a, &'a str, Data, DataIndexes<'a>> {
         let indexes = DataIndexes {
-            name: MultiIndex::new(|d| d.name.as_bytes().to_vec(), "data", "data__name"),
+            name: MultiIndex::new(|_pk, d| d.name.as_bytes().to_vec(), "data", "data__name"),
             age: UniqueIndex::new(|d| d.age, "data__age"),
             name_lastname: UniqueIndex::new(
                 |d| index_string_tuple(&d.name, &d.last_name),
@@ -677,7 +677,11 @@ mod test {
         let mut height = 2;
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map =
             IndexedSnapshotMap::new("data", "checks", "changes", Strategy::EveryBlock, indexes);
@@ -743,7 +747,11 @@ mod test {
         let mut height = 2;
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map =
             IndexedSnapshotMap::new("data", "checks", "changes", Strategy::EveryBlock, indexes);
@@ -1134,7 +1142,11 @@ mod test {
         let mut store = MockStorage::new();
 
         let indexes = DataCompositeMultiIndex {
-            name_age: MultiIndex::new(|d| index_tuple(&d.name, d.age), "data", "data__name_age"),
+            name_age: MultiIndex::new(
+                |_pk, d| index_tuple(&d.name, d.age),
+                "data",
+                "data__name_age",
+            ),
         };
         let map =
             IndexedSnapshotMap::new("data", "checks", "changes", Strategy::EveryBlock, indexes);

packages/storage-plus/src/indexes/multi.rs

@@ -24,11 +24,11 @@ use std::marker::PhantomData;
 /// the secondary load of the associated value from the main map.
 ///
 /// The PK type defines the type of Primary Key, both for deserialization, and
-/// more important, type-safe bound key type.
+/// more important, as the type-safe bound key type.
 /// This type must match the encompassing `IndexedMap` primary key type,
 /// or its owned variant.
 pub struct MultiIndex<'a, IK, T, PK> {
-    index: fn(&T) -> IK,
+    index: fn(&[u8], &T) -> IK,
     idx_namespace: &'a [u8],
     // note, we collapse the ik - combining everything under the namespace - and concatenating the pk
     idx_map: Map<'a, Vec<u8>, u32>,
@@ -60,12 +60,12 @@ where
     /// }
     ///
     /// let index: MultiIndex<_, _, String> = MultiIndex::new(
-    ///     |d: &Data| d.age,
+    ///     |_pk: &[u8], d: &Data| d.age,
     ///     "age",
     ///     "age__owner",
     /// );
     /// ```
-    pub fn new(idx_fn: fn(&T) -> IK, pk_namespace: &'a str, idx_namespace: &'a str) -> Self {
+    pub fn new(idx_fn: fn(&[u8], &T) -> IK, pk_namespace: &'a str, idx_namespace: &'a str) -> Self {
         MultiIndex {
             index: idx_fn,
             idx_namespace: idx_namespace.as_bytes(),
@@ -131,12 +131,12 @@ where
     IK: PrimaryKey<'a>,
 {
     fn save(&self, store: &mut dyn Storage, pk: &[u8], data: &T) -> StdResult<()> {
-        let idx = (self.index)(data).joined_extra_key(pk);
+        let idx = (self.index)(pk, data).joined_extra_key(pk);
         self.idx_map.save(store, idx, &(pk.len() as u32))
     }
 
     fn remove(&self, store: &mut dyn Storage, pk: &[u8], old_data: &T) -> StdResult<()> {
-        let idx = (self.index)(old_data).joined_extra_key(pk);
+        let idx = (self.index)(pk, old_data).joined_extra_key(pk);
         self.idx_map.remove(store, idx);
         Ok(())
     }