feedback

core/store/src/trie/insert_delete.rs

@@ -1,4 +1,5 @@
 use super::TrieRefcountDeltaMap;
+use crate::trie::mem::updating::GenericTrieUpdate;
 use crate::trie::nibble_slice::NibbleSlice;
 use crate::trie::{
     Children, NodeHandle, RawTrieNode, RawTrieNodeWithSize, StorageHandle, StorageValueHandle,
@@ -11,7 +12,7 @@ use near_primitives::state::{ValueRef, ValueUpdate};
 
 pub(crate) struct NodesStorage<'a> {
     nodes: Vec<Option<TrieNodeWithSize>>,
-    values: Vec<Option<Vec<u8>>>,
+    pub(crate) values: Vec<Option<Vec<u8>>>,
     pub(crate) refcount_changes: TrieRefcountDeltaMap,
     pub(crate) trie: &'a Trie,
 }
@@ -58,19 +59,6 @@ impl<'a> NodesStorage<'a> {
         StorageHandle(self.nodes.len() - 1)
     }
 
-    pub(crate) fn store_value(&mut self, value: ValueUpdate) -> StorageValueHandle {
-        let ValueUpdate::MemtrieAndDisk(value) = value else {
-            unimplemented!(
-                "NodesStorage for Trie doesn't support value {value:?} \
-                because disk updates must be generated."
-            );
-        };
-
-        let value_len = value.len();
-        self.values.push(Some(value));
-        StorageValueHandle(self.values.len() - 1, value_len)
-    }
-
     pub(crate) fn value_ref(&self, handle: StorageValueHandle) -> &[u8] {
         self.values
             .get(handle.0)
@@ -111,24 +99,20 @@ impl Trie {
             let children_memory_usage = memory_usage - node.memory_usage_direct(memory);
             match node {
                 TrieNode::Empty => {
-                    let value_handle = memory.store_value(value);
-                    let leaf_node = TrieNode::Leaf(
-                        partial.encoded(true).into_vec(),
-                        ValueHandle::InMemory(value_handle),
-                    );
+                    let value_handle = memory.generic_store_value(value);
+                    let leaf_node = TrieNode::Leaf(partial.encoded(true).into_vec(), value_handle);
                     let memory_usage = leaf_node.memory_usage_direct(memory);
                     memory.store_at(handle, TrieNodeWithSize { node: leaf_node, memory_usage });
                     break;
                 }
                 TrieNode::Branch(mut children, existing_value) => {
                     // If the key ends here, store the value in branch's value.
                     if partial.is_empty() {
-                        if let Some(value) = &existing_value {
-                            self.delete_value(memory, value)?;
+                        if let Some(value) = existing_value {
+                            memory.generic_delete_value(value)?;
                         }
-                        let value_handle = memory.store_value(value);
-                        let new_node =
-                            TrieNode::Branch(children, Some(ValueHandle::InMemory(value_handle)));
+                        let value_handle = memory.generic_store_value(value);
+                        let new_node = TrieNode::Branch(children, Some(value_handle));
                         let new_memory_usage =
                             children_memory_usage + new_node.memory_usage_direct(memory);
                         memory.store_at(handle, TrieNodeWithSize::new(new_node, new_memory_usage));
@@ -160,9 +144,9 @@ impl Trie {
                     let common_prefix = partial.common_prefix(&existing_key);
                     if common_prefix == existing_key.len() && common_prefix == partial.len() {
                         // Equivalent leaf.
-                        self.delete_value(memory, &existing_value)?;
-                        let value_handle = memory.store_value(value);
-                        let node = TrieNode::Leaf(key, ValueHandle::InMemory(value_handle));
+                        memory.generic_delete_value(existing_value)?;
+                        let value_handle = memory.generic_store_value(value);
+                        let node = TrieNode::Leaf(key, value_handle);
                         let memory_usage = node.memory_usage_direct(memory);
                         memory.store_at(handle, TrieNodeWithSize { node, memory_usage });
                         break;
@@ -342,7 +326,7 @@ impl Trie {
                 }
                 TrieNode::Leaf(key, value) => {
                     if NibbleSlice::from_encoded(&key).0 == partial {
-                        self.delete_value(memory, &value)?;
+                        memory.generic_delete_value(value)?;
                         memory.store_at(handle, TrieNodeWithSize::empty());
                         break;
                     } else {
@@ -367,7 +351,7 @@ impl Trie {
                             key_deleted = false;
                             break;
                         }
-                        self.delete_value(memory, &value.unwrap())?;
+                        memory.generic_delete_value(value.unwrap())?;
                         Trie::calc_memory_usage_and_store(
                             memory,
                             handle,

core/store/src/trie/mem/loading.rs

@@ -158,9 +158,9 @@ pub fn load_trie_from_flat_state_and_delta(
             for (key, value) in changes.0 {
                 match value {
                     Some(value) => {
-                        trie_update.insert_memtrie_only(&key, value);
+                        trie_update.insert_memtrie_only(&key, value)?;
                     }
-                    None => trie_update.delete(&key),
+                    None => trie_update.delete(&key)?,
                 };
             }
 

core/store/src/trie/mem/resharding.rs

@@ -349,7 +349,7 @@ mod tests {
         let mut memtries = MemTries::new(ShardUId::single_shard());
         let mut update = memtries.update(Trie::EMPTY_ROOT, false).unwrap();
         for (key, value) in initial_entries {
-            update.insert(&key, value);
+            update.insert(&key, value).unwrap();
         }
         let memtrie_changes = update.to_mem_trie_changes_only();
         let state_root = memtries.apply_memtrie_changes(0, &memtrie_changes);