Refactor of state committing

crates/starknet-devnet-core/src/starknet/get_class_impls.rs

@@ -136,8 +136,8 @@ mod tests {
     #[test]
     fn get_class_hash_at_generated_accounts() {
         let (mut starknet, account) = setup(Some(100000000), StateArchiveCapacity::Full);
-        let state_diff = starknet.pending_state.commit_with_diff().unwrap();
-        starknet.generate_new_block(state_diff).unwrap();
+        let state_diff = starknet.commit_with_diff().unwrap();
+        starknet.generate_new_block_and_state(state_diff).unwrap();
 
         let block_number = starknet.get_latest_block().unwrap().block_number();
         let block_id = BlockId::Number(block_number.0);
@@ -150,8 +150,8 @@ mod tests {
     #[test]
     fn get_class_hash_at_generated_accounts_without_state_archive() {
         let (mut starknet, account) = setup(Some(100000000), StateArchiveCapacity::None);
-        let state_diff = starknet.pending_state.commit_with_diff().unwrap();
-        starknet.generate_new_block(state_diff).unwrap();
+        let state_diff = starknet.commit_with_diff().unwrap();
+        starknet.generate_new_block_and_state(state_diff).unwrap();
 
         let block_number = starknet.get_latest_block().unwrap().block_number();
         let block_id = BlockId::Number(block_number.0);
@@ -166,10 +166,11 @@ mod tests {
     #[test]
     fn get_class_at_generated_accounts() {
         let (mut starknet, account) = setup(Some(100000000), StateArchiveCapacity::Full);
-        let state_diff = starknet.pending_state.commit_with_diff().unwrap();
-        starknet.generate_new_block(state_diff).unwrap();
+        let state_diff = starknet.commit_with_diff().unwrap();
+        starknet.generate_new_block_and_state(state_diff).unwrap();
 
         let block_number = starknet.get_latest_block().unwrap().block_number();
+        assert_eq!(block_number.0, 0); // defined via block context in setup
         let block_id = BlockId::Number(block_number.0);
 
         let contract_class = starknet.get_class_at(&block_id, account.account_address).unwrap();

crates/starknet-devnet-core/src/starknet/mod.rs

@@ -287,10 +287,13 @@ impl Starknet {
         }
     }
 
-    /// Transfer data from pending block into new block and save it to blocks collection
-    /// Generates new pending block
-    /// Returns the new block number
-    pub(crate) fn generate_new_block(&mut self, state_diff: StateDiff) -> DevnetResult<Felt> {
+    /// Transfer data from pending block into new block and save it to blocks collection.
+    /// Generates new pending block. Same for pending state.
+    /// Returns the new block number.
+    pub(crate) fn generate_new_block_and_state(
+        &mut self,
+        state_diff: StateDiff,
+    ) -> DevnetResult<Felt> {
         let mut new_block = self.pending_block().clone();
 
         let new_block_number =
@@ -328,8 +331,8 @@ impl Starknet {
 
         // clone_historic() requires self.historic_state, self.historic_state is set in
         // expand_historic(), expand_historic() can be executed from commit_with_diff() - this
-        // is why self.pending_state.commit_with_diff()? is here
-        self.pending_state.commit_with_diff()?;
+        // is why self.commit_with_diff()? is here
+        self.commit_with_diff()?;
 
         // save into blocks state archive
         if self.config.state_archive == StateArchiveCapacity::Full {
@@ -345,6 +348,12 @@ impl Starknet {
         Ok(new_block_hash)
     }
 
+    /// Commits the the changes accumulated in pending state. Check
+    /// `StarknetState::commit_with_diff` for more info.
+    pub fn commit_with_diff(&mut self) -> DevnetResult<StateDiff> {
+        self.pending_state.commit_with_diff()
+    }
+
     /// Handles transaction result either Ok or Error and updates the state accordingly.
     ///
     /// # Arguments
@@ -451,7 +460,7 @@ impl Starknet {
         transaction: &TransactionWithHash,
         tx_info: TransactionExecutionInfo,
     ) -> DevnetResult<()> {
-        let state_diff = self.pending_state.commit_with_diff()?;
+        let state_diff = self.commit_with_diff()?;
 
         let trace = create_trace(
             &mut self.pending_state.state,
@@ -468,7 +477,7 @@ impl Starknet {
 
         // create new block from pending one, only if block on-demand mode is disabled
         if !self.config.blocks_on_demand {
-            self.generate_new_block(state_diff)?;
+            self.generate_new_block_and_state(state_diff)?;
         }
 
         Ok(())
@@ -1189,7 +1198,7 @@ impl Starknet {
 
     /// create new block from pending one
     pub fn create_block(&mut self) -> DevnetResult<(), Error> {
-        self.generate_new_block(StateDiff::default())?;
+        self.generate_new_block_and_state(StateDiff::default())?;
         Ok(())
     }
 
@@ -1459,7 +1468,7 @@ mod tests {
         // blocks collection should not be empty
         assert_eq!(starknet.blocks.hash_to_block.len(), 1);
 
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
         // blocks collection should not be empty
         assert_eq!(starknet.blocks.hash_to_block.len(), 2);
 
@@ -1561,7 +1570,7 @@ mod tests {
         let config =
             StarknetConfig { state_archive: StateArchiveCapacity::Full, ..Default::default() };
         let mut starknet = Starknet::new(&config).unwrap();
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         match starknet.get_mut_state_at(&BlockId::Hash(Felt::from(0).into())) {
             Err(Error::NoBlock) => (),
@@ -1574,7 +1583,7 @@ mod tests {
         let config =
             StarknetConfig { state_archive: StateArchiveCapacity::Full, ..Default::default() };
         let mut starknet = Starknet::new(&config).unwrap();
-        let block_hash = starknet.generate_new_block(StateDiff::default()).unwrap();
+        let block_hash = starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
         starknet.blocks.hash_to_state.remove(&block_hash);
 
         match starknet.get_mut_state_at(&BlockId::Number(1)) {
@@ -1587,7 +1596,7 @@ mod tests {
     fn getting_state_at_without_state_archive() {
         let config = StarknetConfig::default();
         let mut starknet = Starknet::new(&config).unwrap();
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         match starknet.get_mut_state_at(&BlockId::Number(0)) {
             Err(Error::NoStateAtBlock { .. }) => (),
@@ -1696,7 +1705,7 @@ mod tests {
 
         assert_eq!(block_number.0, added_block.header.block_number.0);
 
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         let added_block2 =
             starknet.blocks.get_by_hash(starknet.blocks.last_block_hash.unwrap()).unwrap();
@@ -1710,7 +1719,7 @@ mod tests {
         let config = StarknetConfig::default();
         let mut starknet = Starknet::new(&config).unwrap();
 
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         let num_no_transactions = starknet.get_block_txs_count(&BlockId::Number(1));
 
@@ -1721,7 +1730,7 @@ mod tests {
         // add transaction hash to pending block
         starknet.blocks.pending_block.add_transaction(*tx.get_transaction_hash());
 
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         let num_one_transaction = starknet.get_block_txs_count(&BlockId::Number(2));
 
@@ -1746,7 +1755,7 @@ mod tests {
         .expect("Could not start Devnet");
 
         // generate initial block with empty state
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         // **generate second block**
         // add data to state
@@ -1756,9 +1765,9 @@ mod tests {
             .increment_nonce(dummy_contract_address().try_into().unwrap())
             .unwrap();
         // get state difference
-        let state_diff = starknet.latest_state.commit_with_diff().unwrap();
+        let state_diff = starknet.commit_with_diff().unwrap();
         // generate new block and save the state
-        let second_block = starknet.generate_new_block(state_diff).unwrap();
+        let second_block = starknet.generate_new_block_and_state(state_diff).unwrap();
 
         // **generate third block**
         // add data to state
@@ -1768,9 +1777,9 @@ mod tests {
             .increment_nonce(dummy_contract_address().try_into().unwrap())
             .unwrap();
         // get state difference
-        let state_diff = starknet.pending_state.commit_with_diff().unwrap();
+        let state_diff = starknet.commit_with_diff().unwrap();
         // generate new block and save the state
-        let third_block = starknet.generate_new_block(state_diff).unwrap();
+        let third_block = starknet.generate_new_block_and_state(state_diff).unwrap();
 
         // check modified state at block 1 and 2 to contain the correct value for the nonce
         let second_block_address_nonce = starknet
@@ -1799,9 +1808,9 @@ mod tests {
         let config = StarknetConfig::default();
         let mut starknet = Starknet::new(&config).unwrap();
 
-        starknet.generate_new_block(StateDiff::default()).unwrap();
-        starknet.generate_new_block(StateDiff::default()).unwrap();
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         let latest_block = starknet.get_latest_block();
 
@@ -1812,7 +1821,7 @@ mod tests {
         let config = StarknetConfig::default();
         let mut starknet = Starknet::new(&config).unwrap();
 
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
         starknet
             .blocks
             .pending_block
@@ -1821,7 +1830,7 @@ mod tests {
 
         let sleep_duration_secs = 5;
         thread::sleep(Duration::from_secs(sleep_duration_secs));
-        starknet.generate_new_block(StateDiff::default()).unwrap();
+        starknet.generate_new_block_and_state(StateDiff::default()).unwrap();
 
         let block_timestamp = starknet.get_latest_block().unwrap().header.timestamp;
         // check if the pending_block_timestamp is less than the block_timestamp,

crates/starknet-devnet-core/src/starknet/state_update.rs

@@ -90,11 +90,10 @@ mod tests {
         }
         .into();
 
-        assert_eq!(
-            state_diff.deprecated_declared_classes,
-            expected_state_diff.deprecated_declared_classes
-        );
-        assert_eq!(state_diff.declared_classes, expected_state_diff.declared_classes);
+        let class_diff = (state_diff.deprecated_declared_classes, state_diff.declared_classes);
+        let expected_class_diff =
+            (expected_state_diff.deprecated_declared_classes, expected_state_diff.declared_classes);
+        assert_eq!(class_diff, expected_class_diff);
     }
 
     /// Initializes starknet with account_without_validations
@@ -138,6 +137,10 @@ mod tests {
 
         starknet.restart_pending_block().unwrap();
 
+        // commit the newly declared account class
+        let state_diff = starknet.commit_with_diff().unwrap();
+        starknet.generate_new_block_and_state(state_diff).unwrap();
+
         (starknet, acc.get_address())
     }
 }

crates/starknet-devnet-core/src/state/mod.rs

@@ -71,13 +71,6 @@ impl CommittedClassStorage {
         self.committed.extend(self.staging.drain());
         diff
     }
-
-    /// Skips the staging phase
-    fn insert_and_commit(&mut self, class_hash: ClassHash, contract_class: ContractClass) {
-        assert!(self.staging.is_empty());
-        self.insert(class_hash, contract_class);
-        self.commit();
-    }
 }
 
 pub struct StarknetState {
@@ -117,7 +110,8 @@ impl StarknetState {
         self.rpc_contract_classes.clone()
     }
 
-    pub fn commit_with_diff(&mut self) -> DevnetResult<StateDiff> {
+    /// Commits and returns the state difference accumulated since the previous (historic) state.
+    pub(crate) fn commit_with_diff(&mut self) -> DevnetResult<StateDiff> {
         let diff = StateDiff::generate(&mut self.state, &mut self.rpc_contract_classes)?;
         let new_historic = self.expand_historic(diff.clone())?;
         self.state = CachedState::new(
@@ -330,7 +324,7 @@ impl CustomState for StarknetState {
         };
 
         self.state.state.set_contract_class(class_hash.into(), compiled_class)?;
-        self.rpc_contract_classes.insert_and_commit(class_hash, contract_class);
+        self.rpc_contract_classes.insert(class_hash, contract_class);
         Ok(())
     }