Add proper prefix_range helper when you want to iterate over the prefix space

packages/storage-plus/src/indexed_map.rs

@@ -6,9 +6,10 @@ use serde::de::DeserializeOwned;
 use serde::Serialize;
 
 use crate::indexes::Index;
+use crate::iter_helpers::deserialize_kv;
 use crate::keys::{Prefixer, PrimaryKey};
 use crate::map::Map;
-use crate::prefix::{Bound, Prefix};
+use crate::prefix::{namespaced_prefix_range, Bound, Prefix, PrefixBound};
 use crate::Path;
 
 pub trait IndexList<T> {
@@ -166,6 +167,34 @@ where
     }
 }
 
+#[cfg(feature = "iterator")]
+impl<'a, K, T, I> IndexedMap<'a, K, T, I>
+where
+    K: PrimaryKey<'a>,
+    T: Serialize + DeserializeOwned + Clone,
+    I: IndexList<T>,
+{
+    /// while range assumes you set the prefix to one element and call range over the last one,
+    /// prefix_range accepts bounds for the lowest and highest elements of the Prefix we wish to
+    /// accept, and iterates over those. There are some issues that distinguish these to and blindly
+    /// casting to Vec<u8> doesn't solve them.
+    pub fn prefix_range<'c>(
+        &self,
+        store: &'c dyn Storage,
+        min: Option<PrefixBound<'a, K::Prefix>>,
+        max: Option<PrefixBound<'a, K::Prefix>>,
+        order: cosmwasm_std::Order,
+    ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::Pair<T>>> + 'c>
+    where
+        T: 'c,
+        'a: 'c,
+    {
+        let mapped =
+            namespaced_prefix_range(store, self.pk_namespace, min, max, order).map(deserialize_kv);
+        Box::new(mapped)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -700,4 +729,53 @@ mod test {
         assert_eq!(datas[0], marias[0].1);
         assert_eq!(datas[1], marias[1].1);
     }
+
+    mod inclusive_bound {
+        use super::*;
+        use crate::U64Key;
+
+        struct Indexes<'a> {
+            secondary: MultiIndex<'a, (U64Key, Vec<u8>), u64>,
+        }
+
+        impl<'a> IndexList<u64> for Indexes<'a> {
+            fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<u64>> + '_> {
+                let v: Vec<&dyn Index<u64>> = vec![&self.secondary];
+                Box::new(v.into_iter())
+            }
+        }
+
+        #[test]
+        #[cfg(feature = "iterator")]
+        fn composite_key_failure() {
+            let indexes = Indexes {
+                secondary: MultiIndex::new(
+                    |secondary, k| (U64Key::new(*secondary), k),
+                    "test_map",
+                    "test_map__secondary",
+                ),
+            };
+            let map = IndexedMap::<&str, u64, Indexes>::new("test_map", indexes);
+            let mut store = MockStorage::new();
+
+            map.save(&mut store, "one", &1).unwrap();
+
+            let items: Vec<_> = map
+                .idx
+                .secondary
+                .prefix_range(
+                    &store,
+                    None,
+                    Some(PrefixBound::inclusive(1)),
+                    Order::Ascending,
+                )
+                .collect::<Result<_, _>>()
+                .unwrap();
+
+            // Strip the index from values (for simpler comparison)
+            let items: Vec<_> = items.into_iter().map(|(_, v)| v).collect();
+
+            assert_eq!(items, vec![1]);
+        }
+    }
 }

packages/storage-plus/src/indexes.rs

@@ -8,6 +8,7 @@ use cosmwasm_std::{from_slice, Binary, Order, Pair, StdError, StdResult, Storage
 
 use crate::helpers::namespaces_with_key;
 use crate::map::Map;
+use crate::prefix::{namespaced_prefix_range, PrefixBound};
 use crate::{Bound, Prefix, Prefixer, PrimaryKey, U32Key};
 
 pub fn index_string(data: &str) -> Vec<u8> {
@@ -225,6 +226,26 @@ where
     ) -> Box<dyn Iterator<Item = Vec<u8>> + 'c> {
         self.no_prefix().keys(store, min, max, order)
     }
