apollo_storage: add state_commitment_infos table

Author: itamar-starkware

Cargo.lock

@@ -7,7 +7,7 @@ license-file.workspace = true
 description = "A storage implementation for a Starknet node."
 
 [features]
-os_input = ["dep:blockifier"]
+os_input = ["dep:blockifier", "dep:starknet_committer"]
 storage_cli = ["clap", "reqwest"]
 testing = ["reqwest", "starknet_api/testing", "tempfile", "tower"]
 
@@ -40,6 +40,7 @@ serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true, features = ["arbitrary_precision"] }
 starknet-types-core = { workspace = true, features = ["papyrus-serialization"] }
 starknet_api.workspace = true
+starknet_committer = { workspace = true, optional = true }
 tempfile = { workspace = true, optional = true }
 thiserror.workspace = true
 tokio = { workspace = true, features = ["rt-multi-thread"] }
@@ -69,6 +70,7 @@ rstest.workspace = true
 schemars.workspace = true
 simple_logger.workspace = true
 starknet_api = { workspace = true, features = ["testing"] }
+starknet_committer = { workspace = true, features = ["testing"] }
 tempfile = { workspace = true }
 test-case.workspace = true
 test-log.workspace = true

crates/apollo_storage/Cargo.toml

@@ -46,7 +46,7 @@ use crate::db::table_types::TableType;
 #[cfg(not(feature = "os_input"))]
 const MAX_DBS: u64 = 25;
 #[cfg(feature = "os_input")]
-const MAX_DBS: u64 = 26;
+const MAX_DBS: u64 = 27;
 
 // Note that NO_TLS mode is used by default.
 type EnvironmentKind = WriteMap;

crates/apollo_storage/src/db/mod.rs

@@ -91,6 +91,8 @@ pub mod global_root_marker;
 #[allow(missing_docs)]
 pub mod metrics;
 pub mod partial_block_hash;
+#[cfg(feature = "os_input")]
+pub mod state_commitment_infos;
 pub mod storage_metrics;
 // TODO(yair): Make the compression_utils module pub(crate) or extract it from the crate.
 #[doc(hidden)]
@@ -184,6 +186,8 @@ use crate::header::StorageBlockHeader;
 use crate::metrics::{register_metrics, STORAGE_COMMIT_LATENCY};
 use crate::mmap_file::MMapFileStats;
 use crate::state::data::IndexedDeprecatedContractClass;
+#[cfg(feature = "os_input")]
+use crate::state_commitment_infos::StateCommitmentInfos;
 use crate::storage_reader_server::{
     create_storage_reader_server,
     ServerConfig,
@@ -278,6 +282,8 @@ fn open_storage_internal(
             .create_simple_table("stateless_compiled_class_hash_v2")?,
         #[cfg(feature = "os_input")]
         accessed_keys: db_writer.create_simple_table("accessed_keys")?,
+        #[cfg(feature = "os_input")]
+        state_commitment_infos: db_writer.create_simple_table("state_commitment_infos")?,
     });
     let (file_writers, file_readers) = open_storage_files(
         &storage_config.db_config,
@@ -712,7 +718,9 @@ struct_field_names! {
         stateless_compiled_class_hash_v2: TableIdentifier<ClassHash, NoVersionValueWrapper<CompiledClassHash>, SimpleTable>,
 
         #[cfg(feature = "os_input")]
-        accessed_keys: TableIdentifier<BlockNumber, VersionZeroWrapper<LocationInFile>, SimpleTable>
+        accessed_keys: TableIdentifier<BlockNumber, VersionZeroWrapper<LocationInFile>, SimpleTable>,
+        #[cfg(feature = "os_input")]
+        state_commitment_infos: TableIdentifier<BlockNumber, VersionZeroWrapper<LocationInFile>, SimpleTable>
     }
 }
 
@@ -879,6 +887,8 @@ struct FileHandlers<Mode: TransactionKind> {
     transaction: FileHandler<VersionZeroWrapper<Transaction>, Mode>,
     #[cfg(feature = "os_input")]
     accessed_keys: FileHandler<VersionZeroWrapper<AccessedKeys>, Mode>,
+    #[cfg(feature = "os_input")]
+    state_commitment_infos: FileHandler<VersionZeroWrapper<StateCommitmentInfos>, Mode>,
 }
 
 impl FileHandlers<RW> {
@@ -921,6 +931,14 @@ impl FileHandlers<RW> {
         self.clone().accessed_keys.append(accessed_keys)
     }
 
+    #[cfg(feature = "os_input")]
+    fn append_state_commitment_infos(
+        &self,
+        state_commitment_infos: &StateCommitmentInfos,
+    ) -> LocationInFile {
+        self.clone().state_commitment_infos.append(state_commitment_infos)
+    }
+
     // TODO(dan): Consider 1. flushing only the relevant files, 2. flushing concurrently.
     #[latency_histogram("storage_file_handler_flush_latency_seconds", false)]
     fn flush(&self) {
@@ -933,6 +951,8 @@ impl FileHandlers<RW> {
         self.transaction.flush();
         #[cfg(feature = "os_input")]
         self.accessed_keys.flush();
+        #[cfg(feature = "os_input")]
+        self.state_commitment_infos.flush();
     }
 }
 
@@ -948,6 +968,8 @@ impl<Mode: TransactionKind> FileHandlers<Mode> {
             ("transaction".to_string(), self.transaction.stats()),
             #[cfg(feature = "os_input")]
             ("accessed_keys".to_string(), self.accessed_keys.stats()),
+            #[cfg(feature = "os_input")]
+            ("state_commitment_infos".to_string(), self.state_commitment_infos.stats()),
         ])
     }
 
