blockifier: update patricia resource calc according to new OS impl (#10073)

Author: Yoni-Starkware

crates/blockifier/src/bouncer.rs

@@ -646,6 +646,8 @@ impl Bouncer {
         self.visited_storage_entries.extend(&tx_execution_summary.visited_storage_entries);
         // Note: cancelling writes (0 -> 1 -> 0) will not be removed, but it's fine since fee was
         // charged for them.
+        // Also, `get_patricia_update_resources` relies on this property - each cell must
+        // be counted at most once as modified.
         self.state_changes_keys.extend(state_changes_keys);
         self.accumulated_weights.class_hashes_to_migrate.extend(tx_weights.class_hashes_to_migrate);
     }
@@ -846,7 +848,11 @@ pub fn get_tx_weights<S: StateReader>(
         map_class_hash_to_casm_hash_computation_resources(state_reader, executed_class_hashes)?;
 
     // Patricia update + transaction resources.
-    let patrticia_update_resources = get_particia_update_resources(n_visited_storage_entries);
+    let patrticia_update_resources = get_patricia_update_resources(
+        n_visited_storage_entries,
+        // TODO(Yoni): consider counting here the global contract tree and the aliases as well.
+        state_changes_keys.storage_keys.len(),
+    );
     let vm_resources = &patrticia_update_resources + &tx_resources.computation.total_vm_resources();
 
     // Builtin gas costs for stone and for stwo.
@@ -939,12 +945,23 @@ pub fn map_class_hash_to_casm_hash_computation_resources<S: StateReader>(
         .collect()
 }
 
-/// Returns the estimated Cairo resources for Patricia tree updates, or hash invocations
-/// (done by the OS), required for accessing (read/write) the given storage entries.
-// For each tree: n_visited_leaves * log(n_initialized_leaves)
-// as the height of a Patricia tree with N uniformly distributed leaves is ~log(N),
-// and number of visited leaves includes reads and writes.
-pub fn get_particia_update_resources(n_visited_storage_entries: usize) -> ExecutionResources {
+/// Returns the estimated Cairo resources for Patricia tree updates given the accessed and
+/// modified storage entries.
+///
+/// Each access (read or write) requires a traversal of the previous tree, and a write access
+/// requires an additional traversal of the new tree.
+///
+/// Note:
+///   1. n_visited_storage_entries includes both read and write accesses, and may overlap with
+///      n_first_time_modified_storage_entries (if the first access to a cell was write) and may not
+///      (if a cell was read by a previous transaction and is now modified for the first time).
+///   2. In practice, the OS performs a multi-update, which is more efficient than performing
+///      separate updates. However, we use this conservative estimate for simplicity.
+pub fn get_patricia_update_resources(
+    n_visited_storage_entries: usize,
+    n_first_time_modified_storage_entries: usize,
+) -> ExecutionResources {
+    // The height of a Patricia tree with N uniformly distributed leaves is ~log(N).
     const TREE_HEIGHT_UPPER_BOUND: usize = 24;
     // TODO(Yoni, 1/5/2024): re-estimate this.
     const STEPS_IN_TREE_PER_HEIGHT: usize = 16;
@@ -959,8 +976,8 @@ pub fn get_particia_update_resources(n_visited_storage_entries: usize) -> Execut
         n_memory_holes: 0,
     };
 
-    // Multiply by 2 since each storage entry is accessed in both the old and new tree.
-    &resources_per_tree_access * (n_visited_storage_entries * 2)
+    // One traversal per access (read or write), and an additional one per write access.
+    &resources_per_tree_access * (n_visited_storage_entries + n_first_time_modified_storage_entries)
 }
 
 pub fn verify_tx_weights_within_max_capacity<S: StateReader>(

crates/blockifier/src/bouncer_test.rs

@@ -17,7 +17,7 @@ use super::BouncerConfig;
 use crate::blockifier::transaction_executor::TransactionExecutorError;
 use crate::bouncer::{
     builtins_to_gas,
-    get_particia_update_resources,
+    get_patricia_update_resources,
     get_tx_weights,
     map_class_hash_to_casm_hash_computation_resources,
     verify_tx_weights_within_max_capacity,
@@ -448,11 +448,11 @@ fn test_bouncer_try_update_n_txs(
             n_events: 10,
             state_diff_size: 14,
             sierra_gas: GasAmount(
-                542410,
+                406810,
             ),
             n_txs: 20,
             proving_gas: GasAmount(
-                702922,
+                527194,
             ),
         }
     "#
@@ -651,7 +651,7 @@ fn test_proving_gas_minus_sierra_gas_equals_builtin_gas(
     let n_visited_storage_entries = if casm_hash_computation_builtins.is_empty() { 0 } else { 1 };
 
     let mut additional_os_resources =
-        get_particia_update_resources(n_visited_storage_entries).prover_builtins();
+        get_patricia_update_resources(n_visited_storage_entries, 0).prover_builtins();
     add_maps(&mut additional_os_resources, &casm_hash_computation_builtins);
 
     let result = get_tx_weights(