Skip to content

Commit

Permalink
Merge pull request #207 from CosmWasm/unique-index-generic-key
Browse files Browse the repository at this point in the history 
Generalize UniqueIndex keys
  • Loading branch information
@maurolacy
maurolacy authored Jan 4, 2021
2 parents 0e4be71 + fcc1d7c commit ffd92a0
Show file tree
Hide file tree
Showing 5 changed files with 272 additions and 108 deletions.
279 changes: 194 additions & 85 deletions packages/storage-plus/src/indexed_map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,26 +156,29 @@ where
mod test {
use super::*;

use crate::indexes::{index_int, index_string, MultiIndex, UniqueIndex};
use crate::indexes::{index_string, index_string_tuple, MultiIndex, UniqueIndex};
use crate::{PkOwned, U32Key};
use cosmwasm_std::testing::MockStorage;
use cosmwasm_std::{MemoryStorage, Order};
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
struct Data {
pub name: String,
pub age: i32,
pub last_name: String,
pub age: u32,
}

struct DataIndexes<'a> {
pub name: MultiIndex<'a, Data>,
pub age: UniqueIndex<'a, Data>,
pub age: UniqueIndex<'a, U32Key, Data>,
pub name_lastname: UniqueIndex<'a, (PkOwned, PkOwned), Data>,
}

// Future Note: this can likely be macro-derived
impl<'a> IndexList<Data> for DataIndexes<'a> {
fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<Data>> + '_> {
let v: Vec<&dyn Index<Data>> = vec![&self.name, &self.age];
let v: Vec<&dyn Index<Data>> = vec![&self.name, &self.age, &self.name_lastname];
Box::new(v.into_iter())
}
}
Expand All @@ -184,35 +187,87 @@ mod test {
fn build_map<'a>() -> IndexedMap<'a, &'a [u8], Data, DataIndexes<'a>> {
let indexes = DataIndexes {
name: MultiIndex::new(|d| index_string(&d.name), "data", "data__name"),
age: UniqueIndex::new(|d| index_int(d.age), "data__age"),
age: UniqueIndex::new(|d| U32Key::new(d.age), "data__age"),
name_lastname: UniqueIndex::new(
|d| index_string_tuple(&d.name, &d.last_name),
"data__name_lastname",
),
};
IndexedMap::new("data", indexes)
}

fn save_data<'a>(
store: &mut MockStorage,
map: &mut IndexedMap<'a, &'a [u8], Data, DataIndexes<'a>>,
) -> (Vec<&'a [u8]>, Vec<Data>) {
let mut pks = vec![];
let mut datas = vec![];
let data = Data {
name: "Maria".to_string(),
last_name: "Doe".to_string(),
age: 42,
};
let pk: &[u8] = b"1";
map.save(store, pk, &data).unwrap();
pks.push(pk);
datas.push(data);

// same name (multi-index), different last name, different age => ok
let data = Data {
name: "Maria".to_string(),
last_name: "Williams".to_string(),
age: 23,
};
let pk: &[u8] = b"2";
map.save(store, pk, &data).unwrap();
pks.push(pk);
datas.push(data);

// different name, different last name, different age => ok
let data = Data {
name: "John".to_string(),
last_name: "Wayne".to_string(),
age: 32,
};
let pk: &[u8] = b"3";
map.save(store, pk, &data).unwrap();
pks.push(pk);
datas.push(data);

let data = Data {
name: "Maria Luisa".to_string(),
last_name: "Rodriguez".to_string(),
age: 12,
};
let pk: &[u8] = b"4";
map.save(store, pk, &data).unwrap();
pks.push(pk);
datas.push(data);

(pks, datas)
}

