Merge pull request #213 from CosmWasm/triple-primary-key-2

Triple primary key 2

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<u8> = vec![0, 0, 48, 57];
+ let two: Vec<u8> = 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<u8> = b"begin".to_vec();
+ let two: Vec<u8> = vec![0, 0, 48, 57];
+ let three: Vec<u8> = b"end".to_vec();
+ assert_eq!(
+ triple.prefix(),
+ vec![one.as_slice(), two.as_slice(), three.as_slice()]
+ );
+ }
 }

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<Vec<_>> = 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();