fix entry block-id/hash cache refactor

Author: kantai

stackslib/src/chainstate/stacks/index/cache.rs

@@ -138,9 +138,23 @@ impl<T: MarfTrieId> TrieCacheState<T> {
         self.block_hash_cache.get(&block_id).cloned()
     }
 
-    /// Return the HashMap Entry for a block_id
-    pub fn entry_block_hash(&mut self, block_id: u32) -> Entry<'_, u32, T> {
-        self.block_hash_cache.entry(block_id)
+    /// Get cached entry for a block hash, given its ID, or, if not
+    ///  found, use `lookup` to get the corresponding block hash and
+    ///  store it in the cache
+    pub fn get_block_hash_caching<E, F: FnOnce(u32) -> Result<T, E>>(
+        &mut self,
+        id: u32,
+        lookup: F,
+    ) -> Result<&T, E> {
+        match self.block_hash_cache.entry(id) {
+            Entry::Occupied(occupied_entry) => Ok(occupied_entry.into_mut()),
+            Entry::Vacant(vacant_entry) => {
+                let block_hash = lookup(id)?;
+                let block_hash_ref = vacant_entry.insert(block_hash.clone());
+                self.block_id_cache.insert(block_hash, id);
+                Ok(block_hash_ref)
+            }
+        }
     }
 
     /// Cache a block hash, given its ID
@@ -315,9 +329,15 @@ impl<T: MarfTrieId> TrieCache<T> {
         self.state_mut().load_block_hash(block_id)
     }
 
-    /// Get entry for a block hash, given its ID
-    pub fn entry_block_hash<'a>(&'a mut self, block_id: u32) -> Entry<'a, u32, T> {
-        self.state_mut().entry_block_hash(block_id)
+    /// Get cached entry for a block hash, given its ID, or, if not
+    ///  found, use `lookup` to get the corresponding block hash and
+    ///  store it in the cache
+    pub fn get_block_hash_caching<E, F: FnOnce(u32) -> Result<T, E>>(
+        &mut self,
+        id: u32,
+        lookup: F,
+    ) -> Result<&T, E> {
+        self.state_mut().get_block_hash_caching(id, lookup)
     }
 
     /// Store a block's ID and hash to teh cache.

stackslib/src/chainstate/stacks/index/storage.rs

@@ -73,11 +73,8 @@ impl<T: MarfTrieId> BlockMap for TrieFileStorage<T> {
     }
 
     fn get_block_hash_caching(&mut self, id: u32) -> Result<&T, Error> {
-        if !self.is_block_hash_cached(id) {
-            let block_hash = self.get_block_hash(id)?;
-            self.cache.store_block_hash(id, block_hash);
-        }
-        self.cache.ref_block_hash(id).ok_or(Error::NotFoundError)
+        self.cache
+            .get_block_hash_caching(id, |id| trie_sql::get_block_hash(&self.db, id))
     }
 
     fn is_block_hash_cached(&self, id: u32) -> bool {
@@ -110,14 +107,8 @@ impl<T: MarfTrieId> BlockMap for TrieStorageConnection<'_, T> {
     }
 
     fn get_block_hash_caching<'a>(&'a mut self, id: u32) -> Result<&'a T, Error> {
-        match self.cache.entry_block_hash(id) {
-            Entry::Occupied(occupied_entry) => Ok(occupied_entry.into_mut()),
-            Entry::Vacant(vacant_entry) => {
-                let block_hash = trie_sql::get_block_hash(&self.db, id)?;
-                let block_hash_ref = vacant_entry.insert(block_hash);
-                Ok(block_hash_ref)
-            }
-        }
+        self.cache
+            .get_block_hash_caching(id, |id| trie_sql::get_block_hash(&self.db, id))
     }
 
     fn is_block_hash_cached(&self, id: u32) -> bool {
@@ -174,11 +165,8 @@ impl<T: MarfTrieId> BlockMap for TrieSqlHashMapCursor<'_, T> {
     }
 
     fn get_block_hash_caching(&mut self, id: u32) -> Result<&T, Error> {
-        if !self.is_block_hash_cached(id) {
-            let block_hash = self.get_block_hash(id)?;
-            self.cache.store_block_hash(id, block_hash);
-        }
-        self.cache.ref_block_hash(id).ok_or(Error::NotFoundError)
+        self.cache
+            .get_block_hash_caching(id, |id| trie_sql::get_block_hash(&self.db, id))
     }
 
     fn is_block_hash_cached(&self, id: u32) -> bool {