From 916f1d2b20e29d49b6b0c1e1436829efd71d2dbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Zemanovi=C4=8D?= Date: Tue, 7 Mar 2023 09:25:32 +0000 Subject: [PATCH] core/storage_api: add LazySet --- .../storage_api/collections/lazy_set.rs | 301 ++++++++++++++++++ .../src/ledger/storage_api/collections/mod.rs | 2 + 2 files changed, 303 insertions(+) create mode 100644 core/src/ledger/storage_api/collections/lazy_set.rs diff --git a/core/src/ledger/storage_api/collections/lazy_set.rs b/core/src/ledger/storage_api/collections/lazy_set.rs new file mode 100644 index 00000000000..c7414f5c497 --- /dev/null +++ b/core/src/ledger/storage_api/collections/lazy_set.rs @@ -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 { + key: storage::Key, + phantom_k: PhantomData, +} + +/// Possible sub-keys of a [`LazySet`] +#[derive(Clone, Debug)] +pub enum SubKey { + /// 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 { + /// Insert a key `K` in a [`LazySet`]. + Insert(K), + /// Remove a key `K` from a [`LazyMap`]. + 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 = std::result::Result; + +impl LazyCollection for LazySet +where + K: storage::KeySeg + Debug, +{ + type Action = Action; + type SubKey = SubKey; + type SubKeyWithData = Action; + 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> { + 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, + storage_key: &storage::Key, + sub_key: Self::SubKey, + ) -> storage_api::Result> + 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, + ) -> storage_api::Result> { + Ok(keys) + } +} + +// `LazySet` methods +impl LazySet +where + K: storage::KeySeg, +{ + /// Returns whether the set contains a value. + pub fn contains(&self, storage: &S, key: &K) -> Result + 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(&self, storage: &mut S, key: K) -> Result + 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(&self, storage: &mut S, key: &K) -> Result + 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(&self, storage: &S) -> Result + 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(&self, storage: &S) -> Result + 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`, + /// 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> + '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: &ENV, + storage_key: &storage::Key, + parsed_key: K, +) -> storage_api::Result>> +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::::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(()) + } +} diff --git a/core/src/ledger/storage_api/collections/mod.rs b/core/src/ledger/storage_api/collections/mod.rs index 688b76bd499..3c824c35f4a 100644 --- a/core/src/ledger/storage_api/collections/mod.rs +++ b/core/src/ledger/storage_api/collections/mod.rs @@ -14,9 +14,11 @@ use derivative::Derivative; use thiserror::Error; pub mod lazy_map; +pub mod lazy_set; pub mod lazy_vec; pub use lazy_map::LazyMap; +pub use lazy_set::LazySet; pub use lazy_vec::LazyVec; use crate::ledger::storage_api;