Merge pull request #1854 from eqlabs/krisztian/reverse-state-update

feat: revert Merkle trie changes upon reorg

crates/common/src/lib.rs

@@ -554,6 +554,10 @@ impl ContractAddress {
 
         ContractAddress::new_or_panic(contract_address)
     }
+
+    pub fn is_system_contract(&self) -> bool {
+        *self == ContractAddress::ONE
+    }
 }
 
 #[derive(Clone, Debug, PartialEq)]

crates/common/src/state_update.rs

@@ -443,6 +443,21 @@ impl StateUpdate {
     }
 }
 
+#[derive(Debug, PartialEq)]
+pub enum ReverseContractUpdate {
+    Deleted,
+    Updated(ContractUpdate),
+}
+
+impl ReverseContractUpdate {
+    pub fn update_mut(&mut self) -> Option<&mut ContractUpdate> {
+        match self {
+            Self::Deleted => None,
+            Self::Updated(update) => Some(update),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/merkle-tree/src/contract.rs

@@ -165,6 +165,12 @@ impl<'tx> StorageCommitmentTree<'tx> {
         self.tree.set(&self.storage, key, value.0)
     }
 
+    pub fn get(&self, address: &ContractAddress) -> anyhow::Result<Option<ContractStateHash>> {
+        let key = address.view_bits().to_owned();
+        let value = self.tree.get(&self.storage, key)?;
+        Ok(value.map(ContractStateHash))
+    }
+
     /// 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>)> {

crates/merkle-tree/src/contract_state.rs

@@ -3,8 +3,8 @@ use std::collections::HashMap;
 use crate::ContractsStorageTree;
 use anyhow::Context;
 use pathfinder_common::{
-    BlockNumber, ClassHash, ContractAddress, ContractNonce, ContractRoot, ContractStateHash,
-    StorageAddress, StorageValue,
+    state_update::ReverseContractUpdate, BlockNumber, ClassHash, ContractAddress, ContractNonce,
+    ContractRoot, ContractStateHash, StorageAddress, StorageValue,
 };
 use pathfinder_crypto::{hash::pedersen_hash, Felt};
 use pathfinder_storage::{Node, Transaction};
@@ -86,7 +86,7 @@ pub fn update_contract_state(
         (current_root, Default::default())
     };
 
-    let class_hash = if contract_address == ContractAddress::ONE {
+    let class_hash = if contract_address.is_system_contract() {
         // This is a special system contract at address 0x1, which doesn't have a class hash.
         ClassHash::ZERO
     } else if let Some(class_hash) = new_class_hash {
@@ -137,6 +137,93 @@ pub fn calculate_contract_state_hash(
     ContractStateHash(hash)
 }
 
+/// Reverts Merkle tree state for a contract.
+///
+/// Takes Merkle tree state at `head` and applies reverse updates.
+pub fn revert_contract_state(
+    transaction: &Transaction<'_>,
+    contract_address: ContractAddress,
+    head: BlockNumber,
+    target_block: BlockNumber,
+    contract_update: ReverseContractUpdate,
+) -> anyhow::Result<ContractStateHash> {
+    tracing::debug!(%contract_address, "Rolling back");
+
+    match contract_update {
+        ReverseContractUpdate::Deleted => {
+            tracing::debug!(%contract_address, "Contract has been deleted");
+            Ok(ContractStateHash::ZERO)
+        }
+        ReverseContractUpdate::Updated(update) => {
+            let class_hash = match update.class {
+                Some(class_hash) => class_hash.class_hash(),
+                None => {
+                    if contract_address.is_system_contract() {
+                        // system contracts have no class hash
+                        ClassHash::ZERO
+                    } else {
+                        transaction
+                            .contract_class_hash(target_block.into(), contract_address)?
+                            .unwrap()
+                    }
+                }
+            };
+
+            let nonce = match update.nonce {
+                Some(nonce) => nonce,
+                None => transaction
+                    .contract_nonce(contract_address, target_block.into())
+                    .context("Getting contract nonce")?
+                    .unwrap_or_default(),
+            };
+
+            // Apply storage updates
+            let root = if !update.storage.is_empty() {
+                let mut tree = ContractsStorageTree::load(transaction, contract_address, head)
+                    .context("Loading contract state")?;
+
+                for (address, value) in update.storage {
+                    tree.set(address, value)
+                        .context("Updating contract state")?;
+                }
+
+                let (root, nodes_added) = tree.commit().context("Committing contract state")?;
+
+                let root_index = if !root.0.is_zero() && !nodes_added.is_empty() {
+                    let root_index = transaction
+                        .insert_contract_trie(root, &nodes_added)
+                        .context("Persisting contract trie")?;
+                    Some(root_index)
+                } else {
+                    None
+                };
+
+                transaction
+                    .insert_or_update_contract_root(target_block, contract_address, root_index)
+                    .context("Inserting contract's root index")?;
+
+                root
+            } else {
+                transaction
+                    .contract_root(head, contract_address)?
+                    .context("Fetching current contract root")?
+            };
+
+            let state_hash = if contract_address.is_system_contract() && root == ContractRoot::ZERO
+            {
+                // special case: if the contract trie is empty the system contract should be deleted
+                ContractStateHash::ZERO
+            } else {
+                calculate_contract_state_hash(class_hash, root, nonce)
+            };
+
+            tracing::debug!(%state_hash, %contract_address, "Contract state rolled back");
+
+            Ok(state_hash)
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::calculate_contract_state_hash;

crates/merkle-tree/src/tree.rs

@@ -437,8 +437,11 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
     }
 
     /// Returns the value stored at key, or `None` if it does not exist.
-    #[cfg(test)]
-    fn get(&self, storage: &impl Storage, key: BitVec<u8, Msb0>) -> anyhow::Result<Option<Felt>> {
+    pub fn get(
+        &self,
+        storage: &impl Storage,
+        key: BitVec<u8, Msb0>,
+    ) -> anyhow::Result<Option<Felt>> {
         let node = self.traverse(storage, &key)?;
         let node = node.last();
 

crates/pathfinder/examples/test_state_rollback.rs

@@ -0,0 +1,59 @@
+use std::{num::NonZeroU32, time::Instant};
+
+use anyhow::Context;
+use pathfinder_common::BlockNumber;
+use pathfinder_storage::{BlockId, JournalMode, Storage};
+
+fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .compact()
+        .init();
+
+    let database_path = std::env::args().nth(1).unwrap();
+    let storage = Storage::migrate(database_path.into(), JournalMode::WAL, 1)?
+        .create_pool(NonZeroU32::new(10).unwrap())?;
+    let mut db = storage
+        .connection()
+        .context("Opening database connection")?;
+
+    let latest_block = {
+        let tx = db.transaction().unwrap();
+        let (latest_block, _) = tx.block_id(BlockId::Latest)?.unwrap();
+        latest_block.get()
+    };
+    let from: u64 = std::env::args()
+        .nth(2)
+        .map(|s| str::parse(&s).unwrap())
+        .unwrap();
+    let to: u64 = std::env::args()
+        .nth(3)
+        .map(|s| str::parse(&s).unwrap())
+        .unwrap();
+    assert!(from <= latest_block);
+    assert!(from > to);
+    let from = BlockNumber::new_or_panic(from);
+    let to = BlockNumber::new_or_panic(to);
+
+    tracing::info!(%from, %to, "Testing state rollback");
+
+    let started = Instant::now();
+
+    let tx = db.transaction()?;
+
+    let to_header = tx.block_header(to.into()).unwrap().unwrap();
+
+    pathfinder_lib::state::revert::revert_starknet_state(&tx, from, to, to_header, true)?;
+
+    tracing::info!(
+        from=%from,
+        to=%to,
+        total=?started.elapsed(),
+        "Finished state rollback"
+    );
+
+    // Explicitly do _not_ commit transaction
+    drop(tx);
+
+    Ok(())
+}

crates/pathfinder/src/state.rs

@@ -1,4 +1,4 @@
 pub mod block_hash;
 mod sync;
 
-pub use sync::{l1, l2, sync, Gossiper, SyncContext};
+pub use sync::{l1, l2, revert, sync, Gossiper, SyncContext};

crates/pathfinder/src/state/sync.rs

@@ -2,6 +2,7 @@ mod class;
 pub mod l1;
 pub mod l2;
 mod pending;
+pub mod revert;
 
 use anyhow::Context;
 use pathfinder_common::prelude::*;
@@ -974,6 +975,17 @@ async fn l2_reorg(connection: &mut Connection, reorg_tail: BlockNumber) -> anyho
             .increment_reorg_counter()
             .context("Incrementing reorg counter")?;
 
+        // Roll back Merkle trie updates.
+        //
+        // If we're rolling back genesis then there will be no blocks left so state will be empty.
+        if let Some(target_block) = reorg_tail.parent() {
+            let target_header = transaction
+                .block_header(target_block.into())
+                .context("Fetching target block header")?
+                .context("Expected target header to exist")?;
+            revert::revert_starknet_state(&transaction, head, target_block, target_header, false)?;
+        }
+
         // Purge each block one at a time.
         //
         // This is done 1-by-1 to allow sending the reorg'd block data