diff --git a/packages/storage-plus/src/indexed_map.rs b/packages/storage-plus/src/indexed_map.rs index 1ac69c3f7..bbcfcfd9c 100644 --- a/packages/storage-plus/src/indexed_map.rs +++ b/packages/storage-plus/src/indexed_map.rs @@ -126,6 +126,11 @@ where pub fn prefix(&self, p: K::Prefix) -> Prefix { Prefix::new(self.pk_namespace, &p.prefix()) } + + // use sub_prefix to scan -> range + pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix { + Prefix::new(self.pk_namespace, &p.prefix()) + } } // short-cut for simple keys, rather than .prefix(()).range(...) diff --git a/packages/storage-plus/src/indexes.rs b/packages/storage-plus/src/indexes.rs index 12a721b32..edbfa59de 100644 --- a/packages/storage-plus/src/indexes.rs +++ b/packages/storage-plus/src/indexes.rs @@ -189,6 +189,10 @@ where Prefix::new_de_fn(self.idx_namespace, &p.prefix(), deserialize_unique_kv) } + pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix { + Prefix::new_de_fn(self.idx_namespace, &p.prefix(), deserialize_unique_kv) + } + /// returns all items that match this secondary index, always by pk Ascending pub fn item(&self, store: &dyn Storage, idx: K) -> StdResult>> { let data = self diff --git a/packages/storage-plus/src/keys.rs b/packages/storage-plus/src/keys.rs index 4e136acaf..0384babd7 100644 --- a/packages/storage-plus/src/keys.rs +++ b/packages/storage-plus/src/keys.rs @@ -7,9 +7,10 @@ use crate::Endian; // pub trait PrimaryKey<'a>: Copy { pub trait PrimaryKey<'a>: Clone { type Prefix: Prefixer<'a>; + type SubPrefix: Prefixer<'a>; /// returns a slice of key steps, which can be optionally combined - fn key<'b>(&'b self) -> Vec<&'b [u8]>; + fn key(&self) -> Vec<&[u8]>; fn joined_key(&self) -> Vec { let keys = self.key(); @@ -24,8 +25,9 @@ pub trait PrimaryKey<'a>: Clone { impl<'a> PrimaryKey<'a> for &'a [u8] { type Prefix = (); + type SubPrefix = (); - fn key<'b>(&'b self) -> Vec<&'b [u8]> { + fn key(&self) -> Vec<&[u8]> { // this is simple, we don't add more prefixes vec![self] } @@ -38,8 +40,9 @@ impl<'a> PrimaryKey<'a> for &'a [u8] { // Provide a string version of this to raw encode strings impl<'a> PrimaryKey<'a> for &'a str { type Prefix = (); + type SubPrefix = (); - fn key<'b>(&'b self) -> Vec<&'b [u8]> { + fn key(&self) -> Vec<&[u8]> { // this is simple, we don't add more prefixes vec![self.as_bytes()] } @@ -52,8 +55,9 @@ impl<'a> PrimaryKey<'a> for &'a str { // use generics for combining there - so we can use &[u8], PkOwned, or IntKey impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a>> PrimaryKey<'a> for (T, U) { type Prefix = T; + type SubPrefix = (); - fn key<'b>(&'b self) -> Vec<&'b [u8]> { + fn key(&self) -> Vec<&[u8]> { let mut keys = self.0.key(); keys.extend(&self.1.key()); keys @@ -72,6 +76,7 @@ impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a> + Prefixer<'a>, V: PrimaryKey<'a> for (T, U, V) { type Prefix = (T, U); + type SubPrefix = T; fn key(&self) -> Vec<&[u8]> { let mut keys = self.0.key(); @@ -97,23 +102,23 @@ impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a> + Prefixer<'a>, V: // pub trait Prefixer<'a>: Copy { pub trait Prefixer<'a> { /// returns 0 or more namespaces that should length-prefixed and concatenated for range searches - fn prefix<'b>(&'b self) -> Vec<&'b [u8]>; + fn prefix(&self) -> Vec<&[u8]>; } impl<'a> Prefixer<'a> for () { - fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { + fn prefix(&self) -> Vec<&[u8]> { vec![] } } impl<'a> Prefixer<'a> for &'a [u8] { - fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { + fn prefix(&self) -> Vec<&[u8]> { vec![self] } } impl<'a, T: Prefixer<'a>, U: Prefixer<'a>> Prefixer<'a> for (T, U) { - fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { + fn prefix(&self) -> Vec<&[u8]> { let mut res = self.0.prefix(); res.extend(self.1.prefix().into_iter()); res @@ -121,7 +126,7 @@ impl<'a, T: Prefixer<'a>, U: Prefixer<'a>> Prefixer<'a> for (T, U) { } impl<'a, T: Prefixer<'a>, U: Prefixer<'a>, V: Prefixer<'a>> Prefixer<'a> for (T, U, V) { - fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { + fn prefix(&self) -> Vec<&[u8]> { let mut res = self.0.prefix(); res.extend(self.1.prefix().into_iter()); res.extend(self.2.prefix().into_iter()); @@ -131,7 +136,7 @@ impl<'a, T: Prefixer<'a>, U: Prefixer<'a>, V: Prefixer<'a>> Prefixer<'a> for (T, // Provide a string version of this to raw encode strings impl<'a> Prefixer<'a> for &'a str { - fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { + fn prefix(&self) -> Vec<&[u8]> { vec![self.as_bytes()] } } @@ -151,8 +156,9 @@ pub struct PkOwned(pub Vec); impl<'a> PrimaryKey<'a> for PkOwned { type Prefix = (); + type SubPrefix = (); - fn key<'b>(&'b self) -> Vec<&'b [u8]> { + fn key(&self) -> Vec<&[u8]> { vec![&self.0] } @@ -162,7 +168,7 @@ impl<'a> PrimaryKey<'a> for PkOwned { } impl<'a> Prefixer<'a> for PkOwned { - fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { + fn prefix(&self) -> Vec<&[u8]> { vec![&self.0] } } @@ -170,8 +176,9 @@ impl<'a> Prefixer<'a> for PkOwned { // this auto-implements PrimaryKey for all the IntKey types (and more!) impl<'a, T: AsRef + From + Clone> PrimaryKey<'a> for T { type Prefix = (); + type SubPrefix = (); - fn key<'b>(&'b self) -> Vec<&'b [u8]> { + fn key(&self) -> Vec<&[u8]> { self.as_ref().key() } @@ -182,7 +189,7 @@ impl<'a, T: AsRef + From + Clone> PrimaryKey<'a> for T { // this auto-implements Prefixer for all the IntKey types (and more!) impl<'a, T: AsRef> Prefixer<'a> for T { - fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { + fn prefix(&self) -> Vec<&[u8]> { self.as_ref().prefix() } } diff --git a/packages/storage-plus/src/map.rs b/packages/storage-plus/src/map.rs index 6bfe03294..4e4c81f0c 100644 --- a/packages/storage-plus/src/map.rs +++ b/packages/storage-plus/src/map.rs @@ -42,6 +42,11 @@ where Prefix::new(self.namespace, &p.prefix()) } + #[cfg(feature = "iterator")] + pub fn sub_prefix(&self, p: K::SubPrefix) -> Prefix { + Prefix::new(self.namespace, &p.prefix()) + } + pub fn save(&self, store: &mut dyn Storage, k: K, data: &T) -> StdResult<()> { self.key(k).save(store, data) } @@ -104,6 +109,8 @@ mod test { use serde::{Deserialize, Serialize}; use std::ops::Deref; + #[cfg(feature = "iterator")] + use crate::iter_helpers::to_length_prefixed; use crate::U8Key; use cosmwasm_std::testing::MockStorage; #[cfg(feature = "iterator")] @@ -290,6 +297,9 @@ mod test { TRIPLE .save(&mut store, (b"owner", 9u8.into(), "recipient2"), &3000) .unwrap(); + TRIPLE + .save(&mut store, (b"owner", 10u8.into(), "recipient3"), &3000) + .unwrap(); TRIPLE .save(&mut store, (b"owner2", 9u8.into(), "recipient"), &5000) .unwrap(); @@ -308,6 +318,32 @@ mod test { (b"recipient2".to_vec(), 3000) ] ); + + // let's iterate over a sub prefix + let all: StdResult> = TRIPLE + .sub_prefix(b"owner") + .range(&store, None, None, Order::Ascending) + .collect(); + let all = all.unwrap(); + assert_eq!(3, all.len()); + // FIXME: range() works, but remaining keys are still encoded + assert_eq!( + all, + vec![ + ( + [to_length_prefixed(b"\x09"), b"recipient".to_vec()].concat(), + 1000 + ), + ( + [to_length_prefixed(b"\x09"), b"recipient2".to_vec()].concat(), + 3000 + ), + ( + [to_length_prefixed(b"\x0a"), b"recipient3".to_vec()].concat(), + 3000 + ) + ] + ); } #[test]