@@ -1017,6 +1039,17 @@ impl<Mode: TransactionKind> FileHandlers<Mode> {
             msg: format!("AccessedKeys at location {location:?} not found."),
         })
     }
+
+    #[cfg(feature = "os_input")]
+    // Returns the commitment infos at the given location or an error in case they don't exist.
+    pub(crate) fn get_state_commitment_infos_unchecked(
+        &self,
+        location: LocationInFile,
+    ) -> StorageResult<StateCommitmentInfos> {
+        self.state_commitment_infos.get(location)?.ok_or(StorageError::DBInconsistency {
+            msg: format!("StateCommitmentInfos at location {location:?} not found."),
+        })
+    }
 }
 
 fn open_storage_files(
@@ -1055,6 +1088,9 @@ fn open_storage_files(
     #[cfg(feature = "os_input")]
     let (accessed_keys_writer, accessed_keys_reader) =
         open_storage_file!("accessed_keys", AccessedKeys)?;
+    #[cfg(feature = "os_input")]
+    let (state_commitment_infos_writer, state_commitment_infos_reader) =
+        open_storage_file!("state_commitment_infos", StateCommitmentInfos)?;
 
     Ok((
         FileHandlers {
@@ -1066,6 +1102,8 @@ fn open_storage_files(
             transaction: transaction_writer,
             #[cfg(feature = "os_input")]
             accessed_keys: accessed_keys_writer,
+            #[cfg(feature = "os_input")]
+            state_commitment_infos: state_commitment_infos_writer,
         },
         FileHandlers {
             thin_state_diff: thin_state_diff_reader,
@@ -1076,6 +1114,8 @@ fn open_storage_files(
             transaction: transaction_reader,
             #[cfg(feature = "os_input")]
             accessed_keys: accessed_keys_reader,
+            #[cfg(feature = "os_input")]
+            state_commitment_infos: state_commitment_infos_reader,
         },
     ))
 }
@@ -1098,4 +1138,7 @@ pub enum OffsetKind {
     /// An accessed-keys file.
     #[cfg(feature = "os_input")]
     AccessedKeys,
+    /// A state-commitment-infos file.
+    #[cfg(feature = "os_input")]
+    StateCommitmentInfos,
 }

crates/apollo_storage/src/lib.rs

@@ -372,6 +372,8 @@ auto_storage_serde! {
         Transaction = 5,
         #[cfg(feature = "os_input")]
         AccessedKeys = 6,
+        #[cfg(feature = "os_input")]
+        StateCommitmentInfos = 7,
     }
     pub struct PartialBlockHashComponents {
         pub header_commitments: BlockHeaderCommitments,

crates/apollo_storage/src/serialization/serializers.rs

@@ -0,0 +1,94 @@
+//! Storage for the per-block OS-input commitment infos (state-trie commitment data for the OS).
+
+use starknet_api::block::BlockNumber;
+pub use starknet_committer::patricia_merkle_tree::types::StateCommitmentInfos;
+
+#[cfg(test)]
+#[path = "state_commitment_infos_test.rs"]
+mod state_commitment_infos_test;
+
+use crate::compression_utils::{compress, decompress};
+use crate::db::serialization::{StorageSerde, StorageSerdeError};
+use crate::db::table_types::Table;
+use crate::db::{TransactionKind, RW};
+use crate::{OffsetKind, StorageResult, StorageTxn};
+
+impl StorageSerde for StateCommitmentInfos {
+    fn serialize_into(&self, res: &mut impl std::io::Write) -> Result<(), StorageSerdeError> {
+        let bytes = serde_json::to_vec(self)?;
+        let compressed = compress(bytes.as_slice())?;
+        compressed.serialize_into(res)
+    }
+
+    fn deserialize_from(bytes: &mut impl std::io::Read) -> Option<Self> {
+        let compressed = Vec::<u8>::deserialize_from(bytes)?;
+        let data = decompress(compressed.as_slice()).ok()?;
+        serde_json::from_slice(&data).ok()
+    }
+}
+
+/// Interface for reading the OS-input commitment infos from storage.
+pub trait StateCommitmentInfosStorageReader<Mode: TransactionKind> {
+    /// Returns the commitment infos for the given block, or `None` if not stored.
+    fn get_state_commitment_infos(
+        &self,
+        block_number: BlockNumber,
+    ) -> StorageResult<Option<StateCommitmentInfos>>;
+}
+
+/// Interface for writing the OS-input commitment infos to storage.
+pub trait StateCommitmentInfosStorageWriter
+where
+    Self: Sized,
+{
+    /// Appends the commitment infos for the given block to storage.
+    fn append_state_commitment_infos(
+        self,
+        block_number: BlockNumber,
+        state_commitment_infos: &StateCommitmentInfos,
+    ) -> StorageResult<Self>;
+
+    /// Removes the commitment infos for the given block from storage.
+    /// If no entry exists for the block, returns without error.
+    fn revert_state_commitment_infos(self, block_number: BlockNumber) -> StorageResult<Self>;
+}
+
+impl<Mode: TransactionKind> StateCommitmentInfosStorageReader<Mode> for StorageTxn<'_, Mode> {
+    fn get_state_commitment_infos(
+        &self,
+        block_number: BlockNumber,
+    ) -> StorageResult<Option<StateCommitmentInfos>> {
+        let table = self.open_table(&self.tables.state_commitment_infos)?;
+        let Some(location) = table.get(&self.txn, &block_number)? else {
+            return Ok(None);
+        };
+        Ok(Some(self.file_handlers.get_state_commitment_infos_unchecked(location)?))
+    }
+}
+
+impl StateCommitmentInfosStorageWriter for StorageTxn<'_, RW> {
+    fn append_state_commitment_infos(
+        self,
+        block_number: BlockNumber,
+        state_commitment_infos: &StateCommitmentInfos,
+    ) -> StorageResult<Self> {
+        let file_offset_table = self.txn.open_table(&self.tables.file_offsets)?;
+        let state_commitment_infos_table = self.open_table(&self.tables.state_commitment_infos)?;
+
+        let location = self.file_handlers.append_state_commitment_infos(state_commitment_infos);
+        state_commitment_infos_table.upsert(&self.txn, &block_number, &location)?;
+        file_offset_table.upsert(
+            &self.txn,
+            &OffsetKind::StateCommitmentInfos,
+            &location.next_offset(),
+        )?;
+
+        Ok(self)
+    }
+
+    fn revert_state_commitment_infos(self, block_number: BlockNumber) -> StorageResult<Self> {
+        let state_commitment_infos_table = self.open_table(&self.tables.state_commitment_infos)?;
+        state_commitment_infos_table.delete(&self.txn, &block_number)?;
+        Ok(self)
+    }
+}

crates/apollo_storage/src/state_commitment_infos.rs

@@ -0,0 +1,73 @@
+use std::collections::HashMap;
+
+use starknet_api::block::BlockNumber;
+use starknet_committer::patricia_merkle_tree::types::{CommitmentInfo, StateCommitmentInfos};
+
+use crate::state_commitment_infos::{
+    StateCommitmentInfosStorageReader,
+    StateCommitmentInfosStorageWriter,
+};
+use crate::test_utils::get_test_storage;
+
+fn sample_state_commitment_infos() -> StateCommitmentInfos {
+    StateCommitmentInfos {
+        contracts_trie_commitment_info: CommitmentInfo::default(),
+        classes_trie_commitment_info: CommitmentInfo::default(),
+        storage_tries_commitment_infos: HashMap::new(),
+    }
+}
+
+#[test]
+fn append_and_get_state_commitment_infos() {
+    let (reader, mut writer) = get_test_storage().0;
+    let height = BlockNumber(5);
+    let state_commitment_infos = sample_state_commitment_infos();
+
+    // No infos stored for the height yet.
+    assert_eq!(reader.begin_ro_txn().unwrap().get_state_commitment_infos(height).unwrap(), None);
+
+    writer
+        .begin_rw_txn()
+        .unwrap()
+        .append_state_commitment_infos(height, &state_commitment_infos)
+        .unwrap()
+        .commit()
+        .unwrap();
+
+    assert_eq!(
+        reader.begin_ro_txn().unwrap().get_state_commitment_infos(height).unwrap(),
+        Some(state_commitment_infos)
+    );
+    // A different height is still empty.
+    assert_eq!(
+        reader.begin_ro_txn().unwrap().get_state_commitment_infos(BlockNumber(6)).unwrap(),
+        None
+    );
+}
+
+#[test]
+fn revert_state_commitment_infos() {
+    let (reader, mut writer) = get_test_storage().0;
+    let height = BlockNumber(5);
+
+    writer
+        .begin_rw_txn()
+        .unwrap()
+        .append_state_commitment_infos(height, &sample_state_commitment_infos())
+        .unwrap()
+        .revert_state_commitment_infos(height)
+        .unwrap()
+        .commit()
+        .unwrap();
+
+    assert_eq!(reader.begin_ro_txn().unwrap().get_state_commitment_infos(height).unwrap(), None);
+
+    // Reverting a height with no stored infos is a no-op.
+    writer
+        .begin_rw_txn()
+        .unwrap()
+        .revert_state_commitment_infos(BlockNumber(99))
+        .unwrap()
+        .commit()
+        .unwrap();
+}