+
+    /// while range assumes you set the prefix to one element and call range over the last one,
+    /// prefix_range accepts bounds for the lowest and highest elements of the Prefix we wish to
+    /// accept, and iterates over those. There are some issues that distinguish these to and blindly
+    /// casting to Vec<u8> doesn't solve them.
+    pub fn prefix_range<'c>(
+        &'c self,
+        store: &'c dyn Storage,
+        min: Option<PrefixBound<'a, K::Prefix>>,
+        max: Option<PrefixBound<'a, K::Prefix>>,
+        order: cosmwasm_std::Order,
+    ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::Pair<T>>> + 'c>
+    where
+        T: 'c,
+        'a: 'c,
+    {
+        let mapped = namespaced_prefix_range(store, self.idx_namespace, min, max, order)
+            .map(move |kv| (deserialize_multi_kv)(store, self.pk_namespace, kv));
+        Box::new(mapped)
+    }
 }
 
 /// UniqueRef stores Binary(Vec[u8]) representation of private key and index value

packages/storage-plus/src/keys.rs

@@ -81,6 +81,11 @@ impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a> + Prefixer<'a>, V:
 pub trait Prefixer<'a> {
     /// returns 0 or more namespaces that should length-prefixed and concatenated for range searches
     fn prefix(&self) -> Vec<&[u8]>;
+
+    fn joined_prefix(&self) -> Vec<u8> {
+        let prefixes = self.prefix();
+        namespaces_with_key(&prefixes, &[])
+    }
 }
 
 impl<'a> Prefixer<'a> for () {

packages/storage-plus/src/map.rs

@@ -4,10 +4,14 @@ use std::marker::PhantomData;
 
 use crate::helpers::query_raw;
 #[cfg(feature = "iterator")]
+use crate::iter_helpers::deserialize_kv;
+#[cfg(feature = "iterator")]
 use crate::keys::Prefixer;
 use crate::keys::PrimaryKey;
 use crate::path::Path;
 #[cfg(feature = "iterator")]
+use crate::prefix::{namespaced_prefix_range, PrefixBound};
+#[cfg(feature = "iterator")]
 use crate::prefix::{Bound, Prefix};
 use cosmwasm_std::{from_slice, Addr, QuerierWrapper, StdError, StdResult, Storage};
 
@@ -113,6 +117,8 @@ where
 impl<'a, K, T> Map<'a, K, T>
 where
     T: Serialize + DeserializeOwned,
+    // TODO: this should only be when K::Prefix == ()
+    // Other cases need to call prefix() first
     K: PrimaryKey<'a>,
 {
     pub fn range<'c>(
@@ -142,6 +148,33 @@ where
     }
 }
 
