Merge pull request #116 from CosmWasm/cleaner-ranges

Use Include/Exclude Bounds to define range searches

packages/storage-plus/README.md

@@ -288,9 +288,100 @@ fn demo() -> StdResult<()> {
 
 ### Prefix 
 
-We
+In addition to getting one particular item out of a map, we can iterate over the map
+(or a subset of the map). This let's us answer questions like "show me all tokens",
+and we provide some nice `Bound`s helpers to easily allow pagination or custom ranges.
 
-**TODO**
+The general format is to get a `Prefix` by calling `map.prefix(k)`, where `k` is exactly
+one less item than the normal key (If `map.key()` took `(&[u8], &[u8])`, then `map.prefix()` takes `&[u8]`.
+If `map.key()` took `&[u8]`, `map.prefix()` takes `()`). Once we have a prefix space, we can iterate
+over all items with `range(store, min, max, order)`. It supports `Order::Ascending` or `Order::Descending`.
+`min` is the lower bound and `max` is the higher bound.
+
+```rust
+#[derive(Copy, Clone, Debug)]
+pub enum Bound<'a> {
+    Inclusive(&'a [u8]),
+    Exclusive(&'a [u8]),
+    None,
+}
+```
+
+If the `min` and `max` bounds, it will return all items under this prefix. You can use `.take(n)` to
+limit the results to `n` items and start doing pagination. You can also set the `min` bound to
+eg. `Bound::Exclusive(last_value)` to start iterating over all items *after* the last value. Combined with
+`take`, we easily have pagination support. You can also use `Bound::Inclusive(x)` when you want to include any
+perfect matches. To better understand the API, please read the following example:
+
+```rust
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
+struct Data {
+    pub name: String,
+    pub age: i32,
+}
+
+const PEOPLE: Map<&[u8], Data> = Map::new(b"people");
+const ALLOWANCE: Map<(&[u8], &[u8]), u64> = Map::new(b"allow");
+
+fn demo() -> StdResult<()> {
+    let mut store = MockStorage::new();
+
+    // save and load on two keys
+    let data = Data { name: "John".to_string(), age: 32 };
+    PEOPLE.save(&mut store, b"john", &data)?;
+    let data2 = Data { name: "Jim".to_string(), age: 44 };
+    PEOPLE.save(&mut store, b"jim", &data2)?;
+
+    // iterate over them all
+    let all: StdResult<Vec<_>> = PEOPLE
+        .range(&store, Bound::None, Bound::None, Order::Ascending)
+        .collect();
+    assert_eq!(
+        all?,
+        vec![(b"jim".to_vec(), data2), (b"john".to_vec(), data.clone())]
+    );
+
+    // or just show what is after jim
+    let all: StdResult<Vec<_>> = PEOPLE
+        .range(
+            &store,
+            Bound::Exclusive(b"jim"),
+            Bound::None,
+            Order::Ascending,
+        )
+        .collect();
+    assert_eq!(all?, vec![(b"john".to_vec(), data)]);
+
+    // save and load on three keys, one under different owner
+    ALLOWANCE.save(&mut store, (b"owner", b"spender"), &1000)?;
+    ALLOWANCE.save(&mut store, (b"owner", b"spender2"), &3000)?;
+    ALLOWANCE.save(&mut store, (b"owner2", b"spender"), &5000)?;
+
+    // get all under one key
+    let all: StdResult<Vec<_>> = ALLOWANCE
+        .prefix(b"owner")
+        .range(&store, Bound::None, Bound::None, Order::Ascending)
+        .collect();
+    assert_eq!(
+        all?,
+        vec![(b"spender".to_vec(), 1000), (b"spender2".to_vec(), 3000)]
+    );
+
+    // Or ranges between two items (even reverse)
+    let all: StdResult<Vec<_>> = ALLOWANCE
+        .prefix(b"owner")
+        .range(
+            &store,
+            Bound::Exclusive(b"spender1"),
+            Bound::Inclusive(b"spender2"),
+            Order::Descending,
+        )
+        .collect();
+    assert_eq!(all?, vec![(b"spender2".to_vec(), 3000)]);
+
+    Ok(())
+}
+```
 
 ## Indexed Map
 

packages/storage-plus/src/iter_helpers.rs

@@ -2,8 +2,8 @@
 
 use serde::de::DeserializeOwned;
 
+use cosmwasm_std::KV;
 use cosmwasm_std::{from_slice, StdResult};
-use cosmwasm_std::{Order, Storage, KV};
 
 use crate::helpers::encode_length;
 
@@ -15,69 +15,26 @@ pub(crate) fn deserialize_kv<T: DeserializeOwned>(kv: KV) -> StdResult<KV<T>> {
 
 /// Calculates the raw key prefix for a given namespace as documented
 /// in https://github.com/webmaster128/key-namespacing#length-prefixed-keys
+#[allow(dead_code)]
 pub(crate) fn to_length_prefixed(namespace: &[u8]) -> Vec<u8> {
     let mut out = Vec::with_capacity(namespace.len() + 2);
     out.extend_from_slice(&encode_length(namespace));
     out.extend_from_slice(namespace);
     out
 }
 
-pub(crate) fn range_with_prefix<'a, S: Storage>(
-    storage: &'a S,
-    namespace: &[u8],
-    start: Option<&[u8]>,
-    end: Option<&[u8]>,
-    order: Order,
-) -> Box<dyn Iterator<Item = KV> + 'a> {
-    // prepare start, end with prefix
-    let start = match start {
-        Some(s) => concat(namespace, s),
-        None => namespace.to_vec(),
-    };
-    let end = match end {
-        Some(e) => concat(namespace, e),
-        // end is updating last byte by one
-        None => namespace_upper_bound(namespace),
-    };
-
-    // 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 prefix = namespace.to_vec();
-    let mapped = base_iterator.map(move |(k, v)| (trim(&prefix, &k), v));
-    Box::new(mapped)
-}
-
 #[inline]
