-
Notifications
You must be signed in to change notification settings - Fork 951
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
07b8ff2
commit 916f1d2
Showing
2 changed files
with
303 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,301 @@ | ||
//! Lazy set. | ||
|
||
use std::fmt::Debug; | ||
use std::marker::PhantomData; | ||
|
||
use thiserror::Error; | ||
|
||
use super::super::Result; | ||
use super::{LazyCollection, ReadError}; | ||
use crate::ledger::storage_api::{self, ResultExt, StorageRead, StorageWrite}; | ||
use crate::ledger::vp_env::VpEnv; | ||
use crate::types::storage::{self, DbKeySeg, KeySeg}; | ||
|
||
/// A lazy set, implemented as a `LazyMap` with unit for a value | ||
#[derive(Debug)] | ||
pub struct LazySet<K> { | ||
key: storage::Key, | ||
phantom_k: PhantomData<K>, | ||
} | ||
|
||
/// Possible sub-keys of a [`LazySet`] | ||
#[derive(Clone, Debug)] | ||
pub enum SubKey<K> { | ||
/// Literal set key | ||
Data(K), | ||
} | ||
|
||
/// Possible actions that can modify a [`LazySet`]. This roughly corresponds to | ||
/// the methods that have `StorageWrite` access. | ||
#[derive(Clone, Debug)] | ||
pub enum Action<K> { | ||
/// Insert a key `K` in a [`LazySet<K>`]. | ||
Insert(K), | ||
/// Remove a key `K` from a [`LazyMap<K>`]. | ||
Remove(K), | ||
} | ||
|
||
#[allow(missing_docs)] | ||
#[derive(Error, Debug)] | ||
pub enum ValidationError { | ||
#[error("Invalid storage key {0}")] | ||
InvalidSubKey(storage::Key), | ||
} | ||
|
||
/// [`LazySet`] validation result | ||
pub type ValidationResult<T> = std::result::Result<T, ValidationError>; | ||
|
||
impl<K> LazyCollection for LazySet<K> | ||
where | ||
K: storage::KeySeg + Debug, | ||
{ | ||
type Action = Action<K>; | ||
type SubKey = SubKey<K>; | ||
type SubKeyWithData = Action<K>; | ||
type Value = (); | ||
|
||
/// Create or use an existing map with the given storage `key`. | ||
fn open(key: storage::Key) -> Self { | ||
Self { | ||
key, | ||
phantom_k: PhantomData, | ||
} | ||
} | ||
|
||
fn is_valid_sub_key( | ||
&self, | ||
key: &storage::Key, | ||
) -> storage_api::Result<Option<Self::SubKey>> { | ||
let suffix = match key.split_prefix(&self.key) { | ||
None => { | ||
// not matching prefix, irrelevant | ||
return Ok(None); | ||
} | ||
Some(None) => { | ||
// no suffix, invalid | ||
return Err(ValidationError::InvalidSubKey(key.clone())) | ||
.into_storage_result(); | ||
} | ||
Some(Some(suffix)) => suffix, | ||
}; | ||
|
||
// Match the suffix against expected sub-keys | ||
match &suffix.segments[..] { | ||
[DbKeySeg::StringSeg(sub)] => { | ||
if let Ok(key) = storage::KeySeg::parse(sub.clone()) { | ||
Ok(Some(SubKey::Data(key))) | ||
} else { | ||
Err(ValidationError::InvalidSubKey(key.clone())) | ||
.into_storage_result() | ||
} | ||
} | ||
_ => Err(ValidationError::InvalidSubKey(key.clone())) | ||
.into_storage_result(), | ||
} | ||
} | ||
|
||
fn read_sub_key_data<ENV>( | ||
env: &ENV, | ||
storage_key: &storage::Key, | ||
sub_key: Self::SubKey, | ||
) -> storage_api::Result<Option<Self::SubKeyWithData>> | ||
where | ||
ENV: for<'a> VpEnv<'a>, | ||
{ | ||
let SubKey::Data(key) = sub_key; | ||
determine_action(env, storage_key, key) | ||
} | ||
|
||
fn validate_changed_sub_keys( | ||
keys: Vec<Self::SubKeyWithData>, | ||
) -> storage_api::Result<Vec<Self::Action>> { | ||
Ok(keys) | ||
} | ||
} | ||
|
||
// `LazySet` methods | ||
impl<K> LazySet<K> | ||
where | ||
K: storage::KeySeg, | ||
{ | ||
/// Returns whether the set contains a value. | ||
pub fn contains<S>(&self, storage: &S, key: &K) -> Result<bool> | ||
where | ||
S: StorageRead, | ||
{ | ||
storage.has_key(&self.get_key(key)) | ||
} | ||
|
||
/// Get the storage sub-key of a given raw key | ||
pub fn get_key(&self, key: &K) -> storage::Key { | ||
let key_str = key.to_db_key(); | ||
self.key.push(&key_str).unwrap() | ||
} | ||
|
||
/// Inserts a key into the set. | ||
/// | ||
/// If the set did not have this key present, `false` is returned. | ||
/// If the set did have this key present, `true`. | ||
/// value is returned. Unlike in `std::collection::HashSet`, the key is also | ||
/// updated; this matters for types that can be `==` without being | ||
/// identical. | ||
pub fn insert<S>(&self, storage: &mut S, key: K) -> Result<bool> | ||
where | ||
S: StorageWrite + StorageRead, | ||
{ | ||
let present = self.contains(storage, &key)?; | ||
|
||
let key = self.get_key(&key); | ||
storage.write(&key, ())?; | ||
|
||
Ok(present) | ||
} | ||
|
||
/// Removes a key from the set, returning `true` if the key | ||
/// was in the set. | ||
pub fn remove<S>(&self, storage: &mut S, key: &K) -> Result<bool> | ||
where | ||
S: StorageWrite + StorageRead, | ||
{ | ||
let present = self.contains(storage, key)?; | ||
|
||
let key = self.get_key(key); | ||
storage.delete(&key)?; | ||
|
||
Ok(present) | ||
} | ||
|
||
/// Returns whether the set contains no elements. | ||
pub fn is_empty<S>(&self, storage: &S) -> Result<bool> | ||
where | ||
S: StorageRead, | ||
{ | ||
let mut iter = storage_api::iter_prefix_bytes(storage, &self.key)?; | ||
Ok(iter.next().is_none()) | ||
} | ||
|
||
/// Reads the number of elements in the map. | ||
/// | ||
/// Note that this function shouldn't be used in transactions and VPs code | ||
/// on unbounded maps to avoid gas usage increasing with the length of the | ||
/// set. | ||
#[allow(clippy::len_without_is_empty)] | ||
pub fn len<S>(&self, storage: &S) -> Result<u64> | ||
where | ||
S: StorageRead, | ||
{ | ||
let iter = storage_api::iter_prefix_bytes(storage, &self.key)?; | ||
iter.count().try_into().into_storage_result() | ||
} | ||
|
||
/// An iterator visiting all keas. The iterator element type is `Result<K>`, | ||
/// because iterator's call to `next` may fail with e.g. out of gas. | ||
/// | ||
/// Note that this function shouldn't be used in transactions and VPs code | ||
/// on unbounded sets to avoid gas usage increasing with the length of the | ||
/// set. | ||
pub fn iter<'iter>( | ||
&self, | ||
storage: &'iter impl StorageRead, | ||
) -> Result<impl Iterator<Item = Result<K>> + 'iter> { | ||
let iter = storage_api::iter_prefix(storage, &self.key)?; | ||
Ok(iter.map(|key_val_res| { | ||
let (key, ()) = key_val_res?; | ||
let last_key_seg = key | ||
.last() | ||
.ok_or(ReadError::UnexpectedlyEmptyStorageKey) | ||
.into_storage_result()?; | ||
let key = K::parse(last_key_seg.raw()).into_storage_result()?; | ||
Ok(key) | ||
})) | ||
} | ||
} | ||
|
||
/// Determine what action was taken from the pre/post state | ||
pub fn determine_action<ENV, K>( | ||
env: &ENV, | ||
storage_key: &storage::Key, | ||
parsed_key: K, | ||
) -> storage_api::Result<Option<Action<K>>> | ||
where | ||
ENV: for<'a> VpEnv<'a>, | ||
{ | ||
let pre = env.read_pre(storage_key)?; | ||
let post = env.read_post(storage_key)?; | ||
Ok(match (pre, post) { | ||
(None, None) => { | ||
// If the key was inserted and then deleted in the same tx, we don't | ||
// need to validate it as it's not visible to any VPs | ||
None | ||
} | ||
(None, Some(())) => Some(Action::Insert(parsed_key)), | ||
(Some(()), None) => Some(Action::Remove(parsed_key)), | ||
(Some(()), Some(())) => { | ||
// Because the value for set is a unit, we can skip this too | ||
None | ||
} | ||
}) | ||
} | ||
|
||
#[cfg(test)] | ||
mod test { | ||
use super::*; | ||
use crate::ledger::storage::testing::TestWlStorage; | ||
|
||
#[test] | ||
fn test_lazy_set_basics() -> storage_api::Result<()> { | ||
let mut storage = TestWlStorage::default(); | ||
|
||
let key = storage::Key::parse("test").unwrap(); | ||
let lazy_set = LazySet::<u32>::open(key); | ||
|
||
// The map should be empty at first | ||
assert!(lazy_set.is_empty(&storage)?); | ||
assert!(lazy_set.len(&storage)? == 0); | ||
assert!(!lazy_set.contains(&storage, &0)?); | ||
assert!(!lazy_set.contains(&storage, &1)?); | ||
assert!(lazy_set.iter(&storage)?.next().is_none()); | ||
assert!(!lazy_set.remove(&mut storage, &0)?); | ||
assert!(!lazy_set.remove(&mut storage, &1)?); | ||
|
||
// Insert a new value and check that it's added | ||
let key = 123; | ||
lazy_set.insert(&mut storage, key)?; | ||
|
||
let key2 = 456; | ||
lazy_set.insert(&mut storage, key2)?; | ||
|
||
assert!(!lazy_set.contains(&storage, &0)?); | ||
assert!(lazy_set.contains(&storage, &key)?); | ||
assert!(!lazy_set.is_empty(&storage)?); | ||
assert!(lazy_set.len(&storage)? == 2); | ||
let mut set_it = lazy_set.iter(&storage)?; | ||
assert_eq!(set_it.next().unwrap()?, key); | ||
assert_eq!(set_it.next().unwrap()?, key2); | ||
drop(set_it); | ||
|
||
assert!(!lazy_set.contains(&storage, &0)?); | ||
assert!(lazy_set.contains(&storage, &key)?); | ||
assert!(lazy_set.contains(&storage, &key2)?); | ||
|
||
// Remove the values and check the map contents | ||
let removed = lazy_set.remove(&mut storage, &key)?; | ||
assert!(removed); | ||
assert!(!lazy_set.is_empty(&storage)?); | ||
assert!(lazy_set.len(&storage)? == 1); | ||
assert!(!lazy_set.contains(&storage, &0)?); | ||
assert!(!lazy_set.contains(&storage, &1)?); | ||
assert!(!lazy_set.contains(&storage, &123)?); | ||
assert!(lazy_set.contains(&storage, &456)?); | ||
assert!(!lazy_set.contains(&storage, &key)?); | ||
assert!(lazy_set.contains(&storage, &key2)?); | ||
assert!(lazy_set.iter(&storage)?.next().is_some()); | ||
assert!(!lazy_set.remove(&mut storage, &key)?); | ||
let removed = lazy_set.remove(&mut storage, &key2)?; | ||
assert!(removed); | ||
assert!(lazy_set.is_empty(&storage)?); | ||
assert!(lazy_set.len(&storage)? == 0); | ||
|
||
Ok(()) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters