Skip to content

Commit

Permalink
some improvements to bounded vec API (paritytech#10590)
Browse files Browse the repository at this point in the history 
* some improvements to bounded vec

* revert license tweak

* more tests

* fix

* Update frame/support/src/storage/bounded_vec.rs

Co-authored-by: Bastian Köcher <[email protected]>

* add the same stuff for btree map and set as well

Co-authored-by: Bastian Köcher <[email protected]>
  • Loading branch information
@kianenigma @bkchr
kianenigma and bkchr authored Jan 6, 2022
1 parent d642f5e commit 55f7969
Show file tree
Hide file tree
Showing 5 changed files with 301 additions and 32 deletions.
93 changes: 88 additions & 5 deletions frame/support/src/storage/bounded_btree_map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,10 @@

//! Traits, types and structs to support a bounded BTreeMap.
use crate::{storage::StorageDecodeLength, traits::Get};
use crate::{
storage::StorageDecodeLength,
traits::{Get, TryCollect},
};
use codec::{Decode, Encode, MaxEncodedLen};
use sp_std::{
borrow::Borrow, collections::btree_map::BTreeMap, convert::TryFrom, marker::PhantomData,
Expand Down Expand Up @@ -69,6 +72,11 @@ where
K: Ord,
S: Get<u32>,
{
/// Create `Self` from `t` without any checks.
fn unchecked_from(t: BTreeMap<K, V>) -> Self {
Self(t, Default::default())
}

/// Create a new `BoundedBTreeMap`.
///
/// Does not allocate.
Expand Down Expand Up @@ -183,16 +191,23 @@ where
}
}

impl<K, V, S> PartialEq for BoundedBTreeMap<K, V, S>
impl<K, V, S1, S2> PartialEq<BoundedBTreeMap<K, V, S1>> for BoundedBTreeMap<K, V, S2>
where
BTreeMap<K, V>: PartialEq,
S1: Get<u32>,
S2: Get<u32>,
{
fn eq(&self, other: &Self) -> bool {
self.0 == other.0
fn eq(&self, other: &BoundedBTreeMap<K, V, S1>) -> bool {
S1::get() == S2::get() && self.0 == other.0
}
}

impl<K, V, S> Eq for BoundedBTreeMap<K, V, S> where BTreeMap<K, V>: Eq {}
impl<K, V, S> Eq for BoundedBTreeMap<K, V, S>
where
BTreeMap<K, V>: Eq,
S: Get<u32>,
{
}

impl<K, V, S> PartialEq<BTreeMap<K, V>> for BoundedBTreeMap<K, V, S>
where
Expand All @@ -206,6 +221,7 @@ where
impl<K, V, S> PartialOrd for BoundedBTreeMap<K, V, S>
where
BTreeMap<K, V>: PartialOrd,
S: Get<u32>,
{
fn partial_cmp(&self, other: &Self) -> Option<sp_std::cmp::Ordering> {
self.0.partial_cmp(&other.0)
Expand All @@ -215,6 +231,7 @@ where
impl<K, V, S> Ord for BoundedBTreeMap<K, V, S>
where
BTreeMap<K, V>: Ord,
S: Get<u32>,
{
fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering {
self.0.cmp(&other.0)
Expand Down Expand Up @@ -302,6 +319,23 @@ impl<K, V, S> codec::EncodeLike<BTreeMap<K, V>> for BoundedBTreeMap<K, V, S> whe
{
}

impl<I, K, V, Bound> TryCollect<BoundedBTreeMap<K, V, Bound>> for I
where
K: Ord,
I: ExactSizeIterator + Iterator<Item = (K, V)>,
Bound: Get<u32>,
{
type Error = &'static str;

fn try_collect(self) -> Result<BoundedBTreeMap<K, V, Bound>, Self::Error> {
if self.len() > Bound::get() as usize {
Err("iterator length too big")
} else {
Ok(BoundedBTreeMap::<K, V, Bound>::unchecked_from(self.collect::<BTreeMap<K, V>>()))
}
}
}

#[cfg(test)]
pub mod test {
use super::*;
Expand Down Expand Up @@ -452,4 +486,53 @@ pub mod test {
assert_eq!(zero_key.1, false);
assert_eq!(*zero_value, 6);
}

#[test]
fn can_be_collected() {
let b1 = boundedmap_from_keys::<u32, ConstU32<5>>(&[1, 2, 3, 4]);
let b2: BoundedBTreeMap<u32, (), ConstU32<5>> =
b1.iter().map(|(k, v)| (k + 1, *v)).try_collect().unwrap();
assert_eq!(b2.into_iter().map(|(k, _)| k).collect::<Vec<_>>(), vec![2, 3, 4, 5]);

// can also be collected into a collection of length 4.
let b2: BoundedBTreeMap<u32, (), ConstU32<4>> =
b1.iter().map(|(k, v)| (k + 1, *v)).try_collect().unwrap();
assert_eq!(b2.into_iter().map(|(k, _)| k).collect::<Vec<_>>(), vec![2, 3, 4, 5]);

// can be mutated further into iterators that are `ExactSizedIterator`.
let b2: BoundedBTreeMap<u32, (), ConstU32<5>> =
b1.iter().map(|(k, v)| (k + 1, *v)).rev().skip(2).try_collect().unwrap();
// note that the binary tree will re-sort this, so rev() is not really seen
assert_eq!(b2.into_iter().map(|(k, _)| k).collect::<Vec<_>>(), vec![2, 3]);

let b2: BoundedBTreeMap<u32, (), ConstU32<5>> =
b1.iter().map(|(k, v)| (k + 1, *v)).take(2).try_collect().unwrap();
assert_eq!(b2.into_iter().map(|(k, _)| k).collect::<Vec<_>>(), vec![2, 3]);

// but these worn't work
let b2: Result<BoundedBTreeMap<u32, (), ConstU32<3>>, _> =
b1.iter().map(|(k, v)| (k + 1, *v)).try_collect();
assert!(b2.is_err());

let b2: Result<BoundedBTreeMap<u32, (), ConstU32<1>>, _> =
b1.iter().map(|(k, v)| (k + 1, *v)).skip(2).try_collect();
assert!(b2.is_err());
}

#[test]
fn eq_works() {
// of same type
let b1 = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2]);
let b2 = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2]);
assert_eq!(b1, b2);

// of different type, but same value and bound.
crate::parameter_types! {
B1: u32 = 7;
B2: u32 = 7;
}
let b1 = boundedmap_from_keys::<u32, B1>(&[1, 2]);
let b2 = boundedmap_from_keys::<u32, B2>(&[1, 2]);
assert_eq!(b1, b2);
}
}
118 changes: 100 additions & 18 deletions frame/support/src/storage/bounded_btree_set.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,10 @@

//! Traits, types and structs to support a bounded `BTreeSet`.
use crate::{storage::StorageDecodeLength, traits::Get};
use crate::{
storage::StorageDecodeLength,
traits::{Get, TryCollect},
};
use codec::{Decode, Encode, MaxEncodedLen};
use sp_std::{
borrow::Borrow, collections::btree_set::BTreeSet, convert::TryFrom, marker::PhantomData,
Expand Down Expand Up @@ -68,6 +71,11 @@ where
T: Ord,
S: Get<u32>,
{
/// Create `Self` from `t` without any checks.
fn unchecked_from(t: BTreeSet<T>) -> Self {
Self(t, Default::default())
}

/// Create a new `BoundedBTreeSet`.
///
/// Does not allocate.
Expand Down Expand Up @@ -168,20 +176,28 @@ where
}
}

impl<T, S> PartialEq for BoundedBTreeSet<T, S>
impl<T, S1, S2> PartialEq<BoundedBTreeSet<T, S1>> for BoundedBTreeSet<T, S2>
where
BTreeSet<T>: PartialEq,
S1: Get<u32>,
S2: Get<u32>,
{
fn eq(&self, other: &Self) -> bool {
self.0 == other.0
fn eq(&self, other: &BoundedBTreeSet<T, S1>) -> bool {
S1::get() == S2::get() && self.0 == other.0
}
}

impl<T, S> Eq for BoundedBTreeSet<T, S> where BTreeSet<T>: Eq {}
impl<T, S> Eq for BoundedBTreeSet<T, S>
where
BTreeSet<T>: Eq,
S: Get<u32>,
{
}

impl<T, S> PartialEq<BTreeSet<T>> for BoundedBTreeSet<T, S>
where
BTreeSet<T>: PartialEq,
S: Get<u32>,
{
fn eq(&self, other: &BTreeSet<T>) -> bool {
self.0 == *other
Expand All @@ -191,6 +207,7 @@ where
impl<T, S> PartialOrd for BoundedBTreeSet<T, S>
where
BTreeSet<T>: PartialOrd,
S: Get<u32>,
{
fn partial_cmp(&self, other: &Self) -> Option<sp_std::cmp::Ordering> {
self.0.partial_cmp(&other.0)
Expand All @@ -200,6 +217,7 @@ where
impl<T, S> Ord for BoundedBTreeSet<T, S>
where
BTreeSet<T>: Ord,
S: Get<u32>,
{
fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering {
self.0.cmp(&other.0)
Expand Down Expand Up @@ -283,6 +301,23 @@ impl<T, S> StorageDecodeLength for BoundedBTreeSet<T, S> {}

impl<T, S> codec::EncodeLike<BTreeSet<T>> for BoundedBTreeSet<T, S> where BTreeSet<T>: Encode {}

impl<I, T, Bound> TryCollect<BoundedBTreeSet<T, Bound>> for I
where
T: Ord,
I: ExactSizeIterator + Iterator<Item = T>,
Bound: Get<u32>,
{
type Error = &'static str;

fn try_collect(self) -> Result<BoundedBTreeSet<T, Bound>, Self::Error> {
if self.len() > Bound::get() as usize {
Err("iterator length too big")
} else {
Ok(BoundedBTreeSet::<T, Bound>::unchecked_from(self.collect::<BTreeSet<T>>()))
}
}
}

#[cfg(test)]
pub mod test {
use super::*;
Expand All @@ -298,39 +333,39 @@ pub mod test {
FooDoubleMap => DoubleMap<(u32, Twox128), (u32, Twox128), BoundedBTreeSet<u32, ConstU32<7>>>
}

fn map_from_keys<T>(keys: &[T]) -> BTreeSet<T>
fn set_from_keys<T>(keys: &[T]) -> BTreeSet<T>
where
T: Ord + Copy,
{
keys.iter().copied().collect()
}

fn boundedmap_from_keys<T, S>(keys: &[T]) -> BoundedBTreeSet<T, S>
fn boundedset_from_keys<T, S>(keys: &[T]) -> BoundedBTreeSet<T, S>
where
T: Ord + Copy,
S: Get<u32>,
{
map_from_keys(keys).try_into().unwrap()
set_from_keys(keys).try_into().unwrap()
}

#[test]
fn decode_len_works() {
TestExternalities::default().execute_with(|| {
let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
Foo::put(bounded);
assert_eq!(Foo::decode_len().unwrap(), 3);
});

TestExternalities::default().execute_with(|| {
let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
FooMap::insert(1, bounded);
assert_eq!(FooMap::decode_len(1).unwrap(), 3);
assert!(FooMap::decode_len(0).is_none());
assert!(FooMap::decode_len(2).is_none());
});

TestExternalities::default().execute_with(|| {
let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
FooDoubleMap::insert(1, 1, bounded);
assert_eq!(FooDoubleMap::decode_len(1, 1).unwrap(), 3);
assert!(FooDoubleMap::decode_len(2, 1).is_none());
Expand All @@ -341,17 +376,17 @@ pub mod test {

#[test]
fn try_insert_works() {
let mut bounded = boundedmap_from_keys::<u32, ConstU32<4>>(&[1, 2, 3]);
let mut bounded = boundedset_from_keys::<u32, ConstU32<4>>(&[1, 2, 3]);
bounded.try_insert(0).unwrap();
assert_eq!(*bounded, map_from_keys(&[1, 0, 2, 3]));
assert_eq!(*bounded, set_from_keys(&[1, 0, 2, 3]));

assert!(bounded.try_insert(9).is_err());
assert_eq!(*bounded, map_from_keys(&[1, 0, 2, 3]));
assert_eq!(*bounded, set_from_keys(&[1, 0, 2, 3]));
}

#[test]
fn deref_coercion_works() {
let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
// these methods come from deref-ed vec.
assert_eq!(bounded.len(), 3);
assert!(bounded.iter().next().is_some());
Expand All @@ -360,7 +395,7 @@ pub mod test {

#[test]
fn try_mutate_works() {
let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3, 4, 5, 6]);
let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3, 4, 5, 6]);
let bounded = bounded
.try_mutate(|v| {
v.insert(7);
Expand All @@ -376,8 +411,8 @@ pub mod test {

#[test]
fn btree_map_eq_works() {
let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3, 4, 5, 6]);
assert_eq!(bounded, map_from_keys(&[1, 2, 3, 4, 5, 6]));
let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3, 4, 5, 6]);
assert_eq!(bounded, set_from_keys(&[1, 2, 3, 4, 5, 6]));
}

#[test]
Expand Down Expand Up @@ -433,4 +468,51 @@ pub mod test {
assert_eq!(zero_item.0, 0);
assert_eq!(zero_item.1, false);
}

#[test]
fn can_be_collected() {
let b1 = boundedset_from_keys::<u32, ConstU32<5>>(&[1, 2, 3, 4]);
let b2: BoundedBTreeSet<u32, ConstU32<5>> = b1.iter().map(|k| k + 1).try_collect().unwrap();
assert_eq!(b2.into_iter().collect::<Vec<_>>(), vec![2, 3, 4, 5]);

// can also be collected into a collection of length 4.
let b2: BoundedBTreeSet<u32, ConstU32<4>> = b1.iter().map(|k| k + 1).try_collect().unwrap();
assert_eq!(b2.into_iter().collect::<Vec<_>>(), vec![2, 3, 4, 5]);

// can be mutated further into iterators that are `ExactSizedIterator`.
let b2: BoundedBTreeSet<u32, ConstU32<5>> =
b1.iter().map(|k| k + 1).rev().skip(2).try_collect().unwrap();
// note that the binary tree will re-sort this, so rev() is not really seen
assert_eq!(b2.into_iter().collect::<Vec<_>>(), vec![2, 3]);

let b2: BoundedBTreeSet<u32, ConstU32<5>> =
b1.iter().map(|k| k + 1).take(2).try_collect().unwrap();
assert_eq!(b2.into_iter().collect::<Vec<_>>(), vec![2, 3]);

// but these worn't work
let b2: Result<BoundedBTreeSet<u32, ConstU32<3>>, _> =
b1.iter().map(|k| k + 1).try_collect();
assert!(b2.is_err());

let b2: Result<BoundedBTreeSet<u32, ConstU32<1>>, _> =
b1.iter().map(|k| k + 1).skip(2).try_collect();
assert!(b2.is_err());
}

#[test]
fn eq_works() {
// of same type
let b1 = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2]);
let b2 = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2]);
assert_eq!(b1, b2);

// of different type, but same value and bound.
crate::parameter_types! {
B1: u32 = 7;
B2: u32 = 7;
}
let b1 = boundedset_from_keys::<u32, B1>(&[1, 2]);
let b2 = boundedset_from_keys::<u32, B2>(&[1, 2]);
assert_eq!(b1, b2);
}
}
Loading

0 comments on commit 55f7969

Please sign in to comment.