feat(sync): add unwinding support for chain reorganizations

bin/katana/Cargo.toml

@@ -16,6 +16,7 @@ katana-primitives.workspace = true
 katana-rpc-client.workspace = true
 katana-rpc-types.workspace = true
 katana-utils.workspace = true
+katana-stage.workspace = true
 
 anyhow.workspace = true
 async-trait.workspace = true

bin/katana/src/cli/stage/mod.rs

@@ -1,7 +1,10 @@
 use anyhow::Result;
 use clap::{Args, Subcommand};
 
+use crate::cli::execute_async;
+
 mod checkpoint;
+mod unwind;
 
 #[derive(Debug, Args)]
 #[cfg_attr(test, derive(PartialEq))]
@@ -15,12 +18,15 @@ pub struct StageArgs {
 enum Commands {
     /// Manage stage checkpoints
     Checkpoint(checkpoint::CheckpointArgs),
+    /// Unwind a stage to a previous state
+    Unwind(unwind::UnwindArgs),
 }
 
 impl StageArgs {
     pub fn execute(self) -> Result<()> {
         match self.commands {
             Commands::Checkpoint(args) => args.execute(),
+            Commands::Unwind(args) => execute_async(args.execute())?,
         }
     }
 }

bin/katana/src/cli/stage/unwind.rs

@@ -0,0 +1,38 @@
+use anyhow::Result;
+use clap::Args;
+use katana_primitives::block::BlockNumber;
+use katana_provider::api::stage::StageCheckpointProvider;
+use katana_provider::providers::db::DbProvider;
+use katana_stage::Stage;
+
+use crate::cli::db::open_db_rw;
+
+#[derive(Debug, Args)]
+#[cfg_attr(test, derive(PartialEq))]
+pub struct UnwindArgs {
+    /// The stage ID to unwind
+    #[arg(value_name = "STAGE_ID")]
+    stage_id: String,
+
+    /// The stage ID to unwind to
+    #[arg(value_name = "UNWIND_TO")]
+    unwind_to: BlockNumber,
+
+    /// Path to the database directory.
+    #[arg(short, long)]
+    path: String,
+}
+
+impl UnwindArgs {
+    pub async fn execute(self) -> Result<()> {
+        use katana_stage::StateTrie;
+
+        let provider = DbProvider::new(open_db_rw(&self.path)?);
+        let mut stage = StateTrie::new(&provider);
+
+        stage.unwind(self.unwind_to).await?;
+        provider.set_checkpoint(stage.id(), self.unwind_to)?;
+
+        Ok(())
+    }
+}

crates/core/src/backend/mod.rs

@@ -656,4 +656,12 @@ impl TrieWriter for GenesisTrieWriter {
         trie.commit(block_number);
         Ok(trie.root())
     }
+
+    fn unwind_classes_trie(&self, _: BlockNumber) -> katana_provider::ProviderResult<Felt> {
+        unimplemented!()
+    }
+
+    fn unwind_contracts_trie(&self, _: BlockNumber) -> katana_provider::ProviderResult<Felt> {
+        unimplemented!()
+    }
 }

crates/node/src/full/pool.rs

@@ -2,18 +2,9 @@ use std::future::Future;
 
 use katana_pool::ordering::FiFo;
 use katana_pool::pool::Pool;
-use katana_pool::validation::stateful::TxValidator;
-use katana_pool::PoolTransaction;
 use katana_pool_api::validation::{ValidationOutcome, ValidationResult, Validator};
-use katana_primitives::chain::ChainId;
-use katana_primitives::fee::ResourceBoundsMapping;
-use katana_primitives::transaction::{
-    DeclareTx, DeployAccountTx, ExecutableTxWithHash, InvokeTx, TxHash,
-};
-use katana_primitives::{ContractAddress, Felt};
-use katana_rpc_types::{
-    BroadcastedDeclareTx, BroadcastedDeployAccountTx, BroadcastedInvokeTx, BroadcastedTx,
-};
+use katana_primitives::transaction::ExecutableTxWithHash;
+use katana_rpc_types::BroadcastedTx;
 
 pub type FullNodePool =
     Pool<ExecutableTxWithHash, GatewayProxyValidator, FiFo<ExecutableTxWithHash>>;

crates/node/src/lib.rs

