[resharding] Simplify GenericTrieUpdateRetain interface and move tests (

#12518)

core/store/src/trie/mem/resharding.rs had the implementation for the
resharding retain function for both memtrie and trie.

That didn't feel like a great as this ideally should have been part of
mem_trie_update.rs and trie_storage_update.rs respectively.

Additionally, it felt like we are exposing the inner function
implementation with `self.retain_multi_range_recursive(0, vec![],
&intervals_nibbles)`.

This PR wraps the core functionality of `retain_split_shard` into the
`GenericTrieUpdateRetain` interface while keeping the private functions
of GenericTrieUpdateRetain in a separate inner trait
`GenericTrieUpdateRetainInner` that is not exposed outside the file.

The tests in core/store/src/trie/mem/resharding.rs are moved to a
separate file without any functional change.

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

@@ -3,12 +3,14 @@ use std::collections::{BTreeMap, HashMap};
 use near_primitives::errors::StorageError;
 use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::state::FlatStateValue;
+use near_primitives::types::AccountId;
 
 use crate::trie::ops::insert_delete::GenericTrieUpdateInsertDelete;
 use crate::trie::ops::interface::{
     GenericNodeOrIndex, GenericTrieNode, GenericTrieNodeWithSize, GenericTrieUpdate,
     GenericTrieValue, GenericUpdatedTrieNode, GenericUpdatedTrieNodeWithSize, UpdatedNodeId,
 };
+use crate::trie::ops::resharding::{GenericTrieUpdateRetain, RetainMode};
 use crate::trie::trie_recording::TrieRecorder;
 use crate::trie::{Children, MemTrieChanges, TrieRefcountDeltaMap};
 use crate::{RawTrieNode, RawTrieNodeWithSize, TrieChanges};
@@ -449,6 +451,21 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> {
             mem_trie_changes: Some(mem_trie_changes),
         }
     }
+
+    /// Splits the trie, separating entries by the boundary account.
+    /// Leaves the left or right part of the trie, depending on the retain mode.
+    ///
+    /// Returns the changes to be applied to in-memory trie and the proof of
+    /// the split operation. Doesn't modifies trie itself, it's a caller's
+    /// responsibility to apply the changes.
+    pub fn retain_split_shard(
+        mut self,
+        boundary_account: &AccountId,
+        retain_mode: RetainMode,
+    ) -> TrieChanges {
+        GenericTrieUpdateRetain::retain_split_shard(&mut self, boundary_account, retain_mode);
+        self.to_trie_changes()
+    }
 }
 
 /// Applies the given memtrie changes to the in-memory trie data structure.

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

@@ -125,7 +125,7 @@ impl MemTries {
     }
 
     /// Returns the root node corresponding to the given state root.
-    pub(super) fn get_root(
+    pub fn get_root(
         &self,
         state_root: &CryptoHash,
     ) -> Result<MemTrieNodePtr<HybridArenaMemory>, StorageError> {

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

@@ -12,7 +12,6 @@ pub mod metrics;
 pub(crate) mod nibbles_utils;
 pub mod node;
 mod parallel_loader;
-pub mod resharding;
 
 /// Check this, because in the code we conveniently assume usize is 8 bytes.
 /// In-memory trie can't possibly work under 32-bit anyway.