diff --git a/packages/storage-plus/src/indexed_map.rs b/packages/storage-plus/src/indexed_map.rs index 6aded9c85..b74948694 100644 --- a/packages/storage-plus/src/indexed_map.rs +++ b/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 { @@ -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, +{ + /// 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 doesn't solve them. + pub fn prefix_range<'c>( + &self, + store: &'c dyn Storage, + min: Option>, + max: Option>, + order: cosmwasm_std::Order, + ) -> Box>> + '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,54 @@ 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), u64>, + } + + impl<'a> IndexList for Indexes<'a> { + fn get_indexes(&'_ self) -> Box> + '_> { + let v: Vec<&dyn Index> = vec![&self.secondary]; + Box::new(v.into_iter()) + } + } + + #[test] + #[cfg(feature = "iterator")] + fn composite_key_query() { + 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(); + map.save(&mut store, "two", &2).unwrap(); + + let items: Vec<_> = map + .idx + .secondary + .prefix_range( + &store, + None, + Some(PrefixBound::inclusive(1)), + Order::Ascending, + ) + .collect::>() + .unwrap(); + + // Strip the index from values (for simpler comparison) + let items: Vec<_> = items.into_iter().map(|(_, v)| v).collect(); + + assert_eq!(items, vec![1]); + } + } } diff --git a/packages/storage-plus/src/indexes.rs b/packages/storage-plus/src/indexes.rs index db465102d..fdd2d15a6 100644 --- a/packages/storage-plus/src/indexes.rs +++ b/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 { @@ -225,6 +226,26 @@ where ) -> Box> + '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 doesn't solve them. + pub fn prefix_range<'c>( + &'c self, + store: &'c dyn Storage, + min: Option>, + max: Option>, + order: cosmwasm_std::Order, + ) -> Box>> + '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 diff --git a/packages/storage-plus/src/keys.rs b/packages/storage-plus/src/keys.rs index 305de7876..bdae2bb07 100644 --- a/packages/storage-plus/src/keys.rs +++ b/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 { + let prefixes = self.prefix(); + namespaces_with_key(&prefixes, &[]) + } } impl<'a> Prefixer<'a> for () { diff --git a/packages/storage-plus/src/map.rs b/packages/storage-plus/src/map.rs index 3154a5204..6473cde1a 100644 --- a/packages/storage-plus/src/map.rs +++ b/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 doesn't solve them. + pub fn prefix_range<'c>( + &self, + store: &'c dyn Storage, + min: Option>, + max: Option>, + order: cosmwasm_std::Order, + ) -> Box>> + '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), 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::>>() + .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::>>() + .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::>>() + .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::>>() + .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::>>() + .unwrap(); + assert_eq!(include.len(), 1); + assert_eq!(include, vec![456]); + } } diff --git a/packages/storage-plus/src/prefix.rs b/packages/storage-plus/src/prefix.rs index e505d4760..a1faacc12 100644 --- a/packages/storage-plus/src/prefix.rs +++ b/packages/storage-plus/src/prefix.rs @@ -6,9 +6,9 @@ use std::marker::PhantomData; use cosmwasm_std::{Order, Pair, StdResult, Storage}; use std::ops::Deref; -use crate::helpers::nested_namespaces_with_key; +use crate::helpers::{namespaces_with_key, nested_namespaces_with_key}; use crate::iter_helpers::{concat, deserialize_kv, trim}; -use crate::Endian; +use crate::{Endian, Prefixer}; /// Bound is used to defines the two ends of a range, more explicit than Option /// None means that we don't limit that side of the range at all. @@ -42,8 +42,39 @@ impl Bound { } } +#[derive(Clone, Debug)] +pub enum PrefixBound<'a, K: Prefixer<'a>> { + Inclusive((K, PhantomData<&'a bool>)), + Exclusive((K, PhantomData<&'a bool>)), +} + +impl<'a, K: Prefixer<'a>> PrefixBound<'a, K> { + pub fn inclusive>(k: T) -> Self { + Self::Inclusive((k.into(), PhantomData)) + } + + pub fn exclusive>(k: T) -> Self { + Self::Exclusive((k.into(), PhantomData)) + } + + pub fn to_bound(&self) -> Bound { + match self { + PrefixBound::Exclusive((k, _)) => Bound::Exclusive(k.joined_prefix()), + PrefixBound::Inclusive((k, _)) => Bound::Inclusive(k.joined_prefix()), + } + } +} + type DeserializeFn = fn(&dyn Storage, &[u8], Pair) -> StdResult>; +pub fn default_deserializer( + _: &dyn Storage, + _: &[u8], + raw: Pair, +) -> StdResult> { + deserialize_kv(raw) +} + #[derive(Clone)] pub struct Prefix where @@ -73,9 +104,7 @@ where T: Serialize + DeserializeOwned, { pub fn new(top_name: &[u8], sub_names: &[&[u8]]) -> Self { - Prefix::with_deserialization_function(top_name, sub_names, &[], |_, _, kv| { - deserialize_kv(kv) - }) + Prefix::with_deserialization_function(top_name, sub_names, &[], default_deserializer) } pub fn with_deserialization_function( @@ -148,20 +177,63 @@ fn calc_start_bound(namespace: &[u8], bound: Option) -> Vec { None => namespace.to_vec(), // this is the natural limits of the underlying Storage Some(Bound::Inclusive(limit)) => concat(namespace, &limit), - Some(Bound::Exclusive(limit)) => concat(namespace, &one_byte_higher(&limit)), + Some(Bound::Exclusive(limit)) => concat(namespace, &extend_one_byte(&limit)), } } fn calc_end_bound(namespace: &[u8], bound: Option) -> Vec { match bound { - None => namespace_upper_bound(namespace), + None => increment_last_byte(namespace), + // this is the natural limits of the underlying Storage + Some(Bound::Exclusive(limit)) => concat(namespace, &limit), + Some(Bound::Inclusive(limit)) => concat(namespace, &extend_one_byte(&limit)), + } +} + +pub fn namespaced_prefix_range<'a, 'c, K: Prefixer<'a>>( + storage: &'c dyn Storage, + namespace: &[u8], + start: Option>, + end: Option>, + order: Order, +) -> Box + 'c> { + let prefix = namespaces_with_key(&[namespace], &[]); + let start = calc_prefix_start_bound(&prefix, start); + let end = calc_prefix_end_bound(&prefix, end); + + // get iterator from storage + let base_iterator = storage.range(Some(&start), Some(&end), order); + + // make a copy for the closure to handle lifetimes safely + let mapped = base_iterator.map(move |(k, v)| (trim(&prefix, &k), v)); + Box::new(mapped) +} + +fn calc_prefix_start_bound<'a, K: Prefixer<'a>>( + namespace: &[u8], + bound: Option>, +) -> Vec { + match bound.map(|b| b.to_bound()) { + None => namespace.to_vec(), + // this is the natural limits of the underlying Storage + Some(Bound::Inclusive(limit)) => concat(namespace, &limit), + Some(Bound::Exclusive(limit)) => concat(namespace, &increment_last_byte(&limit)), + } +} + +fn calc_prefix_end_bound<'a, K: Prefixer<'a>>( + namespace: &[u8], + bound: Option>, +) -> Vec { + match bound.map(|b| b.to_bound()) { + None => increment_last_byte(namespace), // this is the natural limits of the underlying Storage Some(Bound::Exclusive(limit)) => concat(namespace, &limit), - Some(Bound::Inclusive(limit)) => concat(namespace, &one_byte_higher(&limit)), + Some(Bound::Inclusive(limit)) => concat(namespace, &increment_last_byte(&limit)), } } -fn one_byte_higher(limit: &[u8]) -> Vec { +fn extend_one_byte(limit: &[u8]) -> Vec { let mut v = limit.to_vec(); v.push(0); v @@ -170,7 +242,7 @@ fn one_byte_higher(limit: &[u8]) -> Vec { /// Returns a new vec of same length and last byte incremented by one /// If last bytes are 255, we handle overflow up the chain. /// If all bytes are 255, this returns wrong data - but that is never possible as a namespace -fn namespace_upper_bound(input: &[u8]) -> Vec { +fn increment_last_byte(input: &[u8]) -> Vec { let mut copy = input.to_vec(); // zero out all trailing 255, increment first that is not such for i in (0..input.len()).rev() {