Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Triple primary key 2 #213

Merged
merged 6 commits into from
Jan 5, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 73 additions & 33 deletions packages/storage-plus/src/keys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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()]
}
Expand Down Expand Up @@ -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]);
Expand Down Expand Up @@ -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);
}

Expand All @@ -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()]
);
}
}
75 changes: 74 additions & 1 deletion packages/storage-plus/src/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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");
Expand All @@ -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,
maurolacy marked this conversation as resolved.
Show resolved Hide resolved
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]
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -237,6 +278,38 @@ mod test {
);
}

#[test]
#[cfg(feature = "iterator")]
fn range_triple_key() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good test. This is what I want to see for the full-stack usage.

Can you update the TRIPLE to (&[u8], Int32Key, Int64Key) and verify this works just as well?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I fixed the prefix implemention in my last commit.
You should be able to use it in the new code.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to reflect why we cannot satisfy Prefix automatically for anything that implements PrimaryKey - I assume this is a subset of the work.

Copy link
Contributor Author

@maurolacy maurolacy Jan 5, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to reflect why we cannot satisfy Prefix automatically for anything that implements PrimaryKey - I assume this is a subset of the work.

I guess the idea is to be able to range() over the last part of a given key; that is, the one that doesn't implement Prefixer.

That's OK IMO. What I don't like is that currently prefix() only works for the "full" prefix, i. e. in case of triple keys, we cannot prefix only on the first element of the triple.
What I'm now thinking is that maybe we can use (abuse?) notation like (T, (U, V)), for definining a key that is internally "flat", but can be prefixed over its first element.

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();
Expand Down