ref: use U256 as key in initial storage changes (#86)

src/processor/snapshot/database.rs

@@ -52,9 +52,7 @@ impl SnapshotDB {
         let processed_value = match value {
             PackingType::NoCompression(v) | PackingType::Transform(v) => v,
             PackingType::Add(_) | PackingType::Sub(_) => {
-                let mut buffer = [0; 32];
-                key.to_big_endian(&mut buffer);
-                let existing_value = if let Some(log) = self.get_storage_log(&buffer)? {
+                let existing_value = if let Some(log) = self.get_storage_log(&key)? {
                     U256::from(log.value.to_fixed_bytes())
                 } else {
                     U256::from(0)
@@ -141,11 +139,13 @@ impl SnapshotDB {
             .map_err(Into::into)
     }
 
-    pub fn get_storage_log(&self, key: &[u8]) -> Result<Option<SnapshotStorageLog>> {
+    pub fn get_storage_log(&self, key: &U256) -> Result<Option<SnapshotStorageLog>> {
         // Unwrapping column family handle here is safe because presence of
         // those CFs is ensured in construction of this DB.
         let storage_logs = self.cf_handle(STORAGE_LOGS).unwrap();
-        self.get_cf(storage_logs, key)
+        let mut byte_key = [0u8; 32];
+        key.to_big_endian(&mut byte_key);
+        self.get_cf(storage_logs, byte_key)
             .map(|v| v.map(|v| bincode::deserialize(&v).unwrap()))
             .map_err(Into::into)
     }

src/processor/snapshot/exporter.rs

@@ -1,6 +1,6 @@
 use std::path::{Path, PathBuf};
 
-use ethers::types::U64;
+use ethers::types::{U256, U64};
 use eyre::Result;
 
 use super::{
@@ -105,7 +105,8 @@ impl SnapshotExporter {
             let mut chunk = SnapshotStorageLogsChunk::default();
             for _ in 0..chunk_size {
                 if let Some(Ok((_, key))) = iterator.next() {
-                    if let Ok(Some(entry)) = self.database.get_storage_log(key.as_ref()) {
+                    let key = U256::from_big_endian(&key);
+                    if let Ok(Some(entry)) = self.database.get_storage_log(&key) {
                         chunk.storage_logs.push(entry);
                     }
                 } else {

src/processor/snapshot/mod.rs

@@ -70,15 +70,14 @@ impl Processor for SnapshotBuilder {
         while let Some(block) = rx.recv().await {
             // Initial calldata.
             for (key, value) in &block.initial_storage_changes {
-                let key = U256::from_little_endian(key);
                 let value = self
                     .database
-                    .process_value(key, *value)
+                    .process_value(*key, *value)
                     .expect("failed to get key from database");
 
                 self.database
                     .insert_storage_log(&mut SnapshotStorageLog {
-                        key,
+                        key: *key,
                         value,
                         miniblock_number_of_initial_write: U64::from(0),
                         l1_batch_number_of_initial_write: U64::from(
@@ -260,11 +259,10 @@ mod tests {
 
             tokio::spawn(async move {
                 let key = U256::from_dec_str("1234").unwrap();
-                let mut key1 = [0u8; 32];
-                key.to_little_endian(&mut key1);
-                let val1 = U256::from_dec_str("5678").unwrap();
+                let val = U256::from_dec_str("5678").unwrap();
                 let mut initial_storage_changes = IndexMap::new();
-                initial_storage_changes.insert(key1, PackingType::NoCompression(val1));
+                initial_storage_changes.insert(key, PackingType::NoCompression(val));
+
                 let repeated_storage_changes = IndexMap::new();
                 let cb = CommitBlock {
                     l1_block_number: Some(1),
@@ -284,13 +282,11 @@ mod tests {
         let db = SnapshotDB::new(PathBuf::from(db_dir.clone())).unwrap();
 
         let key = U256::from_dec_str("1234").unwrap();
-        let mut key1 = [0u8; 32];
-        key.to_big_endian(&mut key1);
-        let Some(sl1) = db.get_storage_log(&key1).unwrap() else {
-            panic!("key1 not found")
+        let Some(log) = db.get_storage_log(&key).unwrap() else {
+            panic!("key not found")
         };
 
-        assert_eq!(sl1.key, key1.into());
+        assert_eq!(log.key, key);
 
         fs::remove_dir_all(db_dir).unwrap();
     }

src/processor/tree/tree_wrapper.rs

@@ -64,15 +64,14 @@ impl TreeWrapper {
         let mut index =
             block.index_repeated_storage_changes - (block.initial_storage_changes.len() as u64);
         for (key, value) in &block.initial_storage_changes {
-            let key = U256::from_little_endian(key);
-            self.insert_known_key(index, key);
-            let value = self.process_value(key, *value);
+            self.insert_known_key(index, *key);
+            let value = self.process_value(*key, *value);
 
-            tree_entries.push(TreeEntry::new(key, index, value));
+            tree_entries.push(TreeEntry::new(*key, index, value));
             self.snapshot
                 .lock()
                 .await
-                .add_key(&key)
+                .add_key(key)
                 .expect("cannot add key");
             index += 1;
         }

state-reconstruct-fetcher/src/types/mod.rs

@@ -87,7 +87,7 @@ pub struct CommitBlock {
     pub new_state_root: Vec<u8>,
     /// Storage write access as a concatenation key-value.
     #[serde(with = "any_key_map")]
-    pub initial_storage_changes: IndexMap<[u8; 32], PackingType>,
+    pub initial_storage_changes: IndexMap<U256, PackingType>,
     /// Storage write access as a concatenation index-value.
     #[serde(with = "any_key_map")]
     pub repeated_storage_changes: IndexMap<u64, PackingType>,
@@ -144,13 +144,10 @@ impl CommitBlock {
                             derived_key,
                             packing_type,
                         } => {
-                            let mut key = [0u8; 32];
-                            derived_key.to_big_endian(&mut key);
-
                             if is_repeated_write {
                                 repeated_storage_changes.insert(derived_key.as_u64(), packing_type);
                             } else {
-                                initial_storage_changes.insert(key, packing_type);
+                                initial_storage_changes.insert(derived_key, packing_type);
                             };
                         }
                     }
@@ -186,13 +183,10 @@ impl CommitBlock {
                     derived_key,
                     packing_type,
                 } => {
-                    let mut key = [0u8; 32];
-                    derived_key.to_big_endian(&mut key);
-
                     if is_repeated_write {
                         repeated_storage_changes.insert(derived_key.as_u64(), packing_type);
                     } else {
-                        initial_storage_changes.insert(key, packing_type);
+                        initial_storage_changes.insert(derived_key, packing_type);
                     };
                 }
             }

state-reconstruct-fetcher/src/types/v1.rs

@@ -24,7 +24,7 @@ pub struct V1 {
     pub priority_operations_hash: Vec<u8>,
     /// Storage write access as a concatenation key-value.
     #[serde(with = "any_key_map")]
-    pub initial_storage_changes: IndexMap<[u8; 32], [u8; 32]>,
+    pub initial_storage_changes: IndexMap<U256, [u8; 32]>,
     /// Storage write access as a concatenation index-value.
     #[serde(with = "any_key_map")]
     pub repeated_storage_changes: IndexMap<u64, [u8; 32]>,
@@ -111,6 +111,7 @@ impl TryFrom<&abi::Token> for V1 {
                 ));
             }
 
+            let key = U256::from_little_endian(&key);
             let _ = blk.initial_storage_changes.insert(key, value);
         }