#[test]
fn store_and_load_by_index() {
let mut store = MockStorage::new();
let mut map = build_map();

// save data
let data = Data {
name: "Maria".to_string(),
age: 42,
};
let pk: &[u8] = b"5627";
map.save(&mut store, pk, &data).unwrap();
let (pks, datas) = save_data(&mut store, &mut map);
let pk = pks[0];
let data = &datas[0];

// load it properly
let loaded = map.load(&store, pk).unwrap();
assert_eq!(data, loaded);
assert_eq!(*data, loaded);

let count = map
.idx
.name
.all_items(&store, &index_string("Maria"))
.unwrap()
.len();
assert_eq!(1, count);
assert_eq!(2, count);

// TODO: we load by wrong keys - get full storage key!

Expand All @@ -223,10 +278,10 @@ mod test {
.name
.all_items(&store, &index_string("Maria"))
.unwrap();
assert_eq!(1, marias.len());
assert_eq!(2, marias.len());
let (k, v) = &marias[0];
assert_eq!(pk, k.as_slice());
assert_eq!(&data, v);
assert_eq!(data, v);

// other index doesn't match (1 byte after)
let count = map
Expand Down Expand Up @@ -256,14 +311,14 @@ mod test {
assert_eq!(0, count);

// match on proper age
let proper = index_int(42);
let aged = map.idx.age.item(&store, &proper).unwrap().unwrap();
let proper = U32Key::new(42);
let aged = map.idx.age.item(&store, proper).unwrap().unwrap();
assert_eq!(pk.to_vec(), aged.0);
assert_eq!(data, aged.1);
assert_eq!(*data, aged.1);

// no match on wrong age
let too_old = index_int(43);
let aged = map.idx.age.item(&store, &too_old).unwrap();
let too_old = U32Key::new(43);
let aged = map.idx.age.item(&store, too_old).unwrap();
assert_eq!(None, aged);
}

Expand All @@ -272,54 +327,62 @@ mod test {
let mut store = MockStorage::new();
let mut map = build_map();

// first data
let data1 = Data {
name: "Maria".to_string(),
age: 42,
};
let pk1: &[u8] = b"5627";
map.save(&mut store, pk1, &data1).unwrap();

// same name (multi-index), different age => ok
let data2 = Data {
name: "Maria".to_string(),
age: 23,
};
let pk2: &[u8] = b"7326";
map.save(&mut store, pk2, &data2).unwrap();
// save data
let (pks, datas) = save_data(&mut store, &mut map);

// different name, same age => error
let data3 = Data {
// different name, different last name, same age => error
let data5 = Data {
name: "Marta".to_string(),
last_name: "Laurens".to_string(),
age: 42,
};
let pk3: &[u8] = b"8263";
let pk5: &[u8] = b"4";

// enforce this returns some error
map.save(&mut store, pk3, &data3).unwrap_err();
map.save(&mut store, pk5, &data5).unwrap_err();

// query by unique key
// match on proper age
let age42 = index_int(42);
let (k, v) = map.idx.age.item(&store, &age42).unwrap().unwrap();
assert_eq!(k.as_slice(), pk1);
assert_eq!(&v.name, "Maria");
assert_eq!(v.age, 42);
let age42 = U32Key::new(42);
let (k, v) = map.idx.age.item(&store, age42.clone()).unwrap().unwrap();
assert_eq!(k.as_slice(), pks[0]);
assert_eq!(v.name, datas[0].name);
assert_eq!(v.age, datas[0].age);

// match on other age
let age23 = index_int(23);
let (k, v) = map.idx.age.item(&store, &age23).unwrap().unwrap();
assert_eq!(k.as_slice(), pk2);
assert_eq!(&v.name, "Maria");
assert_eq!(v.age, 23);
let age23 = U32Key::new(23);
let (k, v) = map.idx.age.item(&store, age23).unwrap().unwrap();
assert_eq!(k.as_slice(), pks[1]);
assert_eq!(v.name, datas[1].name);
assert_eq!(v.age, datas[1].age);

// if we delete the first one, we can add the blocked one
map.remove(&mut store, pk1).unwrap();
map.save(&mut store, pk3, &data3).unwrap();
map.remove(&mut store, pks[0]).unwrap();
map.save(&mut store, pk5, &data5).unwrap();
// now 42 is the new owner
let (k, v) = map.idx.age.item(&store, &age42).unwrap().unwrap();
assert_eq!(k.as_slice(), pk3);
assert_eq!(&v.name, "Marta");
assert_eq!(v.age, 42);
let (k, v) = map.idx.age.item(&store, age42).unwrap().unwrap();
assert_eq!(k.as_slice(), pk5);
assert_eq!(v.name, data5.name);
assert_eq!(v.age, data5.age);
}

#[test]
fn unique_index_enforced_composite_key() {
let mut store = MockStorage::new();
let mut map = build_map();

// save data
save_data(&mut store, &mut map);

// same name, same lastname => error
let data5 = Data {
name: "Maria".to_string(),
last_name: "Doe".to_string(),
age: 24,
};
let pk5: &[u8] = b"5";
// enforce this returns some error
map.save(&mut store, pk5, &data5).unwrap_err();
}

#[test]
Expand All @@ -337,45 +400,91 @@ mod test {
.count()
};

// set up some data
let data1 = Data {
name: "John".to_string(),
age: 22,
};
let pk1: &[u8] = b"john";
map.save(&mut store, pk1, &data1).unwrap();
let data2 = Data {
name: "John".to_string(),
age: 25,
};
let pk2: &[u8] = b"john2";
map.save(&mut store, pk2, &data2).unwrap();
let data3 = Data {
name: "Fred".to_string(),
age: 33,
};
let pk3: &[u8] = b"fred";
map.save(&mut store, pk3, &data3).unwrap();
// save data
let (pks, _) = save_data(&mut store, &mut map);

// find 2 Johns, 1 Fred, and no Mary
assert_eq!(name_count(&map, &store, "John"), 2);
assert_eq!(name_count(&map, &store, "Fred"), 1);
// find 2 Marias, 1 John, and no Mary
assert_eq!(name_count(&map, &store, "Maria"), 2);
assert_eq!(name_count(&map, &store, "John"), 1);
assert_eq!(name_count(&map, &store, "Maria Luisa"), 1);
assert_eq!(name_count(&map, &store, "Mary"), 0);

// remove john 2
map.remove(&mut store, pk2).unwrap();
// change fred to mary
map.update(&mut store, pk3, |d| -> StdResult<_> {
// remove maria 2
map.remove(&mut store, pks[1]).unwrap();

// change john to mary
map.update(&mut store, pks[2], |d| -> StdResult<_> {
let mut x = d.unwrap();
assert_eq!(&x.name, "Fred");
assert_eq!(&x.name, "John");
x.name = "Mary".to_string();
Ok(x)
})
.unwrap();

// find 1 Johns, no Fred, and 1 Mary
assert_eq!(name_count(&map, &store, "John"), 1);
assert_eq!(name_count(&map, &store, "Fred"), 0);
// find 1 maria, 1 maria luisa, no john, and 1 mary
assert_eq!(name_count(&map, &store, "Maria"), 1);
assert_eq!(name_count(&map, &store, "Maria Luisa"), 1);
assert_eq!(name_count(&map, &store, "John"), 0);
assert_eq!(name_count(&map, &store, "Mary"), 1);
}

#[test]
fn unique_index_simple_key_range() {
let mut store = MockStorage::new();
let mut map = build_map();

// save data
let (pks, datas) = save_data(&mut store, &mut map);

let res: StdResult<Vec<_>> = map
.idx
.age
.range(&store, None, None, Order::Ascending)
.collect();
let ages = res.unwrap();

let count = ages.len();
assert_eq!(4, count);

// The pks, sorted by age ascending
assert_eq!(pks[3].to_vec(), ages[0].0);
assert_eq!(pks[1].to_vec(), ages[1].0);
assert_eq!(pks[2].to_vec(), ages[2].0);
assert_eq!(pks[0].to_vec(), ages[3].0);

// The associated data
assert_eq!(datas[3], ages[0].1);
assert_eq!(datas[1], ages[1].1);
assert_eq!(datas[2], ages[2].1);
assert_eq!(datas[0], ages[3].1);
}

#[test]
fn unique_index_composite_key_range() {
let mut store = MockStorage::new();
let mut map = build_map();

// save data
let (pks, datas) = save_data(&mut store, &mut map);

let res: StdResult<Vec<_>> = map
.idx
.name_lastname
.prefix(PkOwned(b"Maria".to_vec()))
.range(&store, None, None, Order::Ascending)
.collect();
let marias = res.unwrap();

// Only two people are called "Maria"
let count = marias.len();
assert_eq!(2, count);

// The pks
assert_eq!(pks[0].to_vec(), marias[0].0);
assert_eq!(pks[1].to_vec(), marias[1].0);

// The associated data
assert_eq!(datas[0], marias[0].1);
assert_eq!(datas[1], marias[1].1);
}
}
Loading

0 comments on commit ffd92a0

Please sign in to comment.