feat: improve tree_key code

morph-dev · Apr 13, 2024 · 2558448 · 2558448
1 parent a0a4f20
commit 2558448
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 19 deletions.
diff --git a/verkle/src/lib.rs b/verkle/src/lib.rs
@@ -5,12 +5,12 @@ use banderwagon::Fr;
 use stem::Stem;
 pub use trie::Trie;
 
-pub mod account;
 mod committer;
 mod constants;
 pub mod crs;
 pub mod nodes;
 pub mod stem;
+pub mod storage;
 pub mod trie;
 mod utils;
 

diff --git a/verkle/src/account.rs → verkle/src/storage.rs b/verkle/src/account.rs → verkle/src/storage.rs
@@ -1,5 +1,5 @@
-use alloy_primitives::{Address, B128, B256, U256};
-use banderwagon::Fr;
+use alloy_primitives::{Address, B256, U256};
+use banderwagon::{Fr, PrimeField};
 
 use crate::{
     committer::DEFAULT_COMMITER,
@@ -13,16 +13,19 @@ use crate::{
     TrieKey, TrieValue,
 };
 
+type Address32 = B256;
+
 pub struct AccountStorageLayout {
-    pub address: Address,
+    address32: Address32,
     base_storage_stem: Stem,
 }
 
 impl AccountStorageLayout {
     pub fn new(address: Address) -> Self {
+        let address32 = Address32::left_padding_from(address.as_slice());
         Self {
-            address,
-            base_storage_stem: tree_key(address, U256::ZERO, 0).into(),
+            address32,
+            base_storage_stem: tree_key(&address32, U256::ZERO, 0).into(),
         }
     }
 
@@ -53,7 +56,7 @@ impl AccountStorageLayout {
             MAIN_STORAGE_OFFSET + storage_key
         };
         tree_key(
-            self.address,
+            &self.address32,
             pos / VERKLE_NODE_WIDTH_U256,
             (pos % VERKLE_NODE_WIDTH_U256).byte(0),
         )
@@ -62,7 +65,7 @@ impl AccountStorageLayout {
     pub fn code_key(&self, chunk_id: usize) -> TrieKey {
         let pos = CODE_OFFSET + U256::from(chunk_id);
         tree_key(
-            self.address,
+            &self.address32,
             pos / VERKLE_NODE_WIDTH_U256,
             (pos % VERKLE_NODE_WIDTH_U256).byte(0),
         )
@@ -82,6 +85,7 @@ impl AccountStorageLayout {
             value.resize(32, 0);
             result.push((self.code_key(chunk_id), TrieValue::from_le_slice(&value)));
 
+            // update remaining_push_data for next chunk
             for chunk_byte in chunk {
                 if remaining_push_data > 0 {
                     remaining_push_data -= 1;
@@ -94,18 +98,16 @@ impl AccountStorageLayout {
     }
 }
 
-fn tree_key(address: Address, tree_index: U256, sub_index: u8) -> TrieKey {
-    let address_bytes = *B256::left_padding_from(address.as_slice());
+fn tree_key(address: &Address32, tree_index: U256, sub_index: u8) -> TrieKey {
     let tree_index_bytes = tree_index.to_le_bytes::<32>();
 
     let scalars = [
-        2u128 + 256 * 64,
-        u128::from_le_bytes(B128::from_slice(&address_bytes[..16]).0),
-        u128::from_le_bytes(B128::from_slice(&address_bytes[16..]).0),
-        u128::from_le_bytes(B128::from_slice(&tree_index_bytes[..16]).0),
-        u128::from_le_bytes(B128::from_slice(&tree_index_bytes[16..]).0),
-    ]
-    .map(Fr::from);
+        Fr::from(2u128 + 256 * 64),
+        Fr::from_le_bytes_mod_order(&address[..16]),
+        Fr::from_le_bytes_mod_order(&address[16..]),
+        Fr::from_le_bytes_mod_order(&tree_index_bytes[..16]),
+        Fr::from_le_bytes_mod_order(&tree_index_bytes[16..]),
+    ];
     let commitment = DEFAULT_COMMITER.commit_lagrange(&scalars);
     let hash_commitment = commitment.map_to_scalar_field();
 

diff --git a/verkle/src/trie.rs b/verkle/src/trie.rs
@@ -3,8 +3,8 @@ use anyhow::Result;
 use banderwagon::Fr;
 
 use crate::{
-    account::AccountStorageLayout,
     nodes::Node,
+    storage::AccountStorageLayout,
     utils::{b256_to_fr, fr_to_b256},
     Db, TrieKey, TrieValue,
 };

diff --git a/verkle/tests/verkle_test_vectors.rs b/verkle/tests/verkle_test_vectors.rs
@@ -6,7 +6,7 @@ mod verkle_test_vectors {
     use anyhow::Result;
     use banderwagon::{CanonicalDeserialize, Element, Fr};
     use db::memory_db::MemoryDb;
-    use verkle::{account::AccountStorageLayout, Trie, TrieValue};
+    use verkle::{storage::AccountStorageLayout, Trie, TrieValue};
 
     fn element_to_fr(str: &str) -> Fr {
         Element::deserialize_compressed(hex::decode(str).unwrap().as_slice())