@@ -31,7 +31,7 @@ use katana_metrics::sys::DiskReporter;
 use katana_metrics::{Report, Server as MetricsServer};
 use katana_pool::ordering::FiFo;
 use katana_pool::TxPool;
-use katana_primitives::env::{FeeTokenAddressses, VersionedConstantsOverrides};
+use katana_primitives::env::VersionedConstantsOverrides;
 #[cfg(feature = "cartridge")]
 use katana_rpc::cartridge::CartridgeApi;
 use katana_rpc::cors::Cors;

crates/rpc/rpc/src/starknet/blockifier.rs

@@ -9,7 +9,7 @@ use katana_executor::implementation::blockifier::utils::{self, block_context_fro
 use katana_executor::{ExecutionError, ExecutionFlags, ExecutionResult, ResultAndStates};
 use katana_primitives::env::{BlockEnv, VersionedConstantsOverrides};
 use katana_primitives::transaction::ExecutableTxWithHash;
-use katana_primitives::{chain, Felt};
+use katana_primitives::Felt;
 use katana_provider::api::state::StateProvider;
 use katana_rpc_api::error::starknet::{ContractErrorData, StarknetApiError};
 use katana_rpc_types::{FeeEstimate, FunctionCall};

crates/rpc/rpc/src/starknet/mod.rs

@@ -6,7 +6,7 @@ use std::sync::Arc;
 
 use katana_chain_spec::ChainSpec;
 use katana_core::backend::storage::Database;
-use katana_pool::{TransactionPool, TxPool};
+use katana_pool::TransactionPool;
 use katana_primitives::block::{BlockHashOrNumber, BlockIdOrTag, FinalityStatus};
 use katana_primitives::class::{ClassHash, CompiledClass};
 use katana_primitives::contract::{ContractAddress, Nonce, StorageKey, StorageValue};

crates/rpc/rpc/src/starknet/read.rs

@@ -5,7 +5,6 @@ use std::sync::Arc;
 use anyhow::anyhow;
 use jsonrpsee::core::{async_trait, RpcResult};
 use jsonrpsee::types::ErrorObjectOwned;
-use katana_executor::ExecutorFactory;
 #[cfg(feature = "cartridge")]
 use katana_genesis::allocation::GenesisAccountAlloc;
 use katana_pool::TransactionPool;

crates/rpc/rpc/src/starknet/write.rs

@@ -48,8 +48,8 @@ where
                 .into_inner(this.inner.chain_spec.id())
                 .map_err(|_| StarknetApiError::InvalidContractClass)?;
 
-            let class_hash = tx.class_hash();
-            let tx = ExecutableTxWithHash::new(ExecutableTx::Declare(tx));
+            let _ = tx.class_hash();
+            let _ = ExecutableTxWithHash::new(ExecutableTx::Declare(tx));
             // let transaction_hash = this.inner.pool.add_transaction(tx).await?;
 
             // Ok(AddDeclareTransactionResponse { transaction_hash, class_hash })
@@ -68,9 +68,9 @@ where
             }
 
             let tx = tx.into_inner(this.inner.chain_spec.id());
-            let contract_address = tx.contract_address();
+            let _ = tx.contract_address();
 
-            let tx = ExecutableTxWithHash::new(ExecutableTx::DeployAccount(tx));
+            let _ = ExecutableTxWithHash::new(ExecutableTx::DeployAccount(tx));
             // let transaction_hash = this.inner.pool.add_transaction(tx).await?;
 
             // Ok(AddDeployAccountTransactionResponse { transaction_hash, contract_address })

crates/storage/db/src/trie/mod.rs