-fn trim(namespace: &[u8], key: &[u8]) -> Vec<u8> {
+pub(crate) fn trim(namespace: &[u8], key: &[u8]) -> Vec<u8> {
     key[namespace.len()..].to_vec()
 }
 
 #[inline]
-fn concat(namespace: &[u8], key: &[u8]) -> Vec<u8> {
+pub(crate) fn concat(namespace: &[u8], key: &[u8]) -> Vec<u8> {
     let mut k = namespace.to_vec();
     k.extend_from_slice(key);
     k
 }
 
-/// 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<u8> {
-    let mut copy = input.to_vec();
-    // zero out all trailing 255, increment first that is not such
-    for i in (0..input.len()).rev() {
-        if copy[i] == 255 {
-            copy[i] = 0;
-        } else {
-            copy[i] += 1;
-            break;
-        }
-    }
-    copy
-}
-
 #[cfg(test)]
 mod test {
     use super::*;

packages/storage-plus/src/map.rs

@@ -2,10 +2,12 @@ use serde::de::DeserializeOwned;
 use serde::Serialize;
 use std::marker::PhantomData;
 
+#[cfg(feature = "iterator")]
+use crate::keys::Prefixer;
 use crate::keys::PrimaryKey;
 use crate::path::Path;
 #[cfg(feature = "iterator")]
-use crate::{Prefix, Prefixer};
+use crate::prefix::{Bound, Prefix};
 use cosmwasm_std::{StdError, StdResult, Storage};
 
 pub struct Map<'a, K, T> {
@@ -72,7 +74,7 @@ where
     }
 }
 
-/// short-cut for simple keys, rather than .prefix(()).range(...)
+// short-cut for simple keys, rather than .prefix(()).range(...)
 #[cfg(feature = "iterator")]
 impl<'a, T> Map<'a, &'a [u8], T>
 where
@@ -81,18 +83,16 @@ where
     // I would prefer not to copy code from Prefix, but no other way
     // with lifetimes (create Prefix inside function and return ref = no no)
     pub fn range<'c, S: Storage>(
-        &'c self,
+        &self,
         store: &'c S,
