Implement Entry API for storage2::LazyHashMap (#480)

* [core] Rename Entry to Internal Entry

* [core] Add Entry API for storage2::LazyHashMap

* [core] Add storage2::LazyHashMap::len()

* [core] Migrate tests to use storage2::LazyHashMap::len()

* [core] Implement FromIterator and Extend for storage2::LazyHashMap

* [core] Implement macro to generate LazyHashMap + HashMap Entry API tests

* [core] Remove redundant storage2::HashMap Entry API tests

* [core] Make storage2::HashMap Entry API use storage2::LazyHashMap's Entry API

* [core] Move parameterized Entry API tests into separate file

* [core] Rename InternalEntry to StorageEntry

* [core] Make lazy_hmap module public

* [core] Generate Entry API benches for LazyHashMap and HashMap from macro

* [core] Minor streamlining

* [core] Display hashmap variant in benchmark description

* [core] Fix comment

* [core] Fix typos

* [core] Make more use of BTreeMap Entry API

* [core] Replace unwrap with expect

* [core] Improve comment

* [core] Handle loading from storage

* [core] Restrict unsafe

* [core] Less ops for case "entry not in cache, but in storage"

* [core] Rename len()

* [core] Fix typo

* [core] Fix visibility

* [core] Address comments

* [core] Add test to verify that cache is marked as 'Mutated'

* [core] Shorten code with utility function

* [core] Improve naming

core/benches/bench_hashmap.rs

@@ -22,118 +22,115 @@ use criterion::{
 
 use ink_core::{
     env,
-    storage2::{
-        collections::{
-            hashmap::Entry,
-            HashMap as StorageHashMap,
-        },
-        traits::{
+    storage2::traits::{
         KeyPtr,
         SpreadLayout,
     },
-    },
 };
 use ink_primitives::Key;
 
-criterion_group!(
+#[cfg(test)]
+macro_rules! gen_tests_for_backend {
+    ( $backend:ty ) => {
+        criterion_group!(
             populated_cache,
             bench_insert_populated_cache,
             bench_remove_populated_cache,
-);
-criterion_group!(
+        );
+        criterion_group!(
             empty_cache,
             bench_insert_empty_cache,
             bench_remove_empty_cache,
-);
-criterion_main!(populated_cache, empty_cache,);
+        );
+        criterion_main!(populated_cache, empty_cache,);
 
-/// The number of `ENTIRES` denotes how many test values are put into
-/// the `StorageHashMap` used in these benchmarks.
-const ENTRIES: i32 = 500;
+        /// The number of `ENTIRES` denotes how many test values are put into
+        /// the hashmap used in these benchmarks.
+        const ENTRIES: i32 = 500;
 
-/// Returns some test values for use in benchmarks.
-fn test_values() -> Vec<(i32, i32)> {
+        /// Returns some test values for use in benchmarks.
+        fn test_values() -> Vec<(i32, i32)> {
             (0..ENTRIES)
                 .into_iter()
                 .map(|index| (index, index))
                 .collect::<Vec<_>>()
-}
+        }
 
-/// Creates a `StorageHashMap` from the given slice.
-fn hashmap_from_slice(slice: &[(i32, i32)]) -> StorageHashMap<i32, i32> {
-    slice.iter().copied().collect::<StorageHashMap<i32, i32>>()
-}
+        /// Creates a hashmap from the given slice.
+        fn hashmap_from_slice(slice: &[(i32, i32)]) -> $backend {
+            slice.iter().copied().collect::<$backend>()
+        }
 
-/// Creates a `StorageHashMap` from `test_values()`.
-fn setup_hashmap() -> StorageHashMap<i32, i32> {
+        /// Creates a hashmap from `test_values()`.
+        fn setup_hashmap() -> $backend {
             let test_values = test_values();
             hashmap_from_slice(&test_values[..])
-}
+        }
 
-/// Returns always the same `KeyPtr`.
-fn key_ptr() -> KeyPtr {
+        /// Returns always the same `KeyPtr`.
+        fn key_ptr() -> KeyPtr {
             let root_key = Key::from([0x42; 32]);
             KeyPtr::from(root_key)
-}
+        }
 
-/// Creates a `StorageHashMap` and pushes it to the contract storage.
-fn push_storage_hashmap() {
+        /// Creates a hashmap and pushes it to the contract storage.
+        fn push_storage_hashmap() {
             let hmap = setup_hashmap();
             SpreadLayout::push_spread(&hmap, &mut key_ptr());
-}
+        }
 
-/// Pulls a lazily loading `StorageHashMap` instance from the contract storage.
-fn pull_storage_hashmap() -> StorageHashMap<i32, i32> {
-    <StorageHashMap<i32, i32> as SpreadLayout>::pull_spread(&mut key_ptr())
-}
+        /// Pulls a lazily loading hashmap instance from the contract storage.
+        fn pull_storage_hashmap() -> $backend {
+            <$backend as SpreadLayout>::pull_spread(&mut key_ptr())
+        }
 
-/// Iteratively checks if an entry is in the `StorageHashMap`. If not, it
-/// is inserted. In either case it is incremented afterwards.
-fn insert_and_inc(hmap: &mut StorageHashMap<i32, i32>) {
+        /// Iteratively checks if an entry is in the hashmap. If not, it
+        /// is inserted. In either case it is incremented afterwards.
+        fn insert_and_inc(hmap: &mut $backend) {
             for key in 0..ENTRIES * 2 {
-        if !black_box(hmap.contains_key(&key)) {
-            black_box(hmap.insert(key, key));
+                if !black_box(contains_key(hmap, &key)) {
+                    black_box(insert(hmap, key, key));
                 }
                 *black_box(hmap.get_mut(&key)).unwrap() += 1;
             }
-}
+        }
 
-/// Iteratively checks if an entry is in the `StorageHashMap`. If not, it
-/// is inserted. In either case it is incremented afterwards.
-///
-/// Uses the Entry API.
-fn insert_and_inc_entry_api(hmap: &mut StorageHashMap<i32, i32>) {
+        /// Iteratively checks if an entry is in the hashmap. If not, it
+        /// is inserted. In either case it is incremented afterwards.
+        ///
+        /// Uses the Entry API.
+        fn insert_and_inc_entry_api(hmap: &mut $backend) {
             for key in 0..ENTRIES * 2 {
                 let v = black_box(hmap.entry(key).or_insert(key));
                 *v += 1;
             }
-}
+        }
 
-/// Iteratively checks if an entry is in the `StorageHashMap`. If yes, it
-/// is taken out.
-fn remove(hmap: &mut StorageHashMap<i32, i32>) {
+        /// Iteratively checks if an entry is in the hashmap. If yes, it
+        /// is taken out.
+        fn remove(hmap: &mut $backend) {
             for key in 0..ENTRIES * 2 {
-        if black_box(hmap.contains_key(&key)) {
-            let _ = black_box(hmap.take(&key));
+                if black_box(contains_key(hmap, &key)) {
+                    let _ = black_box(take(hmap, &key));
+                }
             }
         }
-}
 
-/// Iteratively checks if an entry is in the `StorageHashMap`. If yes, it
-/// is taken out.
-///
-/// Uses the Entry API.
-fn remove_entry_api(hmap: &mut StorageHashMap<i32, i32>) {
+        /// Iteratively checks if an entry is in the hashmap. If yes, it
+        /// is taken out.
+        ///
+        /// Uses the Entry API.
+        fn remove_entry_api(hmap: &mut $backend) {
             for key in 0..ENTRIES * 2 {
                 if let Entry::Occupied(o) = black_box(hmap.entry(key)) {
                     o.remove();
                 }
             }
-}
+        }
 
-fn bench_insert_populated_cache(c: &mut Criterion) {
+        fn bench_insert_populated_cache(c: &mut Criterion) {
             let mut group = c.benchmark_group(
-        "Compare: `insert_and_inc` and `insert_and_inc_entry_api` (populated cache)",
+                format!("{} Compare: `insert_and_inc` and `insert_and_inc_entry_api` (populated cache)", stringify!($backend))
             );
             group.bench_function("insert_and_inc", |b| {
                 b.iter_batched_ref(
@@ -150,12 +147,12 @@ fn bench_insert_populated_cache(c: &mut Criterion) {
                 )
             });
             group.finish();
-}
+        }
 
-fn bench_remove_populated_cache(c: &mut Criterion) {
+        fn bench_remove_populated_cache(c: &mut Criterion) {
             let _ = env::test::run_test::<env::DefaultEnvTypes, _>(|_| {
                 let mut group = c.benchmark_group(
-            "Compare: `remove` and `remove_entry_api` (populated cache)",
+                    format!("{} Compare: `remove` and `remove_entry_api` (populated cache)", stringify!($backend))
                 );
                 group.bench_function("remove", |b| {
                     b.iter_batched_ref(
@@ -175,12 +172,12 @@ fn bench_remove_populated_cache(c: &mut Criterion) {
                 Ok(())
             })
             .unwrap();
-}
+        }
 
-fn bench_insert_empty_cache(c: &mut Criterion) {
+        fn bench_insert_empty_cache(c: &mut Criterion) {
             let _ = env::test::run_test::<env::DefaultEnvTypes, _>(|_| {
                 let mut group = c.benchmark_group(
-            "Compare: `insert_and_inc` and `insert_and_inc_entry_api` (empty cache)",
+                    format!("{} Compare: `insert_and_inc` and `insert_and_inc_entry_api` (empty cache)", stringify!($backend))
                 );
                 group.bench_function("insert_and_inc", |b| {
                     b.iter_batched_ref(
@@ -206,12 +203,12 @@ fn bench_insert_empty_cache(c: &mut Criterion) {
                 Ok(())
             })
             .unwrap();
-}
+        }
 
-fn bench_remove_empty_cache(c: &mut Criterion) {
+        fn bench_remove_empty_cache(c: &mut Criterion) {
             let _ = env::test::run_test::<env::DefaultEnvTypes, _>(|_| {
                 let mut group =
-            c.benchmark_group("Compare: `remove` and `remove_entry_api` (empty cache)");
+                    c.benchmark_group(format!("{} Compare: `remove` and `remove_entry_api` (empty cache)", stringify!($backend)));
                 group.bench_function("remove", |b| {
                     b.iter_batched_ref(
                         || {
@@ -236,4 +233,80 @@ fn bench_remove_empty_cache(c: &mut Criterion) {
                 Ok(())
             })
             .unwrap();
+        }
+    };
+}
+
+mod lazyhmap_backend {
+    use super::*;
+    use ink_core::{
+        hash::hasher::Blake2x256Hasher,
+        storage2::lazy::lazy_hmap::{
+            Entry,
+            LazyHashMap,
+        },
+    };
+
+    gen_tests_for_backend!(LazyHashMap<i32, i32, Blake2x256Hasher>);
+
+    pub fn insert(
+        hmap: &mut LazyHashMap<i32, i32, Blake2x256Hasher>,
+        key: i32,
+        value: i32,
+    ) -> Option<i32> {
+        hmap.put_get(&key, Some(value))
+    }
+
+    pub fn take(
+        hmap: &mut LazyHashMap<i32, i32, Blake2x256Hasher>,
+        key: &i32,
+    ) -> Option<i32> {
+        hmap.put_get(key, None)
+    }
+
+    pub fn contains_key(
+        hmap: &LazyHashMap<i32, i32, Blake2x256Hasher>,
+        key: &i32,
+    ) -> bool {
+        hmap.get(key).is_some()
+    }
+
+    pub fn run() {
+        self::main()
+    }
+}
+
+mod hashmap_backend {
+    use super::*;
+    use ink_core::storage2::collections::{
+        hashmap::Entry,
+        HashMap as StorageHashMap,
+    };
+
+    gen_tests_for_backend!(StorageHashMap<i32, i32>);
+
+    pub fn insert(
+        hmap: &mut StorageHashMap<i32, i32>,
+        key: i32,
+        value: i32,
+    ) -> Option<i32> {
+        hmap.insert(key, value)
+    }
+
+    pub fn take(hmap: &mut StorageHashMap<i32, i32>, key: &i32) -> Option<i32> {
+        hmap.take(key)
+    }
+
+    pub fn contains_key(hmap: &StorageHashMap<i32, i32>, key: &i32) -> bool {
+        hmap.contains_key(key)
+    }
+
+    pub fn run() {
+        self::main()
+    }
+}
+
+fn main() {
+    hashmap_backend::run();
+    lazyhmap_backend::run();
 }

core/src/storage2/collections/hashmap/mod.rs

@@ -38,7 +38,12 @@ use crate::{
     },
     storage2::{
         collections::Stash,
-        lazy::LazyHashMap,
+        lazy::lazy_hmap::{
+            Entry as LazyEntry,
+            LazyHashMap,
+            OccupiedEntry as LazyOccupiedEntry,
+            VacantEntry as LazyVacantEntry,
+        },
         traits::PackedLayout,
     },
 };
@@ -101,51 +106,43 @@ struct ValueEntry<V> {
     key_index: KeyIndex,
 }
 
-/// A vacant entry with previous and next vacant indices.
-pub struct OccupiedEntry<'a, K, V, H = Blake2x256Hasher>
+/// An occupied entry that holds the value.
+pub struct OccupiedEntry<'a, K, V>
 where
     K: Ord + Clone + PackedLayout,
     V: PackedLayout,
-    H: Hasher,
-    Key: From<<H as Hasher>::Output>,
 {
-    /// A reference to the used `HashMap` instance.
-    base: &'a mut HashMap<K, V, H>,
-    /// The index of the key associated with this value.
-    key_index: KeyIndex,
-    /// The key stored in this entry.
-    key: K,
+    /// A reference to the `Stash` instance, containing the keys.
+    keys: &'a mut Stash<K>,
+    /// The `LazyHashMap::OccupiedEntry`.
+    values_entry: LazyOccupiedEntry<'a, K, ValueEntry<V>>,
 }
 
 /// A vacant entry with previous and next vacant indices.
-pub struct VacantEntry<'a, K, V, H = Blake2x256Hasher>
+pub struct VacantEntry<'a, K, V>
 where
     K: Ord + Clone + PackedLayout,
     V: PackedLayout,
-    H: Hasher,
-    Key: From<<H as Hasher>::Output>,
 {
-    /// A reference to the used `HashMap` instance.
-    base: &'a mut HashMap<K, V, H>,
-    /// The key stored in this entry.
-    key: K,
+    /// A reference to the `Stash` instance, containing the keys.
+    keys: &'a mut Stash<K>,
+    /// The `LazyHashMap::VacantEntry`.
+    values_entry: LazyVacantEntry<'a, K, ValueEntry<V>>,
 }
 
 /// An entry within the stash.
 ///
 /// The vacant entries within a storage stash form a doubly linked list of
 /// vacant entries that is used to quickly re-use their vacant storage.
-pub enum Entry<'a, K: 'a, V: 'a, H = Blake2x256Hasher>
+pub enum Entry<'a, K: 'a, V: 'a>
 where
     K: Ord + Clone + PackedLayout,
     V: PackedLayout,
-    H: Hasher,
-    Key: From<<H as Hasher>::Output>,
 {
     /// A vacant entry that holds the index to the next and previous vacant entry.
-    Vacant(VacantEntry<'a, K, V, H>),
+    Vacant(VacantEntry<'a, K, V>),
     /// An occupied entry that holds the value.
-    Occupied(OccupiedEntry<'a, K, V, H>),
+    Occupied(OccupiedEntry<'a, K, V>),
 }
 
 impl<K, V, H> HashMap<K, V, H>
@@ -163,11 +160,17 @@ where
         }
     }
 
-    /// Returns the number of key- value pairs stored in the hash map.
+    /// Returns the number of key-value pairs stored in the hash map.
     pub fn len(&self) -> u32 {
         self.keys.len()
     }
 
+    /// Returns the number of key-value pairs stored in the cache.
+    #[cfg(test)]
+    pub(crate) fn len_cached_entries(&self) -> u32 {
+        self.keys.len()
+    }
+
     /// Returns `true` if the hash map is empty.
     pub fn is_empty(&self) -> bool {
         self.keys.is_empty()
@@ -386,33 +389,35 @@ where
     }
 
     /// Gets the given key's corresponding entry in the map for in-place manipulation.
-    pub fn entry(&mut self, key: K) -> Entry<K, V, H> {
-        let v = self.values.get(&key);
-        match v {
-            Some(entry) => {
+    pub fn entry(&mut self, key: K) -> Entry<K, V> {
+        let entry = self.values.entry(key);
+        match entry {
+            LazyEntry::Occupied(o) => {
                 Entry::Occupied(OccupiedEntry {
-                    key,
-                    key_index: entry.key_index,
-                    base: self,
+                    keys: &mut self.keys,
+                    values_entry: o,
+                })
+            }
+            LazyEntry::Vacant(v) => {
+                Entry::Vacant(VacantEntry {
+                    keys: &mut self.keys,
+                    values_entry: v,
                 })
             }
-            None => Entry::Vacant(VacantEntry { key, base: self }),
         }
     }
 }
 
-impl<'a, K, V, H> Entry<'a, K, V, H>
+impl<'a, K, V> Entry<'a, K, V>
 where
     K: Ord + Clone + PackedLayout,
     V: PackedLayout + core::fmt::Debug + core::cmp::Eq + Default,
-    H: Hasher,
-    Key: From<<H as Hasher>::Output>,
 {
     /// Returns a reference to this entry's key.
     pub fn key(&self) -> &K {
         match self {
-            Entry::Occupied(entry) => &entry.key,
-            Entry::Vacant(entry) => &entry.key,
+            Entry::Occupied(entry) => &entry.values_entry.key(),
+            Entry::Vacant(entry) => &entry.values_entry.key(),
         }
     }
 
@@ -420,7 +425,7 @@ where
     /// a reference to the value in the entry.
     pub fn or_default(self) -> &'a V {
         match self {
-            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Occupied(entry) => &mut entry.values_entry.into_mut().value,
             Entry::Vacant(entry) => entry.insert(V::default()),
         }
     }
@@ -429,7 +434,7 @@ where
     /// a mutable reference to the value in the entry.
     pub fn or_insert(self, default: V) -> &'a mut V {
         match self {
-            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Occupied(entry) => &mut entry.values_entry.into_mut().value,
             Entry::Vacant(entry) => entry.insert(default),
         }
     }
@@ -441,7 +446,7 @@ where
         F: FnOnce() -> V,
     {
         match self {
-            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Occupied(entry) => &mut entry.values_entry.into_mut().value,
             Entry::Vacant(entry) => Entry::insert(default(), entry),
         }
     }
@@ -454,8 +459,8 @@ where
         F: FnOnce(&K) -> V,
     {
         match self {
-            Entry::Occupied(entry) => entry.into_mut(),
-            Entry::Vacant(entry) => Entry::insert(default(&entry.key), entry),
+            Entry::Occupied(entry) => &mut entry.values_entry.into_mut().value,
+            Entry::Vacant(entry) => Entry::insert(default(&entry.key()), entry),
         }
     }
 
@@ -468,8 +473,8 @@ where
         match self {
             Entry::Occupied(mut entry) => {
                 {
-                    let v = entry.get_mut();
-                    f(v);
+                    let v = entry.values_entry.get_mut();
+                    f(&mut v.value);
                 }
                 Entry::Occupied(entry)
             }
@@ -478,78 +483,60 @@ where
     }
 
     /// Inserts `value` into `entry`.
-    fn insert(value: V, entry: VacantEntry<'a, K, V, H>) -> &'a mut V {
-        let old_value = entry.base.insert(entry.key.clone(), value);
-        debug_assert!(old_value.is_none());
-        entry
-            .base
-            .get_mut(&entry.key)
-            .expect("encountered invalid vacant entry")
+    fn insert(value: V, entry: VacantEntry<'a, K, V>) -> &'a mut V {
+        entry.insert(value)
     }
 }
 
-impl<'a, K, V, H> VacantEntry<'a, K, V, H>
+impl<'a, K, V> VacantEntry<'a, K, V>
 where
     K: Ord + Clone + PackedLayout,
     V: PackedLayout,
-    H: Hasher,
-    Key: From<<H as Hasher>::Output>,
 {
     /// Gets a reference to the key that would be used when inserting a value through the VacantEntry.
     pub fn key(&self) -> &K {
-        &self.key
+        &self.values_entry.key()
     }
 
     /// Take ownership of the key.
     pub fn into_key(self) -> K {
-        self.key
+        self.values_entry.into_key()
     }
 
-    /// Sets the value of the entry with the VacantEntry's key, and returns a mutable reference to it.
+    /// Sets the value of the entry with the `VacantEntry`'s key, and returns a mutable reference to it.
     pub fn insert(self, value: V) -> &'a mut V {
         // At this point we know that `key` does not yet exist in the map.
-        let key_index = self.base.keys.put(self.key.clone());
-        self.base
-            .values
-            .put(self.key.clone(), Some(ValueEntry { value, key_index }));
-        self.base
-            .get_mut(&self.key)
-            .expect("put was just executed; qed")
+        let key_index = self.keys.put(self.key().to_owned());
+        &mut self
+            .values_entry
+            .insert(ValueEntry { value, key_index })
+            .value
     }
 }
 
-impl<'a, K, V, H> OccupiedEntry<'a, K, V, H>
+impl<'a, K, V> OccupiedEntry<'a, K, V>
 where
     K: Ord + Clone + PackedLayout,
     V: PackedLayout,
-    H: Hasher,
-    Key: From<<H as Hasher>::Output>,
 {
     /// Gets a reference to the key in the entry.
     pub fn key(&self) -> &K {
-        &self.key
+        &self.values_entry.key()
     }
 
     /// Take the ownership of the key and value from the map.
     pub fn remove_entry(self) -> (K, V) {
-        let entry = self
-            .base
-            .values
-            .put_get(&self.key, None)
-            .expect("`key` must exist");
-        self.base
-            .keys
-            .take(self.key_index)
+        let k = self.values_entry.key().to_owned();
+        let v = self.values_entry.remove();
+        self.keys
+            .take(v.key_index)
             .expect("`key_index` must point to a valid key entry");
-        (self.key, entry.value)
+        (k, v.value)
     }
 
     /// Gets a reference to the value in the entry.
     pub fn get(&self) -> &V {
-        &self
-            .base
-            .get(&self.key)
-            .expect("entry behind `OccupiedEntry` must always exist")
+        &self.values_entry.get().value
     }
 
     /// Gets a mutable reference to the value in the entry.
@@ -557,19 +544,12 @@ where