Implement Entry API for storage2::LazyHashMap

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,41 +389,43 @@ 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(),
         }
     }
 
     /// Ensures a value is in the entry by inserting the default value if empty, and returns
     /// 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,98 +483,73 @@ 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.
     ///
     /// If you need a reference to the `OccupiedEntry` which may outlive the destruction of the
     /// `Entry` value, see `into_mut`.
     pub fn get_mut(&mut self) -> &mut V {
-        self.base
-            .get_mut(&self.key)
-            .expect("entry behind `OccupiedEntry` must always exist")
+        &mut self.values_entry.get_mut().value
     }
 
     /// Sets the value of the entry, and returns the entry's old value.
     pub fn insert(&mut self, new_value: V) -> V {
-        let occupied = self
-            .base
-            .values
-            .get_mut(&self.key)
-            .expect("entry behind `OccupiedEntry` must always exist");
-        core::mem::replace(&mut occupied.value, new_value)
+        core::mem::replace(&mut self.values_entry.get_mut().value, new_value)
     }
 
     /// Takes the value out of the entry, and returns it.
@@ -580,8 +560,6 @@ where
     /// Converts the OccupiedEntry into a mutable reference to the value in the entry
     /// with a lifetime bound to the map itself.
     pub fn into_mut(self) -> &'a mut V {
-        self.base
-            .get_mut(&self.key)
-            .expect("entry behind `OccupiedEntry` must always exist")
+        &mut self.values_entry.into_mut().value
     }
 }