diff --git a/packages/storage-plus/src/keys.rs b/packages/storage-plus/src/keys.rs index edaf93196..4e136acaf 100644 --- a/packages/storage-plus/src/keys.rs +++ b/packages/storage-plus/src/keys.rs @@ -22,16 +22,8 @@ pub trait PrimaryKey<'a>: Clone { fn parse_key(serialized: &'a [u8]) -> Self; } -// optional type aliases to refer to them easier -type Pk0 = (); -type Pk1<'a> = &'a [u8]; -type Pk2<'a, T = &'a [u8], U = &'a [u8]> = (T, U); -type Pk3<'a, T = &'a [u8], U = &'a [u8], V = &'a [u8]> = (T, U, V); - -type PkStr<'a> = &'a str; - -impl<'a> PrimaryKey<'a> for Pk1<'a> { - type Prefix = Pk0; +impl<'a> PrimaryKey<'a> for &'a [u8] { + type Prefix = (); fn key<'b>(&'b self) -> Vec<&'b [u8]> { // this is simple, we don't add more prefixes @@ -44,8 +36,8 @@ impl<'a> PrimaryKey<'a> for Pk1<'a> { } // Provide a string version of this to raw encode strings -impl<'a> PrimaryKey<'a> for PkStr<'a> { - type Prefix = Pk0; +impl<'a> PrimaryKey<'a> for &'a str { + type Prefix = (); fn key<'b>(&'b self) -> Vec<&'b [u8]> { // this is simple, we don't add more prefixes @@ -79,7 +71,7 @@ impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a>> PrimaryKey<'a> for impl<'a, T: PrimaryKey<'a> + Prefixer<'a>, U: PrimaryKey<'a> + Prefixer<'a>, V: PrimaryKey<'a>> PrimaryKey<'a> for (T, U, V) { - type Prefix = T; + type Prefix = (T, U); fn key(&self) -> Vec<&[u8]> { let mut keys = self.0.key(); @@ -108,32 +100,37 @@ pub trait Prefixer<'a> { fn prefix<'b>(&'b self) -> Vec<&'b [u8]>; } -impl<'a> Prefixer<'a> for Pk0 { +impl<'a> Prefixer<'a> for () { fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { vec![] } } -impl<'a> Prefixer<'a> for Pk1<'a> { +impl<'a> Prefixer<'a> for &'a [u8] { fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { vec![self] } } -impl<'a> Prefixer<'a> for Pk2<'a> { +impl<'a, T: Prefixer<'a>, U: Prefixer<'a>> Prefixer<'a> for (T, U) { fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { - vec![self.0, self.1] + let mut res = self.0.prefix(); + res.extend(self.1.prefix().into_iter()); + res } } -impl<'a> Prefixer<'a> for Pk3<'a> { - fn prefix(&self) -> Vec<&[u8]> { - vec![self.0, self.1, self.2] +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]> { + let mut res = self.0.prefix(); + res.extend(self.1.prefix().into_iter()); + res.extend(self.2.prefix().into_iter()); + res } } // Provide a string version of this to raw encode strings -impl<'a> Prefixer<'a> for PkStr<'a> { +impl<'a> Prefixer<'a> for &'a str { fn prefix<'b>(&'b self) -> Vec<&'b [u8]> { vec![self.as_bytes()] } @@ -279,19 +276,23 @@ mod test { #[test] fn str_key_works() { - let k: &str = "hello"; + type K<'a> = &'a str; + + let k: K = "hello"; let path = k.key(); assert_eq!(1, path.len()); assert_eq!("hello".as_bytes(), path[0]); let joined = k.joined_key(); - let parsed = PkStr::parse_key(&joined); + let parsed = K::parse_key(&joined); assert_eq!(parsed, "hello"); } #[test] fn nested_str_key_works() { - let k: (&str, &[u8]) = ("hello", b"world"); + type K<'a> = (&'a str, &'a [u8]); + + let k: K = ("hello", b"world"); let path = k.key(); assert_eq!(2, path.len()); assert_eq!("hello".as_bytes(), path[0]); @@ -337,28 +338,45 @@ mod test { #[test] fn parse_joined_keys_pk1() { - let key: Pk1 = b"four"; + type K<'a> = &'a [u8]; + + let key: K = b"four"; let joined = key.joined_key(); assert_eq!(key, joined.as_slice()); - let parsed = Pk1::parse_key(&joined); + let parsed = K::parse_key(&joined); assert_eq!(key, parsed); } #[test] fn parse_joined_keys_pk2() { - let key: Pk2 = (b"four", b"square"); + type K<'a> = (&'a [u8], &'a [u8]); + + let key: K = (b"four", b"square"); let joined = key.joined_key(); assert_eq!(4 + 6 + 2, joined.len()); - let parsed = Pk2::parse_key(&joined); + let parsed = K::parse_key(&joined); assert_eq!(key, parsed); } #[test] fn parse_joined_keys_pk3() { - let key: Pk3 = (b"four", b"square", b"cinco"); + type K<'a> = (&'a str, U32Key, &'a [u8]); + + let key: K = ("four", 15.into(), b"cinco"); let joined = key.joined_key(); - assert_eq!(4 + 6 + 5 + 2 * (3 - 1), joined.len()); - let parsed = Pk3::parse_key(&joined); + assert_eq!(4 + 4 + 5 + 2 * 2, joined.len()); + let parsed = K::parse_key(&joined); + assert_eq!(key, parsed); + } + + #[test] + fn parse_joined_keys_pk3_alt() { + type K<'a> = (&'a str, U64Key, &'a str); + + let key: K = ("one", 222.into(), "three"); + let joined = key.joined_key(); + assert_eq!(3 + 8 + 5 + 2 * 2, joined.len()); + let parsed = K::parse_key(&joined); assert_eq!(key, parsed); } @@ -373,11 +391,33 @@ mod test { #[test] fn parse_joined_keys_string_int() { - let key: (U32Key, &str) = (54321.into(), "random"); + type K<'a> = (U32Key, &'a str); + + let key: K = (54321.into(), "random"); let joined = key.joined_key(); assert_eq!(2 + 4 + 6, joined.len()); - let parsed = <(U32Key, &str)>::parse_key(&joined); + let parsed = K::parse_key(&joined); assert_eq!(key, parsed); assert_eq!("random", parsed.1); } + + #[test] + fn proper_prefixes() { + let simple: &str = "hello"; + assert_eq!(simple.prefix(), vec![b"hello"]); + + let pair: (U32Key, &[u8]) = (12345.into(), b"random"); + let one: Vec = vec![0, 0, 48, 57]; + let two: Vec = b"random".to_vec(); + assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]); + + let triple: (&str, U32Key, &[u8]) = ("begin", 12345.into(), b"end"); + let one: Vec = b"begin".to_vec(); + let two: Vec = vec![0, 0, 48, 57]; + let three: Vec = b"end".to_vec(); + assert_eq!( + triple.prefix(), + vec![one.as_slice(), two.as_slice(), three.as_slice()] + ); + } } diff --git a/packages/storage-plus/src/map.rs b/packages/storage-plus/src/map.rs index 133618002..6bfe03294 100644 --- a/packages/storage-plus/src/map.rs +++ b/packages/storage-plus/src/map.rs @@ -104,6 +104,7 @@ mod test { use serde::{Deserialize, Serialize}; use std::ops::Deref; + use crate::U8Key; use cosmwasm_std::testing::MockStorage; #[cfg(feature = "iterator")] use cosmwasm_std::{Order, StdResult}; @@ -118,6 +119,8 @@ mod test { const ALLOWANCE: Map<(&[u8], &[u8]), u64> = Map::new("allow"); + const TRIPLE: Map<(&[u8], U8Key, &str), u64> = Map::new("triple"); + #[test] fn create_path() { let path = PEOPLE.key(b"john"); @@ -130,10 +133,25 @@ mod test { let path = ALLOWANCE.key((b"john", b"maria")); let key = path.deref(); // this should be prefixed(allow) || prefixed(john) || maria - assert_eq!("allow".len() + "john".len() + "maria".len() + 4, key.len()); + assert_eq!( + "allow".len() + "john".len() + "maria".len() + 2 * 2, + key.len() + ); assert_eq!(b"allow".to_vec().as_slice(), &key[2..7]); assert_eq!(b"john".to_vec().as_slice(), &key[9..13]); assert_eq!(b"maria".to_vec().as_slice(), &key[13..]); + + let path = TRIPLE.key((b"john", 8u8.into(), "pedro")); + let key = path.deref(); + // this should be prefixed(allow) || prefixed(john) || maria + assert_eq!( + "triple".len() + "john".len() + 1 + "pedro".len() + 2 * 3, + key.len() + ); + assert_eq!(b"triple".to_vec().as_slice(), &key[2..8]); + assert_eq!(b"john".to_vec().as_slice(), &key[10..14]); + assert_eq!(8u8.to_be_bytes(), &key[16..17]); + assert_eq!(b"pedro".to_vec().as_slice(), &key[17..]); } #[test] @@ -180,6 +198,29 @@ mod test { assert_eq!(1234, same); } + #[test] + fn triple_keys() { + let mut store = MockStorage::new(); + + // save and load on a triple composite key + let triple = TRIPLE.key((b"owner", 10u8.into(), "recipient")); + assert_eq!(None, triple.may_load(&store).unwrap()); + triple.save(&mut store, &1234).unwrap(); + assert_eq!(1234, triple.load(&store).unwrap()); + + // not under other key + let different = TRIPLE + .may_load(&store, (b"owners", 10u8.into(), "ecipient")) + .unwrap(); + assert_eq!(None, different); + + // matches under a proper copy + let same = TRIPLE + .load(&store, (b"owner", 10u8.into(), "recipient")) + .unwrap(); + assert_eq!(1234, same); + } + #[test] #[cfg(feature = "iterator")] fn range_simple_key() { @@ -237,6 +278,38 @@ mod test { ); } + #[test] + #[cfg(feature = "iterator")] + fn range_triple_key() { + let mut store = MockStorage::new(); + + // save and load on three keys, one under different owner + TRIPLE + .save(&mut store, (b"owner", 9u8.into(), "recipient"), &1000) + .unwrap(); + TRIPLE + .save(&mut store, (b"owner", 9u8.into(), "recipient2"), &3000) + .unwrap(); + TRIPLE + .save(&mut store, (b"owner2", 9u8.into(), "recipient"), &5000) + .unwrap(); + + // let's try to iterate! + let all: StdResult> = TRIPLE + .prefix((b"owner", 9u8.into())) + .range(&store, None, None, Order::Ascending) + .collect(); + let all = all.unwrap(); + assert_eq!(2, all.len()); + assert_eq!( + all, + vec![ + (b"recipient".to_vec(), 1000), + (b"recipient2".to_vec(), 3000) + ] + ); + } + #[test] fn basic_update() { let mut store = MockStorage::new();