[trie] Introduce GenericTrieNode as replacement of GenericUpdatedTrie…

…Node (near#12496)

This PR introduces the GenericTrieNode which is a generalization on top
of GenericUpdatedTrieNode where the nodes need not be an update type
node.

This is useful in scenarios where we are not dealing with just updates,
for example generalizing iterators over tries etc.

Along with this change this PR renames `GenericUpdatedNodeId` to
`UpdatedNodeId` which is more appropriate as for both memtries and store
tries we have UpdatedNodeId as usize.

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

@@ -6,8 +6,8 @@ use near_primitives::state::FlatStateValue;
 
 use crate::trie::ops::insert_delete::GenericTrieUpdateInsertDelete;
 use crate::trie::ops::interface::{
-    GenericNodeOrIndex, GenericTrieUpdate, GenericTrieValue, GenericUpdatedNodeId,
-    GenericUpdatedTrieNode, GenericUpdatedTrieNodeWithSize,
+    GenericNodeOrIndex, GenericTrieNode, GenericTrieNodeWithSize, GenericTrieUpdate,
+    GenericTrieValue, GenericUpdatedTrieNode, GenericUpdatedTrieNodeWithSize, UpdatedNodeId,
 };
 use crate::trie::trie_recording::TrieRecorder;
 use crate::trie::{Children, MemTrieChanges, TrieRefcountDeltaMap};
@@ -18,49 +18,49 @@ use super::flexible_data::children::ChildrenView;
 use super::metrics::MEM_TRIE_NUM_NODES_CREATED_FROM_UPDATES;
 use super::node::{InputMemTrieNode, MemTrieNodeId, MemTrieNodeView};
 
-pub type UpdatedMemTrieNodeId = usize;
-
 pub type OldOrUpdatedNodeId = GenericNodeOrIndex<MemTrieNodeId>;
 
+pub type MemTrieNode = GenericTrieNode<MemTrieNodeId, FlatStateValue>;
+
 pub type UpdatedMemTrieNode = GenericUpdatedTrieNode<MemTrieNodeId, FlatStateValue>;
 
+pub type MemTrieNodeWithSize = GenericTrieNodeWithSize<MemTrieNodeId, FlatStateValue>;
+
 pub type UpdatedMemTrieNodeWithSize = GenericUpdatedTrieNodeWithSize<MemTrieNodeId, FlatStateValue>;
 
-impl UpdatedMemTrieNodeWithSize {
+impl MemTrieNodeWithSize {
     /// Converts an existing in-memory trie node into an updated one that is
     /// equivalent.
     pub fn from_existing_node_view<'a, M: ArenaMemory>(view: MemTrieNodeView<'a, M>) -> Self {
         let memory_usage = view.memory_usage();
         let node = match view {
-            MemTrieNodeView::Leaf { extension, value } => UpdatedMemTrieNode::Leaf {
+            MemTrieNodeView::Leaf { extension, value } => MemTrieNode::Leaf {
                 extension: extension.to_vec().into_boxed_slice(),
                 value: value.to_flat_value(),
             },
-            MemTrieNodeView::Branch { children, .. } => UpdatedMemTrieNode::Branch {
+            MemTrieNodeView::Branch { children, .. } => MemTrieNode::Branch {
                 children: Box::new(Self::convert_children_to_updated(children)),
                 value: None,
             },
-            MemTrieNodeView::BranchWithValue { children, value, .. } => {
-                UpdatedMemTrieNode::Branch {
-                    children: Box::new(Self::convert_children_to_updated(children)),
-                    value: Some(value.to_flat_value()),
-                }
-            }
-            MemTrieNodeView::Extension { extension, child, .. } => UpdatedMemTrieNode::Extension {
+            MemTrieNodeView::BranchWithValue { children, value, .. } => MemTrieNode::Branch {
+                children: Box::new(Self::convert_children_to_updated(children)),
+                value: Some(value.to_flat_value()),
+            },
+            MemTrieNodeView::Extension { extension, child, .. } => MemTrieNode::Extension {
                 extension: extension.to_vec().into_boxed_slice(),
-                child: OldOrUpdatedNodeId::Old(child.id()),
+                child: child.id(),
             },
         };
         Self { node, memory_usage }
     }
 
     fn convert_children_to_updated<'a, M: ArenaMemory>(
         view: ChildrenView<'a, M>,
-    ) -> [Option<OldOrUpdatedNodeId>; 16] {
+    ) -> [Option<MemTrieNodeId>; 16] {
         let mut children = [None; 16];
         for i in 0..16 {
             if let Some(child) = view.get(i) {
-                children[i] = Some(OldOrUpdatedNodeId::Old(child.id()));
+                children[i] = Some(child.id());
             }
         }
         children
@@ -152,27 +152,27 @@ impl<'a, M: ArenaMemory> GenericTrieUpdate<'a, MemTrieNodeId, FlatStateValue>
     fn ensure_updated(
         &mut self,
         node: GenericNodeOrIndex<MemTrieNodeId>,
-    ) -> Result<GenericUpdatedNodeId, StorageError> {
+    ) -> Result<UpdatedNodeId, StorageError> {
         Ok(match node {
             GenericNodeOrIndex::Old(node_id) => self.convert_existing_to_updated(Some(node_id)),
             GenericNodeOrIndex::Updated(node_id) => node_id,
         })
     }
 
-    fn take_node(&mut self, index: UpdatedMemTrieNodeId) -> UpdatedMemTrieNodeWithSize {
+    fn take_node(&mut self, index: UpdatedNodeId) -> UpdatedMemTrieNodeWithSize {
         self.updated_nodes.get_mut(index).unwrap().take().expect("Node taken twice")
     }
 
-    fn place_node_at(&mut self, index: UpdatedMemTrieNodeId, node: UpdatedMemTrieNodeWithSize) {
+    fn place_node_at(&mut self, index: UpdatedNodeId, node: UpdatedMemTrieNodeWithSize) {
         assert!(self.updated_nodes[index].is_none(), "Node placed twice");
         self.updated_nodes[index] = Some(node);
     }
 
-    fn get_node_ref(&self, node_id: GenericUpdatedNodeId) -> &UpdatedMemTrieNodeWithSize {
+    fn get_node_ref(&self, node_id: UpdatedNodeId) -> &UpdatedMemTrieNodeWithSize {
         self.updated_nodes[node_id].as_ref().unwrap()
     }
 
-    fn place_node(&mut self, node: UpdatedMemTrieNodeWithSize) -> GenericUpdatedNodeId {
+    fn place_node(&mut self, node: UpdatedMemTrieNodeWithSize) -> UpdatedNodeId {
         let index = self.updated_nodes.len();
         self.updated_nodes.push(Some(node));
         index
@@ -232,7 +232,7 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> {
     }
 
     /// Creates a new updated node, assigning it a new ID.
-    fn new_updated_node(&mut self, node: UpdatedMemTrieNodeWithSize) -> UpdatedMemTrieNodeId {
+    fn new_updated_node(&mut self, node: UpdatedMemTrieNodeWithSize) -> UpdatedNodeId {
         let index = self.updated_nodes.len();
         self.updated_nodes.push(Some(node));
         index
@@ -245,15 +245,15 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> {
     ///
     /// If the original node is None, it is a marker for the root of an empty
     /// trie.
-    fn convert_existing_to_updated(&mut self, node: Option<MemTrieNodeId>) -> UpdatedMemTrieNodeId {
+    fn convert_existing_to_updated(&mut self, node: Option<MemTrieNodeId>) -> UpdatedNodeId {
         let Some(node) = node else {
             return self.new_updated_node(UpdatedMemTrieNodeWithSize::empty());
         };
         let node_view = node.as_ptr(self.memory).view();
         if let Some(tracked_trie_changes) = self.nodes_tracker.as_mut() {
             tracked_trie_changes.record(&node_view);
         }
-        self.new_updated_node(UpdatedMemTrieNodeWithSize::from_existing_node_view(node_view))
+        self.new_updated_node(MemTrieNodeWithSize::from_existing_node_view(node_view).into())
     }
 
     /// Inserts the given key value pair into the trie.
@@ -279,9 +279,9 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> {
     /// updated nodes. After this function, `ordered_nodes` contains the IDs of
     /// the updated nodes in the order they should be created.
     fn post_order_traverse_updated_nodes(
-        node_id: UpdatedMemTrieNodeId,
+        node_id: UpdatedNodeId,
         updated_nodes: &Vec<Option<UpdatedMemTrieNodeWithSize>>,
-        ordered_nodes: &mut Vec<UpdatedMemTrieNodeId>,
+        ordered_nodes: &mut Vec<UpdatedNodeId>,
     ) {
         let node = updated_nodes[node_id].as_ref().unwrap();
         match &node.node {
@@ -320,9 +320,9 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> {
     /// `updated_nodes` must be indexed by the node IDs in `ordered_nodes`.
     pub(crate) fn compute_hashes_and_serialized_nodes(
         &self,
-        ordered_nodes: &Vec<UpdatedMemTrieNodeId>,
+        ordered_nodes: &Vec<UpdatedNodeId>,
         updated_nodes: &Vec<Option<UpdatedMemTrieNodeWithSize>>,
-    ) -> Vec<(UpdatedMemTrieNodeId, CryptoHash, Vec<u8>)> {
+    ) -> Vec<(UpdatedNodeId, CryptoHash, Vec<u8>)> {
         let memory = self.memory;
         let mut result = Vec::<(CryptoHash, Vec<u8>)>::new();
         for _ in 0..updated_nodes.len() {
@@ -459,15 +459,15 @@ pub(super) fn construct_root_from_changes<A: ArenaMut>(
 ) -> Option<MemTrieNodeId> {
     let mut last_node_id: Option<MemTrieNodeId> = None;
     let map_to_new_node_id = |node_id: OldOrUpdatedNodeId,
-                              old_to_new_map: &HashMap<UpdatedMemTrieNodeId, MemTrieNodeId>|
+                              old_to_new_map: &HashMap<UpdatedNodeId, MemTrieNodeId>|
      -> MemTrieNodeId {
         match node_id {
             OldOrUpdatedNodeId::Updated(node_id) => *old_to_new_map.get(&node_id).unwrap(),
             OldOrUpdatedNodeId::Old(node_id) => node_id,
         }
     };
 
-    let mut updated_to_new_map = HashMap::<UpdatedMemTrieNodeId, MemTrieNodeId>::new();
+    let mut updated_to_new_map = HashMap::<UpdatedNodeId, MemTrieNodeId>::new();
     let updated_nodes = &changes.updated_nodes;
     let node_ids_with_hashes = &changes.node_ids_with_hashes;
     for (node_id, node_hash) in node_ids_with_hashes.iter() {

core/store/src/trie/mod.rs

@@ -18,7 +18,7 @@ pub use crate::trie::trie_storage::{TrieCache, TrieCachingStorage, TrieDBStorage
 use crate::StorageError;
 use borsh::{BorshDeserialize, BorshSerialize};
 pub use from_flat::construct_trie_from_flat;
-use mem::mem_trie_update::{TrackingMode, UpdatedMemTrieNodeId, UpdatedMemTrieNodeWithSize};
+use mem::mem_trie_update::{TrackingMode, UpdatedMemTrieNodeWithSize};
 use mem::mem_tries::MemTries;
 use near_primitives::challenge::PartialState;
 use near_primitives::hash::{hash, CryptoHash};
@@ -31,9 +31,9 @@ use near_primitives::types::{AccountId, StateRoot, StateRootNode};
 use near_schema_checker_lib::ProtocolSchema;
 use near_vm_runner::ContractCode;
 use ops::insert_delete::GenericTrieUpdateInsertDelete;
-use ops::interface::GenericTrieValue;
 #[cfg(test)]
-use ops::interface::{GenericNodeOrIndex, GenericTrieUpdate};
+use ops::interface::{GenericNodeOrIndex, GenericTrieNode, GenericTrieUpdate};
+use ops::interface::{GenericTrieValue, UpdatedNodeId};
 pub use raw_node::{Children, RawTrieNode, RawTrieNodeWithSize};
 use std::cell::RefCell;
 use std::collections::{BTreeMap, HashSet};
@@ -43,9 +43,9 @@ use std::ops::DerefMut;
 use std::str;
 use std::sync::{Arc, RwLock, RwLockReadGuard};
 pub use trie_recording::{SubtreeSize, TrieRecorder, TrieRecorderStats};
-#[cfg(test)]
-use trie_storage_update::UpdatedTrieStorageNode;
-use trie_storage_update::{TrieStorageUpdate, UpdatedTrieStorageNodeWithSize};
+use trie_storage_update::{
+    TrieStorageNodeWithSize, TrieStorageUpdate, UpdatedTrieStorageNodeWithSize,
+};
 
 pub mod accounting_cache;
 mod config;
@@ -198,14 +198,14 @@ impl UpdatedTrieStorageNodeWithSize {
         spaces: &mut String,
     ) -> std::fmt::Result {
         match &self.node {
-            UpdatedTrieStorageNode::Empty => {
+            GenericTrieNode::Empty => {
                 write!(f, "{}Empty", spaces)?;
             }
-            UpdatedTrieStorageNode::Leaf { extension, .. } => {
+            GenericTrieNode::Leaf { extension, .. } => {
                 let slice = NibbleSlice::from_encoded(&extension);
                 write!(f, "{}Leaf({:?}, val)", spaces, slice.0)?;
             }
-            UpdatedTrieStorageNode::Branch { children, value } => {
+            GenericTrieNode::Branch { children, value } => {
                 writeln!(
                     f,
                     "{}Branch({}){{",
@@ -232,7 +232,7 @@ impl UpdatedTrieStorageNodeWithSize {
                 spaces.remove(spaces.len() - 1);
                 write!(f, "{}}}", spaces)?;
             }
-            UpdatedTrieStorageNode::Extension { extension, child } => {
+            GenericTrieNode::Extension { extension, child } => {
                 let slice = NibbleSlice::from_encoded(&extension);
                 writeln!(f, "{}Extension({:?})", spaces, slice)?;
                 spaces.push(' ');
@@ -533,7 +533,7 @@ pub struct MemTrieChanges {
     /// Node ids with hashes of updated nodes.
     /// Should be in the post-order traversal of the updated nodes.
     /// It implies that the root node is the last one in the list.
-    node_ids_with_hashes: Vec<(UpdatedMemTrieNodeId, CryptoHash)>,
+    node_ids_with_hashes: Vec<(UpdatedNodeId, CryptoHash)>,
     updated_nodes: Vec<Option<UpdatedMemTrieNodeWithSize>>,
 }
 
@@ -771,7 +771,7 @@ impl Trie {
         trie
     }
 
-    /// Get statisitics about the recorded trie. Useful for observability and debugging.
+    /// Get statistics about the recorded trie. Useful for observability and debugging.
     /// This scans all of the recorded data, so could potentially be expensive to run.
     pub fn recorder_stats(&self) -> Option<TrieRecorderStats> {
         self.recorder.as_ref().map(|recorder| recorder.borrow().get_stats(&self.root))
@@ -882,23 +882,23 @@ impl Trie {
                     .expect("storage failure")
                     .expect("node cannot be Empty")
                     .1;
-                UpdatedTrieStorageNodeWithSize::from_raw_trie_node_with_size(raw_node)
+                TrieStorageNodeWithSize::from_raw_trie_node_with_size(raw_node).into()
             }
         };
 
         let mut memory_usage_naive = node.memory_usage_direct();
         match &node {
-            UpdatedTrieStorageNode::Empty => {}
-            UpdatedTrieStorageNode::Leaf { .. } => {}
-            UpdatedTrieStorageNode::Branch { children, .. } => {
+            GenericTrieNode::Empty => {}
+            GenericTrieNode::Leaf { .. } => {}
+            GenericTrieNode::Branch { children, .. } => {
                 memory_usage_naive += children
                     .iter()
                     .filter_map(|handle| {
                         handle.as_ref().map(|h| self.memory_usage_verify(trie_update, *h))
                     })
                     .sum::<u64>();
             }
-            UpdatedTrieStorageNode::Extension { child, .. } => {
+            GenericTrieNode::Extension { child, .. } => {
                 memory_usage_naive += self.memory_usage_verify(trie_update, *child);
             }
         };
@@ -1253,7 +1253,7 @@ impl Trie {
             None => Ok(trie_update.store(UpdatedTrieStorageNodeWithSize::empty())),
             Some((_, node)) => {
                 let result = trie_update
-                    .store(UpdatedTrieStorageNodeWithSize::from_raw_trie_node_with_size(node));
+                    .store(TrieStorageNodeWithSize::from_raw_trie_node_with_size(node).into());
                 trie_update.refcount_changes.subtract(*hash, 1);
                 Ok(result)
             }

core/store/src/trie/ops/insert_delete.rs

@@ -3,8 +3,8 @@ use near_primitives::errors::StorageError;
 use crate::NibbleSlice;
 
 use super::interface::{
-    GenericNodeOrIndex, GenericTrieUpdate, GenericTrieValue, GenericUpdatedNodeId,
-    GenericUpdatedTrieNode, GenericUpdatedTrieNodeWithSize, HasValueLength,
+    GenericNodeOrIndex, GenericTrieUpdate, GenericTrieValue, GenericUpdatedTrieNode,
+    GenericUpdatedTrieNodeWithSize, HasValueLength, UpdatedNodeId,
 };
 use super::squash::GenericTrieUpdateSquash;
 
@@ -16,10 +16,10 @@ where
 {
     fn calc_memory_usage_and_store(
         &mut self,
-        node_id: GenericUpdatedNodeId,
+        node_id: UpdatedNodeId,
         children_memory_usage: u64,
         new_node: GenericUpdatedTrieNode<N, V>,
-        old_child: Option<GenericUpdatedNodeId>,
+        old_child: Option<UpdatedNodeId>,
     ) {
         let new_memory_usage =
             children_memory_usage.saturating_add(new_node.memory_usage_direct()).saturating_sub(
@@ -38,7 +38,7 @@ where
     /// created nodes - that's done at the end.
     fn generic_insert(
         &mut self,
-        mut node_id: GenericUpdatedNodeId,
+        mut node_id: UpdatedNodeId,
         key: &[u8],
         value: GenericTrieValue,
     ) -> Result<(), StorageError> {
@@ -304,7 +304,7 @@ where
     /// Deleting a non-existent key is allowed, and is a no-op.
     fn generic_delete(
         &mut self,
-        mut node_id: GenericUpdatedNodeId,
+        mut node_id: UpdatedNodeId,
         key: &[u8],
     ) -> Result<(), StorageError> {
         let mut partial = NibbleSlice::new(key);