From 63e320159e0c006a9975a4499f69f040c6a03740 Mon Sep 17 00:00:00 2001
From: Yoav Gross <yoav@starkware.co>
Date: Sun, 10 Nov 2024 15:38:08 +0200
Subject: [PATCH] feat(blockifier): add n_allocated_keys

---
 crates/blockifier/src/state/cached_state.rs | 33 +++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/crates/blockifier/src/state/cached_state.rs b/crates/blockifier/src/state/cached_state.rs
index 93988e7f66..e89fae21c1 100644
--- a/crates/blockifier/src/state/cached_state.rs
+++ b/crates/blockifier/src/state/cached_state.rs
@@ -401,7 +401,9 @@ impl StateCache {
     /// reads. Assumes (and enforces) all initial reads are cached.
     pub fn to_state_diff(&self) -> StateChanges {
         let state_maps = self.writes.diff(&self.initial_reads);
-        StateChanges { state_maps }
+        let allocated_keys =
+            AllocatedKeys::from_storage_diff(&self.writes.storage, &self.initial_reads.storage);
+        StateChanges { state_maps, allocated_keys }
     }
 
     fn declare_contract(&mut self, class_hash: ClassHash) {
@@ -677,11 +679,38 @@ impl StateChangesKeys {
     }
 }
 
+#[cfg_attr(any(feature = "testing", test), derive(Clone))]
+#[derive(Debug, Default, Eq, PartialEq)]
+pub struct AllocatedKeys(HashSet<StorageEntry>);
+
+impl AllocatedKeys {
+    pub fn extend(&mut self, state_change: &StateChanges) {
+        self.0.extend(&state_change.allocated_keys.0);
+        // TODO: Remove keys that are set back to zero.
+    }
+
+    pub fn count(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Collect entries that turn zero -> nonzero.
+    pub fn from_storage_diff(
+        _updated_storage: &HashMap<StorageEntry, Felt>,
+        _base_storage: &HashMap<StorageEntry, Felt>,
+    ) -> Self {
+        Self(
+            HashSet::new()
+            // TODO: Calculate the difference between the updated_storage and the base_storage.
+        )
+    }
+}
+
 /// Holds the state changes.
 #[cfg_attr(any(feature = "testing", test), derive(Clone))]
 #[derive(Debug, Default, Eq, PartialEq)]
 pub struct StateChanges {
     pub state_maps: StateMaps,
+    pub allocated_keys: AllocatedKeys,
 }
 
 impl StateChanges {
@@ -691,8 +720,8 @@ impl StateChanges {
         let mut merged_state_changes = Self::default();
         for state_change in state_changes {
             merged_state_changes.state_maps.extend(&state_change.state_maps);
+            merged_state_changes.allocated_keys.extend(&state_change);
         }
-
         merged_state_changes
     }