@@ -163,7 +163,7 @@ where
     fn create_batch(&self) -> Self::Batch {}
 
     fn remove_by_prefix(&mut self, _: &DatabaseKey<'_>) -> Result<(), Self::DatabaseError> {
-        Ok(())
+        unimplemented!()
     }
 
     fn get(&self, key: &DatabaseKey<'_>) -> Result<Option<ByteVec>, Self::DatabaseError> {
@@ -173,9 +173,22 @@ where
 
     fn get_by_prefix(
         &self,
-        _: &DatabaseKey<'_>,
+        prefix: &DatabaseKey<'_>,
     ) -> Result<Vec<(ByteVec, ByteVec)>, Self::DatabaseError> {
-        todo!()
+        let mut results = Vec::new();
+
+        let mut cursor = self.tx.cursor::<Tb>()?;
+        let walker = cursor.walk(None)?;
+
+        for entry in walker {
+            let (TrieDatabaseKey { key, .. }, value) = entry?;
+
+            if key.starts_with(prefix.as_slice()) {
+                results.push((key.to_smallvec(), value));
+            }
+        }
+
+        Ok(results)
     }
 
     fn insert(
@@ -300,8 +313,7 @@ where
         &self,
         prefix: &DatabaseKey<'_>,
     ) -> Result<Vec<(ByteVec, ByteVec)>, Self::DatabaseError> {
-        let _ = prefix;
-        todo!()
+        TrieDb::<Tb, Tx>::new(self.tx.clone()).get_by_prefix(prefix)
     }
 
     fn insert(
@@ -509,4 +521,67 @@ mod tests {
             assert_eq!(vec![value0, value1], result);
         }
     }
+
+    #[test]
+    fn revert_to() {
+        let db = test_utils::create_test_db();
+        let db_tx = db.tx_mut().expect("failed to get tx");
+
+        let mut trie = ClassesTrie::new(TrieDbMut::<tables::ClassesTrie, _>::new(&db_tx));
+
+        // Insert values at block 0
+        trie.insert(felt!("0x1"), felt!("0x100"));
+        trie.insert(felt!("0x2"), felt!("0x200"));
+        trie.commit(0);
+        let root_at_block_0 = trie.root();
+
+        // Insert more values at block 1
+        trie.insert(felt!("0x3"), felt!("0x300"));
+        trie.insert(felt!("0x4"), felt!("0x400"));
+        trie.commit(1);
+        let root_at_block_1 = trie.root();
+
+        // Roots should be different
+        assert_ne!(root_at_block_0, root_at_block_1);
+
+        // Insert even more values at block 2
+        trie.insert(felt!("0x5"), felt!("0x500"));
+        trie.commit(2);
+        let root_at_block_2 = trie.root();
+
+        // Roots should be different
+        assert_ne!(root_at_block_1, root_at_block_2);
+        assert_ne!(root_at_block_0, root_at_block_2);
+
+        // Revert to block 1
+        trie.revert_to(1, 2);
+        let root_after_revert = trie.root();
+
+        // After revert, root should match block 1
+        assert_eq!(root_after_revert, root_at_block_1);
+
+        // Revert to block 0
+        trie.revert_to(0, 1);
+        let root_after_second_revert = trie.root();
+
+        // After revert, root should match block 0
+        assert_eq!(root_after_second_revert, root_at_block_0);
+
+        // Insert more values at block 1
+        trie.insert(felt!("0x3"), felt!("0x300"));
+        trie.insert(felt!("0x4"), felt!("0x400"));
+        trie.commit(1);
+        let root_at_block_1_after_insert = trie.root();
+
+        // After insertion, root should match block 1
+        assert_eq!(root_at_block_1_after_insert, root_at_block_1);
+
+        // Insert even more values at block 2
+        trie.insert(felt!("0x5"), felt!("0x500"));
+        trie.commit(2);
+        let root_at_block_2_after_insert = trie.root();
+
+        // After insertion, root should match block 2
+        assert_eq!(root_at_block_2_after_insert, root_at_block_2);
+    }
 }

crates/storage/provider/provider-api/src/trie.rs

@@ -20,4 +20,8 @@ pub trait TrieWriter: Send + Sync {
         block_number: BlockNumber,
         state_updates: &StateUpdates,
     ) -> ProviderResult<Felt>;
+
+    fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult<Felt>;
+
+    fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult<Felt>;
 }

crates/storage/provider/provider/src/lib.rs

@@ -375,6 +375,14 @@ where
     ) -> ProviderResult<Felt> {
         self.provider.trie_insert_contract_updates(block_number, state_updates)
     }
+
+    fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult<Felt> {
+        self.provider.unwind_classes_trie(unwind_to)
+    }
+
+    fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult<Felt> {
+        self.provider.unwind_contracts_trie(unwind_to)
+    }
 }
 
 impl<Db> StageCheckpointProvider for BlockchainProvider<Db>