feat(katana): db storage intergration (#1440)

crates/katana/core/Cargo.toml

@@ -7,6 +7,7 @@ repository.workspace = true
 version.workspace = true
 
 [dependencies]
+katana-db = { path = "../storage/db" }
 katana-executor = { path = "../executor" }
 katana-primitives = { path = "../primitives" }
 katana-provider = { path = "../storage/provider" }
@@ -37,6 +38,7 @@ url.workspace = true
 [dev-dependencies]
 assert_matches.workspace = true
 hex = "0.4.3"
+tempfile = "3.8.1"
 
 [features]
 messaging = [ "ethers" ]

crates/katana/core/src/backend/config.rs

@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use katana_primitives::block::GasPrices;
 use katana_primitives::chain::ChainId;
 use katana_primitives::env::BlockEnv;
@@ -15,6 +17,7 @@ pub struct StarknetConfig {
     pub fork_rpc_url: Option<Url>,
     pub fork_block_number: Option<u64>,
     pub disable_validate: bool,
+    pub db_dir: Option<PathBuf>,
 }
 
 impl StarknetConfig {
@@ -40,6 +43,7 @@ impl Default for StarknetConfig {
             fork_block_number: None,
             env: Environment::default(),
             disable_validate: false,
+            db_dir: None,
         }
     }
 }

crates/katana/core/src/backend/mod.rs

