simplify retrieve_raw_bytes logic

near · Aug 16, 2019 · 0383590 · 0383590
1 parent 4d9d5cf
commit 0383590
Showing 1 changed file with 14 additions and 30 deletions.
diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs
@@ -336,11 +336,7 @@ pub struct TrieRecordingStorage {
 impl TrieStorage for TrieRecordingStorage {
     fn retrieve_raw_bytes(&self, hash: &CryptoHash) -> Option<Vec<u8>> {
         let result = self.storage.retrieve_raw_bytes(hash);
-        // storage has RcTrieNode but record RawTrieNode. Refcount doesn't matter.
-        self.recorded
-            .lock()
-            .expect(POISONED_LOCK_ERR)
-            .insert(*hash, result.as_ref().map(|vec| vec[0..vec.len() - 4].to_vec()));
+        self.recorded.lock().expect(POISONED_LOCK_ERR).insert(*hash, result.clone());
         result
     }
 
@@ -360,12 +356,7 @@ pub struct TrieMemoryPartialStorage {
 
 impl TrieStorage for TrieMemoryPartialStorage {
     fn retrieve_raw_bytes(&self, hash: &CryptoHash) -> Option<Vec<u8>> {
-        // RawTrieNode -> RcTrieNode with a refcount of 0. Refcount doesn't matter
-        self.recorded_storage.get(hash).map(|vec| {
-            let mut vec = vec.clone();
-            vec.extend_from_slice(&[0, 0, 0, 0]);
-            vec
-        })
+        self.recorded_storage.get(hash).cloned()
     }
 
     fn retrieve_rc(&self, _hash: &CryptoHash) -> Option<u32> {
@@ -391,15 +382,16 @@ impl TrieStorage for TrieCachingStorage {
         let result = {
             let mut guard = self.cache.lock().expect(POISONED_LOCK_ERR);
             if let Some(val) = (*guard).cache_get(hash) {
-                val.clone()
+                val.as_ref().map(|vec| vec[0..vec.len() - 4].to_vec())
             } else {
                 let result = if let Ok(Some(bytes)) = self.store.get(COL_STATE, hash.as_ref()) {
                     Some(bytes)
                 } else {
                     None
                 };
-                (*guard).cache_set(*hash, result.clone());
-                result
+                let raw_node = result.as_ref().map(|vec| vec[0..vec.len() - 4].to_vec());
+                (*guard).cache_set(*hash, result);
+                raw_node
             }
         };
         // the only time a node can be missing from storage is if we're validating with partial
@@ -606,21 +598,13 @@ impl Trie {
             Ok(memory.store(TrieNode::Empty))
         } else {
             if let Some(bytes) = self.storage.retrieve_raw_bytes(hash) {
-                match RcTrieNode::decode(&bytes) {
-                    Ok((value, _)) => {
+                match RawTrieNode::decode(&bytes) {
+                    Ok(value) => {
                         let result = memory.store(TrieNode::new(value));
                         memory
                             .refcount_changes
                             .entry(*hash)
-                            .or_insert_with(|| {
-                                (
-                                    RcTrieNode::decode_raw(&bytes)
-                                        .expect("calling after decode()")
-                                        .0
-                                        .to_vec(),
-                                    0,
-                                )
-                            })
+                            .or_insert_with(|| (bytes.to_vec(), 0))
                             .1 -= 1;
                         Ok(result)
                     }
@@ -637,8 +621,8 @@ impl Trie {
             return Ok(TrieNode::Empty);
         }
         if let Some(bytes) = self.storage.retrieve_raw_bytes(hash) {
-            match RcTrieNode::decode(&bytes) {
-                Ok((value, _)) => Ok(TrieNode::new(value)),
+            match RawTrieNode::decode(&bytes) {
+                Ok(value) => Ok(TrieNode::new(value)),
                 Err(_) => Err(format!("Failed to decode node {}", hash)),
             }
         } else {
@@ -654,9 +638,9 @@ impl Trie {
                 return Ok(None);
             }
             let node = match self.storage.retrieve_raw_bytes(&hash) {
-                Some(bytes) => RcTrieNode::decode(&bytes)
-                    .map(|trie_node| trie_node.0)
-                    .map_err(|_| "Failed to decode node".to_string())?,
+                Some(bytes) => {
+                    RawTrieNode::decode(&bytes).map_err(|_| "Failed to decode node".to_string())?
+                }
                 _ => return Err(format!("Node {} not found in storage", hash)),
             };