Merge pull request #1893 from eqlabs/krisztian/merkle-tree-track-remo…

…ved-nodes

feat(merkle-tree): track removed nodes

crates/merkle-tree/src/class.rs

@@ -1,11 +1,9 @@
-use std::collections::HashMap;
-
 use anyhow::Context;
 use pathfinder_common::{
     BlockNumber, ClassCommitment, ClassCommitmentLeafHash, ClassHash, SierraHash,
 };
 use pathfinder_crypto::Felt;
-use pathfinder_storage::{Node, Transaction};
+use pathfinder_storage::{Transaction, TrieUpdate};
 
 use crate::tree::MerkleTree;
 use pathfinder_common::hash::PoseidonHash;
@@ -62,11 +60,11 @@ impl<'tx> ClassCommitmentTree<'tx> {
 
     /// Commits the changes and calculates the new node hashes. Returns the new commitment and
     /// any potentially newly created nodes.
-    pub fn commit(self) -> anyhow::Result<(ClassCommitment, HashMap<Felt, Node>)> {
+    pub fn commit(self) -> anyhow::Result<(ClassCommitment, TrieUpdate)> {
         let update = self.tree.commit(&self.storage)?;
 
-        let commitment = ClassCommitment(update.root);
-        Ok((commitment, update.nodes))
+        let commitment = ClassCommitment(update.root_hash());
+        Ok((commitment, update))
     }
 }
 

crates/merkle-tree/src/contract.rs

@@ -16,8 +16,7 @@ use pathfinder_common::{
     StorageCommitment, StorageValue,
 };
 use pathfinder_crypto::Felt;
-use pathfinder_storage::{Node, Transaction};
-use std::collections::HashMap;
+use pathfinder_storage::{Transaction, TrieUpdate};
 use std::ops::ControlFlow;
 
 /// A [Patricia Merkle tree](MerkleTree) used to calculate commitments to a Starknet contract's storage.
@@ -100,10 +99,10 @@ impl<'tx> ContractsStorageTree<'tx> {
 
     /// Commits the changes and calculates the new node hashes. Returns the new commitment and
     /// any potentially newly created nodes.
-    pub fn commit(self) -> anyhow::Result<(ContractRoot, HashMap<Felt, Node>)> {
+    pub fn commit(self) -> anyhow::Result<(ContractRoot, TrieUpdate)> {
         let update = self.tree.commit(&self.storage)?;
-        let commitment = ContractRoot(update.root);
-        Ok((commitment, update.nodes))
+        let commitment = ContractRoot(update.root_hash());
+        Ok((commitment, update))
     }
 
     /// See [`MerkleTree::dfs`]
@@ -173,10 +172,10 @@ impl<'tx> StorageCommitmentTree<'tx> {
 
     /// Commits the changes and calculates the new node hashes. Returns the new commitment and
     /// any potentially newly created nodes.
-    pub fn commit(self) -> anyhow::Result<(StorageCommitment, HashMap<Felt, Node>)> {
+    pub fn commit(self) -> anyhow::Result<(StorageCommitment, TrieUpdate)> {
         let update = self.tree.commit(&self.storage)?;
-        let commitment = StorageCommitment(update.root);
-        Ok((commitment, update.nodes))
+        let commitment = StorageCommitment(update.root_hash());
+        Ok((commitment, update))
     }
 
     /// Generates a proof for the given `key`. See [`MerkleTree::get_proof`].

crates/merkle-tree/src/contract_state.rs

@@ -7,16 +7,15 @@ use pathfinder_common::{
     ContractRoot, ContractStateHash, StorageAddress, StorageValue,
 };
 use pathfinder_crypto::{hash::pedersen_hash, Felt};
-use pathfinder_storage::{Node, Transaction};
+use pathfinder_storage::{Transaction, TrieUpdate};
 
 #[derive(Debug)]
 pub struct ContractStateUpdateResult {
     pub state_hash: ContractStateHash,
     pub contract_address: ContractAddress,
     root: ContractRoot,
     did_storage_updates: bool,
-    // trie nodes to be inserted into the database
-    nodes: HashMap<Felt, Node>,
+    trie_update: TrieUpdate,
 }
 
 impl ContractStateUpdateResult {
@@ -27,9 +26,9 @@ impl ContractStateUpdateResult {
     pub fn insert(self, block: BlockNumber, transaction: &Transaction<'_>) -> anyhow::Result<()> {
         // Insert nodes only if we made storage updates.
         if self.did_storage_updates {
-            let root_index = if !self.root.0.is_zero() && !self.nodes.is_empty() {
+            let root_index = if !self.root.0.is_zero() && !self.trie_update.nodes_added.is_empty() {
                 let root_index = transaction
-                    .insert_contract_trie(self.root, &self.nodes)
+                    .insert_contract_trie(&self.trie_update)
                     .context("Persisting contract trie")?;
                 Some(root_index)
             } else {
@@ -58,7 +57,7 @@ pub fn update_contract_state(
     block: BlockNumber,
 ) -> anyhow::Result<ContractStateUpdateResult> {
     // Load the contract tree and insert the updates.
-    let (new_root, nodes) = if !updates.is_empty() {
+    let (new_root, trie_update) = if !updates.is_empty() {
         let mut contract_tree = match block.parent() {
             Some(parent) => ContractsStorageTree::load(transaction, contract_address, parent)
                 .context("Loading contract storage tree")?
@@ -72,11 +71,11 @@ pub fn update_contract_state(
                 .set(*key, *value)
                 .context("Update contract storage tree")?;
         }
-        let (contract_root, nodes) = contract_tree
+        let (contract_root, trie_update) = contract_tree
             .commit()
             .context("Apply contract storage tree changes")?;
 
-        (contract_root, nodes)
+        (contract_root, trie_update)
     } else {
         let current_root = transaction
             .contract_root(block, contract_address)
@@ -115,7 +114,7 @@ pub fn update_contract_state(
         state_hash,
         root: new_root,
         did_storage_updates: !updates.is_empty(),
-        nodes,
+        trie_update,
     })
 }
 
@@ -187,11 +186,11 @@ pub fn revert_contract_state(
                         .context("Updating contract state")?;
                 }
 
-                let (root, nodes_added) = tree.commit().context("Committing contract state")?;
+                let (root, trie_update) = tree.commit().context("Committing contract state")?;
 
-                let root_index = if !root.0.is_zero() && !nodes_added.is_empty() {
+                let root_index = if !root.0.is_zero() && !trie_update.nodes_added.is_empty() {
                     let root_index = transaction
-                        .insert_contract_trie(root, &nodes_added)
+                        .insert_contract_trie(&trie_update)
                         .context("Persisting contract trie")?;
                     Some(root_index)
                 } else {

crates/merkle-tree/src/merkle_node.rs

@@ -16,7 +16,7 @@ use pathfinder_common::hash::FeltHash;
 pub enum InternalNode {
     /// A node that has not been fetched from storage yet.
     ///
-    /// As such, all we know is its hash.
+    /// As such, all we know is its index.
     Unresolved(u64),
     /// A branch node with exactly two children.
     Binary(BinaryNode),
@@ -29,6 +29,8 @@ pub enum InternalNode {
 /// Describes the [InternalNode::Binary] variant.
 #[derive(Clone, Debug, PartialEq)]
 pub struct BinaryNode {
+    /// The storage index of this node (if it was loaded from storage).
+    pub storage_index: Option<u64>,
     /// The height of this node in the tree.
     pub height: usize,
     /// [Left](Direction::Left) child.
@@ -39,6 +41,8 @@ pub struct BinaryNode {
 
 #[derive(Clone, Debug, PartialEq)]
 pub struct EdgeNode {
+    /// The storage index of this node (if it was loaded from storage).
+    pub storage_index: Option<u64>,
     /// The starting height of this node in the tree.
     pub height: usize,
     /// The path this edge takes.
@@ -137,6 +141,15 @@ impl InternalNode {
     pub fn is_leaf(&self) -> bool {
         matches!(self, InternalNode::Leaf)
     }
+
+    pub fn storage_index(&self) -> Option<u64> {
+        match self {
+            InternalNode::Unresolved(storage_index) => Some(*storage_index),
+            InternalNode::Binary(binary) => binary.storage_index,
+            InternalNode::Edge(edge) => edge.storage_index,
+            InternalNode::Leaf => None,
+        }
+    }
 }
 
 impl EdgeNode {
@@ -209,6 +222,7 @@ mod tests {
         #[test]
         fn direction() {
             let uut = BinaryNode {
+                storage_index: None,
                 height: 1,
                 left: Rc::new(RefCell::new(InternalNode::Unresolved(1))),
                 right: Rc::new(RefCell::new(InternalNode::Unresolved(2))),
@@ -233,6 +247,7 @@ mod tests {
             let right = Rc::new(RefCell::new(InternalNode::Unresolved(2)));
 
             let uut = BinaryNode {
+                storage_index: None,
                 height: 1,
                 left: left.clone(),
                 right: right.clone(),
@@ -296,6 +311,7 @@ mod tests {
                 let child = Rc::new(RefCell::new(InternalNode::Unresolved(1)));
 
                 let uut = EdgeNode {
+                    storage_index: None,
                     height: 0,
                     path: key.view_bits().to_bitvec(),
                     child,
@@ -312,6 +328,7 @@ mod tests {
                 let path = key.view_bits()[..45].to_bitvec();
 
                 let uut = EdgeNode {
+                    storage_index: None,
                     height: 0,
                     path,
                     child,
@@ -328,6 +345,7 @@ mod tests {
                 let path = key.view_bits()[50..].to_bitvec();
 
                 let uut = EdgeNode {
+                    storage_index: None,
                     height: 50,
                     path,
                     child,
@@ -344,6 +362,7 @@ mod tests {
                 let path = key.view_bits()[230..235].to_bitvec();
 
                 let uut = EdgeNode {
+                    storage_index: None,
                     height: 230,
                     path,
                     child,

crates/merkle-tree/src/transaction.rs

@@ -51,7 +51,9 @@ impl TransactionOrEventTree {
     }
 
     pub fn commit(self) -> anyhow::Result<Felt> {
-        self.tree.commit(&NullStorage {}).map(|update| update.root)
+        self.tree
+            .commit(&NullStorage {})
+            .map(|update| update.root_hash())
     }
 }