Fix revert for trie diffs to include incomplete part as well

blockchain/src/blockchain/rebranch_utils.rs

@@ -200,7 +200,7 @@ impl Blockchain {
                 Err(e) => {
                     // If a block fails to apply here it does not verify fully.
                     // This block and all blocks after this thus should be removed from the store
-                    // as they are not verifying. 
+                    // as they are not verifying.
                     warn!(
                         block = %block.1.head,
                         reason = "failed to apply fork block while rebranching",

blockchain/tests/push_with_chunks.rs

@@ -1,3 +1,4 @@
+use nimiq_account::RevertInfo;
 use nimiq_blockchain_interface::{
     AbstractBlockchain, ChunksPushError, ChunksPushResult, PushResult,
 };
@@ -316,11 +317,26 @@ fn can_rebranch_and_revert_chunks() {
         .unwrap();
     let chunk2b = temp_producer1.get_chunk(chunk2_start, 3);
 
+    let block2a_number = block2a.block_number();
     assert_eq!(
-        temp_producer2.push_with_chunks(block2a, diff2a, vec![chunk2a]),
+        temp_producer2.push_with_chunks(block2a, diff2a.clone(), vec![chunk2a]),
         Ok((PushResult::Extended, Ok(ChunksPushResult::Chunks(1, 0))))
     );
 
+    {
+        let rev_info = temp_producer2
+            .blockchain
+            .read()
+            .chain_store
+            .get_revert_info(block2a_number, None);
+        match rev_info {
+            Some(RevertInfo::Diff(diff)) => {
+                assert_eq!(diff.0.len(), diff2a.0.len());
+            }
+            _ => panic!("Wrong revert info"),
+        }
+    }
+
     assert_eq!(
         temp_producer2.push_with_chunks(block2b, diff2b, vec![chunk2b]),
         Ok((PushResult::Rebranched, Ok(ChunksPushResult::Chunks(1, 0))))

primitives/account/src/accounts.rs

@@ -10,7 +10,7 @@ use nimiq_primitives::{
     trie::{
         error::IncompleteTrie,
         trie_chunk::{TrieChunk, TrieChunkPushResult},
-        trie_diff::TrieDiff,
+        trie_diff::{RevertTrieDiff, TrieDiff},
         trie_proof::TrieProof,
         TrieItem,
     },
@@ -307,7 +307,7 @@ impl Accounts {
         &self,
         txn: &mut WriteTransactionProxy,
         diff: TrieDiff,
-    ) -> Result<TrieDiff, AccountError> {
+    ) -> Result<RevertTrieDiff, AccountError> {
         let diff = self.tree.apply_diff(txn, diff)?;
         self.tree.update_root(txn).ok();
         Ok(diff)
@@ -569,9 +569,9 @@ impl Accounts {
     pub fn revert_diff(
         &self,
         txn: &mut WriteTransactionProxy,
-        diff: TrieDiff,
+        diff: RevertTrieDiff,
     ) -> Result<(), AccountError> {
-        self.tree.apply_diff(txn, diff)?;
+        self.tree.revert_diff(txn, diff)?;
         Ok(())
     }
 

primitives/account/src/receipts.rs

@@ -1,7 +1,7 @@
 use std::{fmt::Debug, io};
 
 use nimiq_database_value::{FromDatabaseValue, IntoDatabaseValue};
-use nimiq_primitives::{account::FailReason, trie::trie_diff::TrieDiff};
+use nimiq_primitives::{account::FailReason, trie::trie_diff::RevertTrieDiff};
 use nimiq_serde::{Deserialize, Serialize};
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -71,7 +71,7 @@ pub struct Receipts {
 #[repr(u8)]
 pub enum RevertInfo {
     Receipts(Receipts),
-    Diff(TrieDiff),
+    Diff(RevertTrieDiff),
 }
 
 impl From<Receipts> for RevertInfo {
@@ -80,8 +80,8 @@ impl From<Receipts> for RevertInfo {
     }
 }
 
-impl From<TrieDiff> for RevertInfo {
-    fn from(diff: TrieDiff) -> RevertInfo {
+impl From<RevertTrieDiff> for RevertInfo {
+    fn from(diff: RevertTrieDiff) -> RevertInfo {
         RevertInfo::Diff(diff)
     }
 }

primitives/src/trie/trie_diff.rs

@@ -112,3 +112,57 @@ impl FromDatabaseValue for TrieDiff {
             .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
     }
 }
+
+/// A diff on an incomplete trie.
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+#[repr(u8)]
+pub enum RevertDiffValue {
+    /// Puts a value into the known part of the trie.
+    Put(Vec<u8>),
+    /// Removes a value from either the known or the unknown part of the trie.
+    Remove,
+    /// Updates a value in the unknown part of the trie.
+    UpdateStump,
+}
+
+impl RevertDiffValue {
+    /// Create an incomplete diff value for a value from the known part of the trie.
+    pub fn known_value(old_value: Option<Vec<u8>>) -> Self {
+        match old_value {
+            Some(v) => RevertDiffValue::Put(v),
+            None => RevertDiffValue::Remove,
+        }
+    }
+
+    /// Create an incomplete diff value for the unknown part of the trie.
+    pub fn unknown_value(new_stump: bool) -> Self {
+        match new_stump {
+            true => RevertDiffValue::Remove,
+            false => RevertDiffValue::UpdateStump,
+        }
+    }
+}
+
+/// A diff on an incomplete trie.
+#[derive(Clone, Debug, Default, Serialize, Deserialize)]
+pub struct RevertTrieDiff(pub BTreeMap<KeyNibbles, RevertDiffValue>);
+
+impl IntoDatabaseValue for RevertTrieDiff {
+    fn database_byte_size(&self) -> usize {
+        self.serialized_size()
+    }
+
+    fn copy_into_database(&self, mut bytes: &mut [u8]) {
+        Serialize::serialize(&self, &mut bytes).unwrap();
+    }
+}
+
+impl FromDatabaseValue for RevertTrieDiff {
+    fn copy_from_database(bytes: &[u8]) -> io::Result<Self>
+    where
+        Self: Sized,
+    {
+        Deserialize::deserialize_from_vec(bytes)
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
+    }
+}

primitives/trie/src/trie.rs

@@ -16,7 +16,7 @@ use nimiq_primitives::{
     trie::{
         error::{IncompleteTrie, MerkleRadixTrieError},
         trie_chunk::{TrieChunk, TrieChunkPushResult, TrieItem},
-        trie_diff::TrieDiff,
+        trie_diff::{RevertDiffValue, RevertTrieDiff, TrieDiff},
         trie_node::{RootData, TrieNode, TrieNodeKind},
         trie_proof::TrieProof,
         trie_proof_node::TrieProofNode,
@@ -144,6 +144,41 @@ impl MerkleRadixTrie {
         self.init_root(txn, true);
     }
 
+    /// Prints a human friendly version of the subtrie for debugging.
+    /// Not to be used on large trees!
+    pub fn debug_print(&self, txn: &TransactionProxy) {
+        let missing_range = self.get_missing_range(txn);
+        println!("-> ROOT");
+        self.debug_print_subtrie(txn, &KeyNibbles::ROOT, &missing_range, 2)
+    }
+
+    /// Prints a human friendly version of the subtrie for debugging.
+    /// Not to be used on large trees!
+    pub fn debug_print_subtrie(
+        &self,
+        txn: &TransactionProxy,
+        key: &KeyNibbles,
+        missing_range: &Option<ops::RangeFrom<KeyNibbles>>,
+        depth: usize,
+    ) {
+        let prefix = "  ".repeat(depth);
+        let node = if let Some(node) = self.get_node(txn, key) {
+            node
+        } else {
+            println!("{}-> MISSING", prefix);
+            return;
+        };
+
+        for child in node.iter_children() {
+            println!("{}-> {} ({})", prefix, child.suffix, child.has_hash());
+            if let Ok(subkey) = child.key(key, missing_range) {
+                self.debug_print_subtrie(txn, &subkey, missing_range, depth + 2);
+            } else {
+                println!("{}  -> MISSING", prefix);
+            }
+        }
+    }
+
     /// Returns the root hash of the Merkle Radix Trie.
     pub fn root_hash_assert(&self, txn: &TransactionProxy) -> Blake2bHash {
         let root = self.get_root(txn).unwrap();
@@ -433,23 +468,27 @@ impl MerkleRadixTrie {
         key: &KeyNibbles,
     ) -> Result<(), MerkleRadixTrieError> {
         let missing_range = self.get_missing_range(txn);
-        self.update_within_missing_part_raw(txn, key, &missing_range)
+        self.update_within_missing_part_raw(txn, key, &missing_range)?;
+        Ok(())
     }
 
     /// Resets the hashes for a path to a node in the missing part of the tree.
     /// This is used to indicate that the corresponding branch's hash has changed,
     /// but we cannot provide the accurate hash until a new chunk has been pushed.
+    ///
+    /// Returns `true` if a new stump has been added, otherwise `false`.
     fn update_within_missing_part_raw(
         &self,
         txn: &mut WriteTransactionProxy,
         key: &KeyNibbles,
         missing_range: &Option<ops::RangeFrom<KeyNibbles>>,
-    ) -> Result<(), MerkleRadixTrieError> {
+    ) -> Result<bool, MerkleRadixTrieError> {
         let mut node = self
             .get_root(txn)
             .expect("The Merkle Radix Trie didn't have a root node!");
 
         let mut root_path = Vec::new();
+        let mut new_stump = false;
         // Descend down the tree and collect nodes to be updated.
         loop {
             // This function should only be called with keys in the missing range.
@@ -469,13 +508,72 @@ impl MerkleRadixTrie {
                     let parent = mem::replace(&mut node, child);
                     root_path.push(parent);
                 }
-                Err(MerkleRadixTrieError::ChildIsStump)
-                | Err(MerkleRadixTrieError::ChildDoesNotExist) => {
+                e @ Err(MerkleRadixTrieError::ChildIsStump)
+                | e @ Err(MerkleRadixTrieError::ChildDoesNotExist) => {
                     // We cannot continue further down as the next node is outside our trie.
                     // Update this node prior to `update_keys` (which only updates the other nodes in the root path).
                     node.put_child_no_hash(key).expect("Prefix must be correct");
                     self.put_node(txn, &node, OldValue::Unchanged);
 
+                    // Mark that a new stump has been added.
+                    if let Err(MerkleRadixTrieError::ChildDoesNotExist) = e {
+                        new_stump = true;
+                    }
+
+                    root_path.push(node);
+                    break;
+                }
+                Err(e) => {
+                    return Err(e);
+                }
+            }
+        }
+
+        self.update_keys(txn, root_path, CountUpdates::default());
+
+        Ok(new_stump)
+    }
+
+    /// Removes the corresponding stump and updates the hashes for a path
+    /// to a node in the missing part of the tree.
+    /// This is used to indicate that the corresponding branch's hash has changed,
+    /// but we cannot provide the accurate hash until a new chunk has been pushed.
+    fn remove_within_missing_part_raw(
+        &self,
+        txn: &mut WriteTransactionProxy,
+        key: &KeyNibbles,
+        missing_range: &Option<ops::RangeFrom<KeyNibbles>>,
+    ) -> Result<(), MerkleRadixTrieError> {
+        let mut node = self
+            .get_root(txn)
+            .expect("The Merkle Radix Trie didn't have a root node!");
+
+        let mut root_path = Vec::new();
+        // Descend down the tree and collect nodes to be updated.
+        loop {
+            // This function should only be called with keys in the missing range.
+            assert_ne!(&node.key, key);
+
+            // Check that the key we are trying to update can exist in this part of the tree.
+            if !node.key.is_prefix_of(key) {
+                return Err(MerkleRadixTrieError::ChildDoesNotExist);
+            }
+
+            // Descend to the corresponding child.
+            match node.child_key(key, missing_range) {
+                Ok(child_key) => {
+                    let child = self
+                        .get_node(txn, &child_key)
+                        .expect("Child should be present");
+                    let parent = mem::replace(&mut node, child);
+                    root_path.push(parent);
+                }
+                Err(MerkleRadixTrieError::ChildIsStump) => {
+                    // We cannot continue further down as the next node is outside our trie.
+                    // Remove this stump.
+                    node.remove_child(key).expect("Prefix must be correct");
+                    self.put_node(txn, &node, OldValue::Unchanged);
+
                     root_path.push(node);
                     break;
                 }
@@ -1196,7 +1294,7 @@ impl MerkleRadixTrie {
         &self,
         txn: &mut WriteTransactionProxy,
         diff: TrieDiff,
-    ) -> Result<TrieDiff, MerkleRadixTrieError> {
+    ) -> Result<RevertTrieDiff, MerkleRadixTrieError> {
         let missing_range = self.get_missing_range(txn);
         let mut result = BTreeMap::default();
         for (key, value) in diff.0 {
@@ -1207,12 +1305,47 @@ impl MerkleRadixTrie {
                 } else {
                     self.remove_raw(txn, &key, &missing_range);
                 };
-                assert!(result.insert(key, old_value).is_none());
+                assert!(result
+                    .insert(key, RevertDiffValue::known_value(old_value))
+                    .is_none());
             } else {
-                self.update_within_missing_part_raw(txn, &key, &missing_range)?;
+                let new_stump = self.update_within_missing_part_raw(txn, &key, &missing_range)?;
+                assert!(result
+                    .insert(key, RevertDiffValue::unknown_value(new_stump))
+                    .is_none());
             }
         }
-        Ok(TrieDiff(result))
+        Ok(RevertTrieDiff(result))
+    }
+
+    pub fn revert_diff(
+        &self,
+        txn: &mut WriteTransactionProxy,
+        diff: RevertTrieDiff,
+    ) -> Result<(), MerkleRadixTrieError> {
+        let missing_range = self.get_missing_range(txn);
+        for (key, value) in diff.0.into_iter().rev() {
+            if self.is_within_complete_part(&key, &missing_range) {
+                match value {
+                    RevertDiffValue::Put(value) => self.put_raw(txn, &key, value, &missing_range),
+                    RevertDiffValue::Remove => self.remove_raw(txn, &key, &missing_range),
+                    RevertDiffValue::UpdateStump => {
+                        unreachable!("key should be in complete part")
+                    }
+                }
+            } else {
+                match value {
+                    RevertDiffValue::Put(_) => unreachable!("key should be in incomplete part"),
+                    RevertDiffValue::Remove => {
+                        self.remove_within_missing_part_raw(txn, &key, &missing_range)?
+                    }
+                    RevertDiffValue::UpdateStump => {
+                        self.update_within_missing_part_raw(txn, &key, &missing_range)?;
+                    }
+                }
+            }
+        }
+        Ok(())
     }
 
     fn is_within_complete_part(