fix: read only nodes along the path on non-existent key deletion

core/store/src/test_utils.rs

@@ -94,11 +94,15 @@ impl TestTriesBuilder {
         self
     }
 
-    pub fn with_flat_storage(mut self) -> Self {
-        self.enable_flat_storage = true;
+    pub fn with_flat_storage(mut self, enable: bool) -> Self {
+        self.enable_flat_storage = enable;
         self
     }
 
+    pub fn with_enabled_flat_storage(self) -> Self {
+        self.with_flat_storage(true)
+    }
+
     pub fn with_in_memory_tries(mut self) -> Self {
         self.enable_in_memory_tries = true;
         self

core/store/src/trie/insert_delete.rs

@@ -320,13 +320,15 @@ impl Trie {
         let mut partial = partial;
         let root_node = handle;
         let mut path: Vec<StorageHandle> = Vec::new();
+        let mut key_deleted = true;
         loop {
             path.push(handle);
             let TrieNodeWithSize { node, memory_usage } = memory.destroy(handle);
             let children_memory_usage = memory_usage - node.memory_usage_direct(memory);
             match node {
                 TrieNode::Empty => {
                     memory.store_at(handle, TrieNodeWithSize::empty());
+                    key_deleted = false;
                     break;
                 }
                 TrieNode::Leaf(key, value) => {
@@ -338,25 +340,33 @@ impl Trie {
                         let leaf_node = TrieNode::Leaf(key, value);
                         let memory_usage = leaf_node.memory_usage_direct(memory);
                         memory.store_at(handle, TrieNodeWithSize::new(leaf_node, memory_usage));
+                        key_deleted = false;
                         break;
                     }
                 }
                 TrieNode::Branch(mut children, value) => {
                     if partial.is_empty() {
-                        if let Some(value) = &value {
-                            self.delete_value(memory, value)?;
-                        }
-                        if children.iter().count() == 0 {
-                            memory.store_at(handle, TrieNodeWithSize::empty());
-                        } else {
-                            Trie::calc_memory_usage_and_store(
-                                memory,
+                        if value.is_none() {
+                            // Key being deleted doesn't exist.
+                            memory.store_at(
                                 handle,
-                                children_memory_usage,
-                                TrieNode::Branch(children, None),
-                                None,
+                                TrieNodeWithSize::new(
+                                    TrieNode::Branch(children, value),
+                                    memory_usage,
+                                ),
                             );
+                            key_deleted = false;
+                            break;
                         }
+                        self.delete_value(memory, &value.unwrap())?;
+                        Trie::calc_memory_usage_and_store(
+                            memory,
+                            handle,
+                            children_memory_usage,
+                            TrieNode::Branch(children, None),
+                            None,
+                        );
+                        // if needed, branch will be squashed at the end of the function.
                         break;
                     } else {
                         let child = &mut children[partial.at(0)];
@@ -386,6 +396,7 @@ impl Trie {
                                     memory_usage,
                                 ),
                             );
+                            key_deleted = false;
                             break;
                         }
                     }
@@ -415,71 +426,85 @@ impl Trie {
                             handle,
                             TrieNodeWithSize::new(TrieNode::Extension(key, child), memory_usage),
                         );
+                        key_deleted = false;
                         break;
                     }
                 }
             }
         }
-        self.fix_nodes(memory, path)?;
+        self.fix_nodes(memory, path, key_deleted)?;
         Ok(root_node)
     }
 
+    /// Iterates over nodes in `path`, changing their types where needed,
+    /// if `key_deleted` is true, so trie structure has to change.
+    /// If `key_deleted` is false, only recomputes memory usages along the path.
     fn fix_nodes(
         &self,
         memory: &mut NodesStorage,
         path: Vec<StorageHandle>,
+        key_deleted: bool,
     ) -> Result<(), StorageError> {
         let mut child_memory_usage = 0;
         for handle in path.into_iter().rev() {
             let TrieNodeWithSize { node, memory_usage } = memory.destroy(handle);
             let memory_usage = memory_usage + child_memory_usage;
-            match node {
-                TrieNode::Empty => {
-                    memory.store_at(handle, TrieNodeWithSize::empty());
-                }
-                TrieNode::Leaf(key, value) => {
-                    memory.store_at(
-                        handle,
-                        TrieNodeWithSize::new(TrieNode::Leaf(key, value), memory_usage),
-                    );
-                }
-                TrieNode::Branch(mut children, value) => {
-                    for child in children.0.iter_mut() {
-                        if let Some(NodeHandle::InMemory(h)) = child {
-                            if let TrieNode::Empty = memory.node_ref(*h).node {
-                                *child = None
+            if key_deleted {
+                match node {
+                    TrieNode::Empty => {
+                        memory.store_at(handle, TrieNodeWithSize::empty());
+                    }
+                    TrieNode::Leaf(key, value) => {
+                        memory.store_at(
+                            handle,
+                            TrieNodeWithSize::new(TrieNode::Leaf(key, value), memory_usage),
+                        );
+                    }
+                    TrieNode::Branch(mut children, value) => {
+                        for child in children.0.iter_mut() {
+                            if let Some(NodeHandle::InMemory(h)) = child {
+                                if let TrieNode::Empty = memory.node_ref(*h).node {
+                                    *child = None
+                                }
                             }
                         }
-                    }
-                    let num_children = children.iter().count();
-                    if num_children == 0 {
-                        if let Some(value) = value {
-                            let empty = NibbleSlice::new(&[]).encoded(true).into_vec();
-                            let leaf_node = TrieNode::Leaf(empty, value);
-                            let memory_usage = leaf_node.memory_usage_direct(memory);
-                            memory.store_at(handle, TrieNodeWithSize::new(leaf_node, memory_usage));
+                        let num_children = children.iter().count();
+                        if num_children == 0 {
+                            if let Some(value) = value {
+                                let empty = NibbleSlice::new(&[]).encoded(true).into_vec();
+                                let leaf_node = TrieNode::Leaf(empty, value);
+                                let memory_usage = leaf_node.memory_usage_direct(memory);
+                                memory.store_at(
+                                    handle,
+                                    TrieNodeWithSize::new(leaf_node, memory_usage),
+                                );
+                            } else {
+                                memory.store_at(handle, TrieNodeWithSize::empty());
+                            }
+                        } else if num_children == 1 && value.is_none() {
+                            // Branch with one child becomes extension
+                            // Extension followed by leaf becomes leaf
+                            // Extension followed by extension becomes extension
+                            let idx = children.iter().next().unwrap().0;
+                            let child = children[idx].take().unwrap();
+                            let key = NibbleSlice::new(&[(idx << 4) as u8])
+                                .encoded_leftmost(1, false)
+                                .into_vec();
+                            self.fix_extension_node(memory, handle, key, child)?;
                         } else {
-                            memory.store_at(handle, TrieNodeWithSize::empty());
+                            let node = TrieNodeWithSize::new(
+                                TrieNode::Branch(children, value),
+                                memory_usage,
+                            );
+                            memory.store_at(handle, node);
                         }
-                    } else if num_children == 1 && value.is_none() {
-                        // Branch with one child becomes extension
-                        // Extension followed by leaf becomes leaf
-                        // Extension followed by extension becomes extension
-                        let idx = children.iter().next().unwrap().0;
-                        let child = children[idx].take().unwrap();
-                        let key = NibbleSlice::new(&[(idx << 4) as u8])
-                            .encoded_leftmost(1, false)
-                            .into_vec();
+                    }
+                    TrieNode::Extension(key, child) => {
                         self.fix_extension_node(memory, handle, key, child)?;
-                    } else {
-                        let node =
-                            TrieNodeWithSize::new(TrieNode::Branch(children, value), memory_usage);
-                        memory.store_at(handle, node);
                     }
                 }
-                TrieNode::Extension(key, child) => {
-                    self.fix_extension_node(memory, handle, key, child)?;
-                }
+            } else {
+                memory.store_at(handle, TrieNodeWithSize { node, memory_usage });
             }
             child_memory_usage = memory.node_ref(handle).memory_usage;
         }

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

@@ -197,7 +197,7 @@ mod tests {
     use rand::{Rng, SeedableRng};
 
     fn check(keys: Vec<Vec<u8>>) {
-        let shard_tries = TestTriesBuilder::new().with_flat_storage().build();
+        let shard_tries = TestTriesBuilder::new().with_enabled_flat_storage().build();
         let shard_uid = ShardUId::single_shard();
         let changes = keys.iter().map(|key| (key.to_vec(), Some(key.to_vec()))).collect::<Vec<_>>();
         let changes = simplify_changes(&changes);

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

@@ -1306,6 +1306,7 @@ mod tests {
                         key.push(nibble0 << 4 | nibble1);
                     }
                 }
+
                 let mut value_length = rand::thread_rng().gen_range(0..=10);
                 if value_length == 10 {
                     value_length = 8000; // make a long value that is not inlined

core/store/src/trie/mod.rs

@@ -80,7 +80,7 @@ pub struct TrieCosts {
 }
 
 /// Whether a key lookup will be performed through flat storage or through iterating the trie
-#[derive(PartialEq, Eq)]
+#[derive(Clone, Copy, PartialEq, Eq)]
 pub enum KeyLookupMode {
     FlatStorage,
     Trie,
@@ -751,9 +751,16 @@ impl Trie {
 
     #[cfg(test)]
     fn memory_usage_verify(&self, memory: &NodesStorage, handle: NodeHandle) -> u64 {
+        // Cannot compute memory usage naively if given only partial storage.
+        if self.storage.as_partial_storage().is_some() {
+            return 0;
+        }
+        // We don't want to impact recorded storage by retrieving nodes for
+        // this sanity check.
         if self.recorder.is_some() {
             return 0;
         }
+
         let TrieNodeWithSize { node, memory_usage } = match handle {
             NodeHandle::InMemory(h) => memory.node_ref(h).clone(),
             NodeHandle::Hash(h) => self.retrieve_node(&h).expect("storage failure").1,
@@ -1656,6 +1663,7 @@ pub mod estimator {
 #[cfg(test)]
 mod tests {
     use assert_matches::assert_matches;
+    use rand::prelude::SliceRandom;
     use rand::Rng;
 
     use crate::test_utils::{
@@ -1838,7 +1846,7 @@ mod tests {
     fn test_contains_key() {
         let sid = ShardUId::single_shard();
         let bid = CryptoHash::default();
-        let tries = TestTriesBuilder::new().with_flat_storage().build();
+        let tries = TestTriesBuilder::new().with_enabled_flat_storage().build();
         let initial = vec![
             (vec![99, 44, 100, 58, 58, 49], Some(vec![1])),
             (vec![99, 44, 100, 58, 58, 50], Some(vec![1])),

core/store/src/trie/state_parts.rs

@@ -1147,7 +1147,7 @@ mod tests {
     fn get_trie_nodes_for_part_with_flat_storage() {
         let value_len = 1000usize;
 
-        let tries = TestTriesBuilder::new().with_flat_storage().build();
+        let tries = TestTriesBuilder::new().with_enabled_flat_storage().build();
         let shard_uid = ShardUId::single_shard();
         let block_hash = CryptoHash::default();
         let part_id = PartId::new(1, 3);