Skip to content

Commit

Permalink
WIP generic lazy map into eager map
Browse files Browse the repository at this point in the history
  • Loading branch information
brentstone committed Jul 18, 2023
1 parent 7466b8a commit b400eb1
Showing 1 changed file with 54 additions and 1 deletion.
55 changes: 54 additions & 1 deletion core/src/ledger/storage_api/collections/lazy_map.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
//! Lazy map.

use std::collections::HashMap;
use std::collections::{BTreeMap, HashMap};
use std::fmt::Debug;
use std::hash::Hash;
use std::marker::PhantomData;
Expand Down Expand Up @@ -101,6 +101,31 @@ pub enum ValidationError {
InvalidNestedSubKey(storage::Key),
}

pub trait EagerMapFromIter<K: Eq + Hash + Ord, V> {
fn from_iter<I>(iter: I) -> Self
where
I: IntoIterator<Item = (K, V)>;
}

impl<K: Eq + Hash + Ord, V> EagerMapFromIter<K, V> for HashMap<K, V> {
fn from_iter<I>(iter: I) -> Self
where
I: IntoIterator<Item = (K, V)>,
{
iter.into_iter().collect()
}
}

impl<K: Eq + Hash + Ord, V> EagerMapFromIter<K, V> for BTreeMap<K, V> {
fn from_iter<I>(iter: I) -> Self
where
K: Eq + Hash + Ord,
I: IntoIterator<Item = (K, V)>,
{
iter.into_iter().collect()
}
}

/// [`LazyMap`] validation result
pub type ValidationResult<T> = std::result::Result<T, ValidationError>;

Expand Down Expand Up @@ -563,6 +588,19 @@ where
}))
}

/// Collect the lazy map into an eager map
pub fn collect<M, S>(&self, storage: &S) -> Result<M>
where
S: StorageRead,
M: EagerMapFromIter<K, V>,
K: Eq + Hash + Ord,
{
let it = self
.iter(storage)?
.map(|res| res.expect("Failed to unwrap a lazy map element"));
Ok(M::from_iter(it))
}

/// Reads a value from storage
fn read_key_val<S>(
storage: &S,
Expand Down Expand Up @@ -629,6 +667,21 @@ mod test {
assert_eq!(lazy_map.get(&storage, &key)?.unwrap(), val);
assert_eq!(lazy_map.get(&storage, &key2)?.unwrap(), val2);

let eager_map: HashMap<_, _> = lazy_map.collect(&storage)?;
assert_eq!(
eager_map,
vec![(123, "Test".to_string()), (456, "Test2".to_string())]
.into_iter()
.collect::<HashMap<_, _>>()
);
let eager_map: BTreeMap<_, _> = lazy_map.collect(&storage)?;
assert_eq!(
eager_map,
vec![(123, "Test".to_string()), (456, "Test2".to_string())]
.into_iter()
.collect::<BTreeMap<_, _>>()
);

// Remove the values and check the map contents
let removed = lazy_map.remove(&mut storage, &key)?.unwrap();
assert_eq!(removed, val);
Expand Down

0 comments on commit b400eb1

Please sign in to comment.