fix: generic trie value updates (#12344)

Preparation step for #12324.

Introduce `generic_store_value` and `generic_delete_value`. These are
different ways for MemTrie and TrieStorage to process value updates.
TLDR: memtries can work with inlined values; trie storages must always
have full values. Also the way to record accessed nodes is different.

Also, solving some issues on the way. For example, interfaces of
`Trie::insert` and `MemTrieUpdate::insert_impl` currently diverge. The
latter needs both `FlatStateValue` for memtrie changes and
`Option<Vec<u8>>` for disk changes. There is a good reason for that - if
memtrie is loaded from flat state, we don't need to produce disk trie
changes, so it is enough to load `FlatStateValue`s.

However, the current interface is quite loose, so it is a net
improvement to make it more strict by requiring everyone to give
`ValueUpdate::MemtrieAndDisk` or `::MemtrieOnly`. The drawback is that
`Trie::insert` technically can receive `ValueUpdate::MemtrieOnly`, it
just will panic. But well... it still looks better than unclear memtrie
interface with two values and hashes on the way.

core/primitives/src/state.rs

@@ -110,3 +110,14 @@ impl FlatStateValue {
         }
     }
 }
+
+/// Value to insert to trie or update existing value in the trie.
+#[derive(Debug, Clone)]
+pub enum GenericTrieValue {
+    /// Value to update both memtrie and trie storage. Full value is required
+    /// for that.
+    MemtrieAndDisk(Vec<u8>),
+    /// Value to update only memtrie. In such case it is enough to have a
+    /// `FlatStateValue`.
+    MemtrieOnly(FlatStateValue),
+}

core/store/src/trie/insert_delete.rs

@@ -1,5 +1,5 @@
 use super::TrieRefcountDeltaMap;
-use crate::trie::mem::updating::GenericTrieUpdateSquash;
+use crate::trie::mem::updating::{GenericTrieUpdate, GenericTrieUpdateSquash};
 use crate::trie::nibble_slice::NibbleSlice;
 use crate::trie::{
     Children, NodeHandle, RawTrieNode, RawTrieNodeWithSize, StorageHandle, StorageValueHandle,
@@ -8,11 +8,11 @@ use crate::trie::{
 use crate::{StorageError, Trie, TrieChanges};
 use borsh::BorshSerialize;
 use near_primitives::hash::{hash, CryptoHash};
-use near_primitives::state::ValueRef;
+use near_primitives::state::{GenericTrieValue, ValueRef};
 
 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,
 }
@@ -59,12 +59,6 @@ impl<'a> NodesStorage<'a> {
         StorageHandle(self.nodes.len() - 1)
     }
 
-    pub(crate) fn store_value(&mut self, value: Vec<u8>) -> StorageValueHandle {
-        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)
@@ -93,11 +87,10 @@ impl Trie {
         memory: &mut NodesStorage,
         node: StorageHandle,
         partial: NibbleSlice<'_>,
-        value: Vec<u8>,
+        value: GenericTrieValue,
     ) -> Result<StorageHandle, StorageError> {
         let root_handle = node;
         let mut handle = node;
-        let mut value = Some(value);
         let mut partial = partial;
         let mut path = Vec::new();
         loop {
@@ -106,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.take().unwrap());
-                    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.take().unwrap());
-                        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));
@@ -155,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.take().unwrap());
-                        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;
@@ -337,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 {
@@ -362,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

@@ -346,7 +346,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);