+#[cfg(feature = "iterator")]
+impl<'a, K, T> Map<'a, K, T>
+where
+    T: Serialize + DeserializeOwned,
+    K: PrimaryKey<'a>,
+{
+    /// while range assumes you set the prefix to one element and call range over the last one,
+    /// prefix_range accepts bounds for the lowest and highest elements of the Prefix we wish to
+    /// accept, and iterates over those. There are some issues that distinguish these to and blindly
+    /// casting to Vec<u8> doesn't solve them.
+    pub fn prefix_range<'c>(
+        &self,
+        store: &'c dyn Storage,
+        min: Option<PrefixBound<'a, K::Prefix>>,
+        max: Option<PrefixBound<'a, K::Prefix>>,
+        order: cosmwasm_std::Order,
+    ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::Pair<T>>> + 'c>
+    where
+        T: 'c,
+        'a: 'c,
+    {
+        let mapped =
+            namespaced_prefix_range(store, self.namespace, min, max, order).map(deserialize_kv);
+        Box::new(mapped)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -150,6 +183,8 @@ mod test {
 
     #[cfg(feature = "iterator")]
     use crate::iter_helpers::to_length_prefixed;
+    #[cfg(feature = "iterator")]
+    use crate::U32Key;
     use crate::U8Key;
     use cosmwasm_std::testing::MockStorage;
     #[cfg(feature = "iterator")]
@@ -662,4 +697,97 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    #[cfg(feature = "iterator")]
+    fn prefixed_range_works() {
+        // this is designed to look as much like a secondary index as possible
+        // we want to query over a range of u32 for the first key and all subkeys
+        const AGES: Map<(U32Key, Vec<u8>), u64> = Map::new("ages");
+
+        let mut store = MockStorage::new();
+        AGES.save(&mut store, (2.into(), vec![1, 2, 3]), &123)
+            .unwrap();
+        AGES.save(&mut store, (3.into(), vec![4, 5, 6]), &456)
+            .unwrap();
+        AGES.save(&mut store, (5.into(), vec![7, 8, 9]), &789)
+            .unwrap();
+        AGES.save(&mut store, (5.into(), vec![9, 8, 7]), &987)
+            .unwrap();
+        AGES.save(&mut store, (7.into(), vec![20, 21, 22]), &2002)
+            .unwrap();
+        AGES.save(&mut store, (8.into(), vec![23, 24, 25]), &2332)
+            .unwrap();
+
+        // typical range under one prefix as a control
+        let fives = AGES
+            .prefix(5.into())
+            .range(&store, None, None, Order::Ascending)
+            .collect::<StdResult<Vec<_>>>()
+            .unwrap();
+        assert_eq!(fives.len(), 2);
+        assert_eq!(fives, vec![(vec![7, 8, 9], 789), (vec![9, 8, 7], 987)]);
+
+        let keys: Vec<_> = AGES
+            .no_prefix()
+            .keys(&store, None, None, Order::Ascending)
+            .collect();
+        println!("keys: {:?}", keys);
+
+        // using inclusive bounds both sides
+        let include = AGES
+            .prefix_range(
+                &store,
+                Some(PrefixBound::inclusive(3)),
+                Some(PrefixBound::inclusive(7)),
+                Order::Ascending,
+            )
+            .map(|r| r.map(|(_, v)| v))
+            .collect::<StdResult<Vec<_>>>()
+            .unwrap();
+        assert_eq!(include.len(), 4);
+        assert_eq!(include, vec![456, 789, 987, 2002]);
+
+        // using exclusive bounds both sides
+        let exclude = AGES
+            .prefix_range(
+                &store,
+                Some(PrefixBound::exclusive(3)),
+                Some(PrefixBound::exclusive(7)),
+                Order::Ascending,
+            )
+            .map(|r| r.map(|(_, v)| v))
+            .collect::<StdResult<Vec<_>>>()
+            .unwrap();
+        assert_eq!(exclude.len(), 2);
+        assert_eq!(exclude, vec![789, 987]);
+
+        // using inclusive in descending
+        let include = AGES
+            .prefix_range(
+                &store,
+                Some(PrefixBound::inclusive(3)),
+                Some(PrefixBound::inclusive(5)),
+                Order::Descending,
+            )
+            .map(|r| r.map(|(_, v)| v))
+            .collect::<StdResult<Vec<_>>>()
+            .unwrap();
+        assert_eq!(include.len(), 3);
+        assert_eq!(include, vec![987, 789, 456]);
+
+        // using exclusive in descending
+        let include = AGES
+            .prefix_range(
+                &store,
+                Some(PrefixBound::exclusive(2)),
+                Some(PrefixBound::exclusive(5)),
+                Order::Descending,
+            )
+            .map(|r| r.map(|(_, v)| v))
+            .collect::<StdResult<Vec<_>>>()
+            .unwrap();
+        assert_eq!(include.len(), 1);
+        assert_eq!(include, vec![456]);
+    }
 }