-        start: Option<&[u8]>,
-        end: Option<&[u8]>,
+        min: Bound<'_>,
+        max: Bound<'_>,
         order: cosmwasm_std::Order,
-    ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::KV<T>>> + 'c> {
-        // put the imports here, so we don't have to feature flag them above
-        use crate::iter_helpers::{deserialize_kv, range_with_prefix, to_length_prefixed};
-
-        let prefix = to_length_prefixed(self.namespace);
-        let mapped = range_with_prefix(store, &prefix, start, end, order).map(deserialize_kv::<T>);
-        Box::new(mapped)
+    ) -> Box<dyn Iterator<Item = StdResult<cosmwasm_std::KV<T>>> + 'c>
+    where
+        T: 'c,
+    {
+        self.prefix(()).range(store, min, max, order)
     }
 }
 
@@ -196,7 +196,9 @@ mod test {
         PEOPLE.save(&mut store, b"jim", &data2).unwrap();
 
         // let's try to iterate!
-        let all: StdResult<Vec<_>> = PEOPLE.range(&store, None, None, Order::Ascending).collect();
+        let all: StdResult<Vec<_>> = PEOPLE
+            .range(&store, Bound::None, Bound::None, Order::Ascending)
+            .collect();
         let all = all.unwrap();
         assert_eq!(2, all.len());
         assert_eq!(
@@ -224,7 +226,7 @@ mod test {
         // let's try to iterate!
         let all: StdResult<Vec<_>> = ALLOWANCE
             .prefix(b"owner")
-            .range(&store, None, None, Order::Ascending)
+            .range(&store, Bound::None, Bound::None, Order::Ascending)
             .collect();
         let all = all.unwrap();
         assert_eq!(2, all.len());
@@ -359,4 +361,71 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    #[cfg(feature = "iterator")]
+    fn readme_with_range() -> StdResult<()> {
+        let mut store = MockStorage::new();
+
+        // save and load on two keys
+        let data = Data {
+            name: "John".to_string(),
+            age: 32,
+        };
+        PEOPLE.save(&mut store, b"john", &data)?;
+        let data2 = Data {
+            name: "Jim".to_string(),
+            age: 44,
+        };
+        PEOPLE.save(&mut store, b"jim", &data2)?;
+
+        // iterate over them all
+        let all: StdResult<Vec<_>> = PEOPLE
+            .range(&store, Bound::None, Bound::None, Order::Ascending)
+            .collect();
+        assert_eq!(
+            all?,
+            vec![(b"jim".to_vec(), data2), (b"john".to_vec(), data.clone())]
+        );
+
+        // or just show what is after jim
+        let all: StdResult<Vec<_>> = PEOPLE
+            .range(
+                &store,
+                Bound::Exclusive(b"jim"),
+                Bound::None,
+                Order::Ascending,
+            )
+            .collect();
+        assert_eq!(all?, vec![(b"john".to_vec(), data)]);
+
+        // save and load on three keys, one under different owner
+        ALLOWANCE.save(&mut store, (b"owner", b"spender"), &1000)?;
+        ALLOWANCE.save(&mut store, (b"owner", b"spender2"), &3000)?;
+        ALLOWANCE.save(&mut store, (b"owner2", b"spender"), &5000)?;
+
+        // get all under one key
+        let all: StdResult<Vec<_>> = ALLOWANCE
+            .prefix(b"owner")
+            .range(&store, Bound::None, Bound::None, Order::Ascending)
+            .collect();
+        assert_eq!(
+            all?,
+            vec![(b"spender".to_vec(), 1000), (b"spender2".to_vec(), 3000)]
+        );
+
+        // Or ranges between two items (even reverse)
+        let all: StdResult<Vec<_>> = ALLOWANCE
+            .prefix(b"owner")
+            .range(
+                &store,
+                Bound::Exclusive(b"spender1"),
+                Bound::Inclusive(b"spender2"),
+                Order::Descending,
+            )
+            .collect();
+        assert_eq!(all?, vec![(b"spender2".to_vec(), 3000)]);
+
+        Ok(())
+    }
 }