Skip to content

Commit

Permalink
ref: use U256 as key in initial storage changes
Browse files Browse the repository at this point in the history
  • Loading branch information
zeapoz committed Apr 15, 2024
1 parent fa16ca6 commit 27a1202
Show file tree
Hide file tree
Showing 6 changed files with 25 additions and 34 deletions.
10 changes: 5 additions & 5 deletions src/processor/snapshot/database.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}
Expand Down
5 changes: 3 additions & 2 deletions src/processor/snapshot/exporter.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::path::{Path, PathBuf};

use ethers::types::U64;
use ethers::types::{U256, U64};
use eyre::Result;

use super::{
Expand Down Expand Up @@ -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 {
Expand Down
20 changes: 8 additions & 12 deletions src/processor/snapshot/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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),
Expand All @@ -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();
}
Expand Down
9 changes: 4 additions & 5 deletions src/processor/tree/tree_wrapper.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
12 changes: 3 additions & 9 deletions state-reconstruct-fetcher/src/types/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down Expand Up @@ -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);
};
}
}
Expand Down Expand Up @@ -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);
};
}
}
Expand Down
3 changes: 2 additions & 1 deletion state-reconstruct-fetcher/src/types/v1.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]>,
Expand Down Expand Up @@ -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);
}

Expand Down

0 comments on commit 27a1202

Please sign in to comment.