@@ -104,6 +104,12 @@ impl Backend {
             .expect("able to create forked blockchain");
 
             (blockchain, forked_chain_id.into())
+        } else if let Some(db_path) = &config.db_dir {
+            (
+                Blockchain::new_with_db(db_path, &block_env)
+                    .expect("able to create blockchain from db"),
+                config.env.chain_id,
+            )
         } else {
             let blockchain = Blockchain::new_with_genesis(InMemoryProvider::new(), &block_env)
                 .expect("able to create blockchain from genesis block");

crates/katana/core/src/backend/storage.rs

@@ -1,11 +1,16 @@
+use std::path::Path;
+
 use anyhow::Result;
+use katana_db::init_db;
+use katana_db::utils::is_database_empty;
 use katana_primitives::block::{
     Block, BlockHash, FinalityStatus, Header, PartialHeader, SealedBlockWithStatus,
 };
 use katana_primitives::env::BlockEnv;
 use katana_primitives::state::StateUpdatesWithDeclaredClasses;
 use katana_primitives::version::CURRENT_STARKNET_VERSION;
 use katana_primitives::FieldElement;
+use katana_provider::providers::db::DbProvider;
 use katana_provider::traits::block::{BlockProvider, BlockWriter};
 use katana_provider::traits::contract::ContractClassWriter;
 use katana_provider::traits::env::BlockEnvProvider;
@@ -87,6 +92,15 @@ impl Blockchain {
         Self::new_with_block_and_state(provider, block, get_genesis_states_for_testing())
     }
 
+    pub fn new_with_db(db_path: impl AsRef<Path>, block_context: &BlockEnv) -> Result<Self> {
+        if is_database_empty(&db_path) {
+            let provider = DbProvider::new(init_db(db_path)?);
+            Ok(Self::new_with_genesis(provider, block_context)?)
+        } else {
+            Ok(Self::new(DbProvider::new(init_db(db_path)?)))
+        }
+    }
+
     // TODO: make this function to just accept a `Header` created from the forked block.
     /// Builds a new blockchain with a forked block.
     pub fn new_from_forked(
@@ -131,19 +145,27 @@ impl Blockchain {
 
 #[cfg(test)]
 mod tests {
-    use katana_primitives::block::{FinalityStatus, GasPrices};
+    use katana_primitives::block::{
+        Block, FinalityStatus, GasPrices, Header, SealedBlockWithStatus,
+    };
     use katana_primitives::env::BlockEnv;
+    use katana_primitives::receipt::{InvokeTxReceipt, Receipt};
+    use katana_primitives::state::StateUpdatesWithDeclaredClasses;
+    use katana_primitives::transaction::{InvokeTx, Tx, TxWithHash};
     use katana_primitives::FieldElement;
     use katana_provider::providers::in_memory::InMemoryProvider;
     use katana_provider::traits::block::{
-        BlockHashProvider, BlockNumberProvider, BlockStatusProvider, HeaderProvider,
+        BlockHashProvider, BlockNumberProvider, BlockProvider, BlockStatusProvider, BlockWriter,
+        HeaderProvider,
     };
     use katana_provider::traits::state::StateFactoryProvider;
+    use katana_provider::traits::transaction::TransactionProvider;
     use starknet::macros::felt;
 
     use super::Blockchain;
     use crate::constants::{
-        ERC20_CONTRACT_CLASS_HASH, FEE_TOKEN_ADDRESS, UDC_ADDRESS, UDC_CLASS_HASH,
+        ERC20_CONTRACT, ERC20_CONTRACT_CLASS_HASH, FEE_TOKEN_ADDRESS, UDC_ADDRESS, UDC_CLASS_HASH,
+        UDC_CONTRACT,
     };
 
     #[test]
@@ -207,4 +229,106 @@ mod tests {
         assert_eq!(header.parent_hash, FieldElement::ZERO);
         assert_eq!(block_status, FinalityStatus::AcceptedOnL1);
     }
+
+    #[test]
+    fn blockchain_from_db() {
+        let db_path = tempfile::TempDir::new().expect("Failed to create temp dir.").into_path();
+
+        let block_env = BlockEnv {
+            number: 0,
+            timestamp: 0,
+            sequencer_address: Default::default(),
+            l1_gas_prices: GasPrices { eth: 0, strk: 0 },
+        };
+
+        let dummy_tx =
+            TxWithHash { hash: felt!("0xbad"), transaction: Tx::Invoke(InvokeTx::default()) };
+
+        let dummy_block = SealedBlockWithStatus {
+            status: FinalityStatus::AcceptedOnL1,
+            block: Block {
+                header: Header {
+                    parent_hash: FieldElement::ZERO,
+                    number: 1,
+                    gas_prices: GasPrices::default(),
+                    timestamp: 123456,
+                    ..Default::default()
+                },
+                body: vec![dummy_tx.clone()],
+            }
+            .seal(),
+        };
+
+        {
+            let blockchain = Blockchain::new_with_db(&db_path, &block_env)
+                .expect("Failed to create db-backed blockchain storage");
+
+            blockchain
+                .provider()
+                .insert_block_with_states_and_receipts(
+                    dummy_block.clone(),
+                    StateUpdatesWithDeclaredClasses::default(),
+                    vec![Receipt::Invoke(InvokeTxReceipt::default())],
+                )
+                .unwrap();
+
+            // assert genesis state is correct
+
+            let state = blockchain.provider().latest().expect("failed to get latest state");
+
+            let actual_udc_class_hash =
+                state.class_hash_of_contract(*UDC_ADDRESS).unwrap().unwrap();
+            let actual_udc_class = state.class(actual_udc_class_hash).unwrap().unwrap();
+
+            let actual_fee_token_class_hash =
+                state.class_hash_of_contract(*FEE_TOKEN_ADDRESS).unwrap().unwrap();
+            let actual_fee_token_class = state.class(actual_fee_token_class_hash).unwrap().unwrap();
+
+            assert_eq!(actual_udc_class_hash, *UDC_CLASS_HASH);
+            assert_eq!(actual_udc_class, (*UDC_CONTRACT).clone());
+
+            assert_eq!(actual_fee_token_class_hash, *ERC20_CONTRACT_CLASS_HASH);
+            assert_eq!(actual_fee_token_class, (*ERC20_CONTRACT).clone());
+        }
+
+        // re open the db and assert the state is the same and not overwritten
+
+        {
+            let blockchain = Blockchain::new_with_db(&db_path, &block_env)
+                .expect("Failed to create db-backed blockchain storage");
+
+            // assert genesis state is correct
+
+            let state = blockchain.provider().latest().expect("failed to get latest state");
+
+            let actual_udc_class_hash =
+                state.class_hash_of_contract(*UDC_ADDRESS).unwrap().unwrap();
+            let actual_udc_class = state.class(actual_udc_class_hash).unwrap().unwrap();
+
+            let actual_fee_token_class_hash =
+                state.class_hash_of_contract(*FEE_TOKEN_ADDRESS).unwrap().unwrap();
+            let actual_fee_token_class = state.class(actual_fee_token_class_hash).unwrap().unwrap();
+
+            assert_eq!(actual_udc_class_hash, *UDC_CLASS_HASH);
+            assert_eq!(actual_udc_class, (*UDC_CONTRACT).clone());
+
+            assert_eq!(actual_fee_token_class_hash, *ERC20_CONTRACT_CLASS_HASH);
+            assert_eq!(actual_fee_token_class, (*ERC20_CONTRACT).clone());
+
+            let block_number = blockchain.provider().latest_number().unwrap();
+            let block_hash = blockchain.provider().latest_hash().unwrap();
+            let block = blockchain
+                .provider()
+                .block_by_hash(dummy_block.block.header.hash)
+                .unwrap()
+                .unwrap();
+
+            let tx = blockchain.provider().transaction_by_hash(dummy_tx.hash).unwrap().unwrap();
+
+            assert_eq!(block_hash, dummy_block.block.header.hash);
+            assert_eq!(block_number, dummy_block.block.header.header.number);
+            assert_eq!(block, dummy_block.block.unseal());
+            assert_eq!(tx, dummy_tx);
+        }
+    }
 }

crates/katana/src/args.rs

@@ -48,10 +48,11 @@ pub struct KatanaArgs {
 
     #[arg(long)]
     #[arg(value_name = "PATH")]
-    #[arg(help = "Dump the state of chain on exit to the given file.")]
-    #[arg(long_help = "Dump the state of chain on exit to the given file. If the value is a \
-                       directory, the state will be written to `<PATH>/state.bin`.")]
-    pub dump_state: Option<PathBuf>,
+    #[arg(help = "Directory path of the database to initialize from.")]
+    #[arg(long_help = "Directory path of the database to initialize from. The path must either \
+                       be an empty directory or a directory which already contains a previously \
+                       initialized Katana database.")]
+    pub db_dir: Option<PathBuf>,
 
     #[arg(long)]
     #[arg(value_name = "URL")]
@@ -237,6 +238,7 @@ impl KatanaArgs {
                     .validate_max_steps
                     .unwrap_or(DEFAULT_VALIDATE_MAX_STEPS),
             },
+            db_dir: self.db_dir.clone(),
         }
     }
 }

crates/katana/storage/db/src/utils.rs

@@ -16,8 +16,10 @@ pub(crate) fn default_page_size() -> usize {
     os_page_size.clamp(MIN_PAGE_SIZE, LIBMDBX_MAX_PAGE_SIZE)
 }
 
-/// Check if a db is empty. It does not provide any information on the
-/// validity of the data in it. We consider a database as non empty when it's a non empty directory.
+/// Check if a database is empty. We consider a database as empty when (1) `path` it's an empty
+/// directory, (2) if `path` doesn't exist, or (3) `path` is not a directory.
+///
+/// It does not provide any information on the validity of the data in the Db if it isn't empty.
 pub fn is_database_empty<P: AsRef<Path>>(path: P) -> bool {
     let path = path.as_ref();
     if !path.exists() {

crates/katana/storage/provider/src/providers/db/mod.rs

@@ -48,7 +48,7 @@ use crate::traits::transaction::{
 };
 use crate::ProviderResult;
 
-/// A provider implementation that uses a database as a backend.
+/// A provider implementation that uses a persistent database as the backend.
 #[derive(Debug)]
 pub struct DbProvider(DbEnv);