fix: leaf commitment update

portal-verkle-trie/src/nodes/leaf.rs

@@ -1,13 +1,13 @@
 use std::array;
 
-use banderwagon::{Element, Fr, PrimeField, Zero};
+use banderwagon::{Element, Fr, PrimeField};
 use verkle_core::{
     constants::{
         EXTENSION_C1_INDEX, EXTENSION_C2_INDEX, EXTENSION_MARKER_INDEX, EXTENSION_STEM_INDEX,
         VERKLE_NODE_WIDTH,
     },
     msm::{DefaultMsm, MultiScalarMultiplicator},
-    Stem, TrieValue,
+    Stem, TrieValue, TrieValueSplit,
 };
 
 pub struct LeafNode {
@@ -74,7 +74,7 @@ impl LeafNode {
         }
     }
 
-    pub fn version(&self) -> u64 {
+    pub fn marker(&self) -> u64 {
         self.marker
     }
 
@@ -92,23 +92,29 @@ impl LeafNode {
 
     pub fn set(&mut self, index: usize, child: TrieValue) {
         let (new_low_value, new_high_value) = child.split();
-        let (old_low_value, old_high_value) = self.children[index]
-            .replace(child)
-            .map_or_else(|| (Fr::zero(), Fr::zero()), |old_child| old_child.split());
+        let old_value = self.children[index].replace(child);
+        let (old_low_value, old_high_value) = old_value.split();
 
         let suffix_index = index % (VERKLE_NODE_WIDTH / 2);
-        let suffix_commitment_diff = DefaultMsm
-            .scalar_mul(2 * suffix_index, new_low_value - old_low_value)
-            + DefaultMsm.scalar_mul(2 * suffix_index + 1, new_high_value - old_high_value);
+        let suffix_commitment_diff = DefaultMsm.commit_sparse(&[
+            (2 * suffix_index, new_low_value - old_low_value),
+            (2 * suffix_index + 1, new_high_value - old_high_value),
+        ]);
 
         if index < VERKLE_NODE_WIDTH / 2 {
+            let old_c1_commitment_hash = self.c1.map_to_scalar_field();
             self.c1 += suffix_commitment_diff;
-            self.commitment +=
-                DefaultMsm.scalar_mul(EXTENSION_C1_INDEX, self.c1.map_to_scalar_field());
+            self.commitment += DefaultMsm.scalar_mul(
+                EXTENSION_C1_INDEX,
+                self.c1.map_to_scalar_field() - old_c1_commitment_hash,
+            );
         } else {
+            let old_c2_commitment_hash = self.c2.map_to_scalar_field();
             self.c2 += suffix_commitment_diff;
-            self.commitment +=
-                DefaultMsm.scalar_mul(EXTENSION_C2_INDEX, self.c2.map_to_scalar_field());
+            self.commitment += DefaultMsm.scalar_mul(
+                EXTENSION_C2_INDEX,
+                self.c2.map_to_scalar_field() - old_c2_commitment_hash,
+            );
         }
     }
 

portal-verkle-trie/src/nodes/portal/leaf_bundle.rs

@@ -79,13 +79,19 @@ impl LeafBundleNode {
         let diff = child - self.children[index];
         self.children[index] = child;
         if index < self.children.len() / 2 {
+            let old_c1_commitment_hash = self.c1.map_to_scalar_field();
             self.c1 += diff;
-            self.commitment +=
-                DefaultMsm.scalar_mul(EXTENSION_C1_INDEX, self.c1.map_to_scalar_field());
+            self.commitment += DefaultMsm.scalar_mul(
+                EXTENSION_C1_INDEX,
+                self.c1.map_to_scalar_field() - old_c1_commitment_hash,
+            );
         } else {
+            let old_c2_commitment_hash = self.c2.map_to_scalar_field();
             self.c2 += diff;
-            self.commitment +=
-                DefaultMsm.scalar_mul(EXTENSION_C2_INDEX, self.c2.map_to_scalar_field());
+            self.commitment += DefaultMsm.scalar_mul(
+                EXTENSION_C2_INDEX,
+                self.c2.map_to_scalar_field() - old_c2_commitment_hash,
+            );
         }
     }
 

portal-verkle-trie/src/nodes/portal/leaf_fragment.rs

@@ -1,8 +1,8 @@
-use banderwagon::{Element, Fr, Zero};
+use banderwagon::Element;
 use verkle_core::{
     constants::PORTAL_NETWORK_NODE_WIDTH,
     msm::{DefaultMsm, MultiScalarMultiplicator},
-    TrieValue,
+    TrieValue, TrieValueSplit,
 };
 
 pub struct LeafFragmentNode {
@@ -51,15 +51,15 @@ impl LeafFragmentNode {
 
     pub fn set(&mut self, child_index: usize, child: TrieValue) {
         let (new_low_value, new_high_value) = child.split();
-        let (old_low_value, old_high_value) =
-            self.children[child_index].replace(child).map_or_else(
-                || (Fr::zero(), Fr::zero()),
-                |old_child_value| old_child_value.split(),
-            );
+        let old_value = self.children[child_index].replace(child);
+        let (old_low_value, old_high_value) = old_value.split();
 
         let (low_index, high_index) = Self::bases_indices(self.parent_index, child_index);
-        self.commitment += DefaultMsm.scalar_mul(low_index, new_low_value - old_low_value);
-        self.commitment += DefaultMsm.scalar_mul(high_index, new_high_value - old_high_value);
+
+        self.commitment += DefaultMsm.commit_sparse(&[
+            (low_index, new_low_value - old_low_value),
+            (high_index, new_high_value - old_high_value),
+        ]);
     }
 
     pub fn get(&self, index: usize) -> Option<&TrieValue> {

verkle-core/src/lib.rs

@@ -2,10 +2,11 @@ pub mod constants;
 pub mod msm;
 pub mod serde;
 mod stem;
+pub mod storage;
 mod trie_key;
 mod trie_value;
 pub mod utils;
 
 pub use stem::Stem;
 pub use trie_key::TrieKey;
-pub use trie_value::TrieValue;
+pub use trie_value::{TrieValue, TrieValueSplit};

verkle-core/src/msm/multiplicator.rs

@@ -1,9 +1,7 @@
 use banderwagon::{msm::MSMPrecompWnaf, Element, Fr, Zero};
 use once_cell::sync::Lazy;
 
-use crate::constants::VERKLE_NODE_WIDTH;
-
-use super::crs::CRS;
+use crate::{constants::VERKLE_NODE_WIDTH, msm::crs::CRS};
 
 static PRECOMP_WNAF_WINDOW_SIZE: usize = 12;
 

verkle-core/src/stem.rs

@@ -1,6 +1,6 @@
 use alloy_primitives::{bytes, wrap_fixed_bytes};
 
-use super::trie_key::TrieKey;
+use crate::TrieKey;
 
 wrap_fixed_bytes!(pub struct Stem<31>;);
 

verkle-core/src/storage.rs

@@ -0,0 +1,112 @@
+use alloy_primitives::{Address, B256, U256};
+use banderwagon::{CanonicalSerialize, Fr, PrimeField};
+
+use crate::{
+    constants::{
+        BALANCE_LEAF_KEY, CODE_KECCAK_LEAF_KEY, CODE_OFFSET, CODE_SIZE_LEAF_KEY,
+        HEADER_STORAGE_OFFSET, MAIN_STORAGE_OFFSET, NONCE_LEAF_KEY, VERKLE_NODE_WIDTH_U256,
+        VERSION_LEAF_KEY,
+    },
+    msm::{DefaultMsm, MultiScalarMultiplicator},
+    Stem, TrieKey, TrieValue,
+};
+
+type Address32 = B256;
+
+pub struct AccountStorageLayout {
+    address32: Address32,
+    base_storage_stem: Stem,
+}
+
+impl AccountStorageLayout {
+    pub fn new(address: Address) -> Self {
+        let address32 = Address32::left_padding_from(address.as_slice());
+        Self {
+            address32,
+            base_storage_stem: tree_key(&address32, &U256::ZERO).into(),
+        }
+    }
+
+    pub fn version_key(&self) -> TrieKey {
+        TrieKey::from_stem_and_last_byte(&self.base_storage_stem, VERSION_LEAF_KEY)
+    }
+
+    pub fn balance_key(&self) -> TrieKey {
+        TrieKey::from_stem_and_last_byte(&self.base_storage_stem, BALANCE_LEAF_KEY)
+    }
+
+    pub fn nonce_key(&self) -> TrieKey {
+        TrieKey::from_stem_and_last_byte(&self.base_storage_stem, NONCE_LEAF_KEY)
+    }
+
+    pub fn code_hash_key(&self) -> TrieKey {
+        TrieKey::from_stem_and_last_byte(&self.base_storage_stem, CODE_KECCAK_LEAF_KEY)
+    }
+
+    pub fn code_size_key(&self) -> TrieKey {
+        TrieKey::from_stem_and_last_byte(&self.base_storage_stem, CODE_SIZE_LEAF_KEY)
+    }
+
+    pub fn storage_slot_key(&self, storage_key: U256) -> TrieKey {
+        let pos = if storage_key < CODE_OFFSET - HEADER_STORAGE_OFFSET {
+            HEADER_STORAGE_OFFSET + storage_key
+        } else {
+            MAIN_STORAGE_OFFSET + storage_key
+        };
+        tree_key(&self.address32, &pos)
+    }
+
+    pub fn code_key(&self, chunk_id: usize) -> TrieKey {
+        let pos = CODE_OFFSET + U256::from(chunk_id);
+        tree_key(&self.address32, &pos)
+    }
+
+    pub fn chunkify_code(&self, code: &[u8]) -> Vec<(TrieKey, TrieValue)> {
+        const PUSH_OFFSET: u8 = 95;
+        const PUSH1: u8 = PUSH_OFFSET + 1;
+        const PUSH32: u8 = PUSH_OFFSET + 32;
+
+        let mut remaining_push_data = 0u8;
+        let mut result = vec![];
+        for (chunk_id, chunk) in code.chunks(31).enumerate() {
+            let mut value = Vec::with_capacity(32);
+            value.push(remaining_push_data.min(31));
+            value.extend(chunk);
+            value.resize(32, 0);
+            result.push((self.code_key(chunk_id), B256::from_slice(&value).into()));
+
+            // update remaining_push_data for next chunk
+            for chunk_byte in chunk {
+                if remaining_push_data > 0 {
+                    remaining_push_data -= 1;
+                } else if (PUSH1..=PUSH32).contains(chunk_byte) {
+                    remaining_push_data = chunk_byte - PUSH_OFFSET;
+                }
+            }
+        }
+        result
+    }
+}
+
+fn tree_key(address: &Address32, storage_pos: &U256) -> TrieKey {
+    let tree_index = storage_pos / VERKLE_NODE_WIDTH_U256;
+    let sub_index = (storage_pos % VERKLE_NODE_WIDTH_U256).byte(0);
+
+    let tree_index_bytes = tree_index.to_le_bytes::<32>();
+
+    let scalars = [
+        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 hash_commitment = DefaultMsm.commit_lagrange(&scalars).map_to_scalar_field();
+
+    let mut key = TrieKey::new(B256::ZERO);
+    hash_commitment
+        .serialize_compressed(key.as_mut_slice())
+        .unwrap();
+    key[31] = sub_index;
+    key
+}

verkle-core/src/trie_key.rs

@@ -1,9 +1,25 @@
 use alloy_primitives::B256;
-use derive_more::{Constructor, Deref, From, Index};
+use derive_more::{Constructor, Deref, DerefMut, From, Index, IndexMut};
+use serde::{Deserialize, Serialize};
 
-use super::stem::Stem;
+use crate::Stem;
 
-#[derive(PartialEq, Eq, Clone, Copy, Constructor, Index, Deref, From)]
+#[derive(
+    Debug,
+    Hash,
+    PartialEq,
+    Eq,
+    Clone,
+    Copy,
+    Constructor,
+    Index,
+    IndexMut,
+    Deref,
+    DerefMut,
+    From,
+    Serialize,
+    Deserialize,
+)]
 pub struct TrieKey(B256);
 
 impl TrieKey {

verkle-core/src/trie_value.rs

@@ -1,5 +1,5 @@
-use alloy_primitives::B256;
-use banderwagon::{Fr, PrimeField};
+use alloy_primitives::{B256, U256};
+use banderwagon::{Fr, PrimeField, Zero};
 use derive_more::{Constructor, Deref, From, Index};
 use serde::{Deserialize, Serialize};
 
@@ -19,12 +19,20 @@ use serde::{Deserialize, Serialize};
 )]
 pub struct TrieValue(B256);
 
-impl TrieValue {
+impl From<U256> for TrieValue {
+    fn from(value: U256) -> Self {
+        Self(value.to_le_bytes().into())
+    }
+}
+
+pub trait TrieValueSplit {
     /// Splits self into low (first 16 bytes) and high (second 16 bytes) values, and converts them
     /// to `Fr` scalar field.
-    ///
-    /// It also adds 2^128 to the low value.
-    pub fn split(&self) -> (Fr, Fr) {
+    fn split(&self) -> (Fr, Fr);
+}
+
+impl TrieValueSplit for TrieValue {
+    fn split(&self) -> (Fr, Fr) {
         let (low_value, high_value) = self.split_at(16);
         let mut low_value = Vec::from(low_value);
         low_value.push(1);
@@ -34,3 +42,12 @@ impl TrieValue {
         )
     }
 }
+
+impl<T: TrieValueSplit> TrieValueSplit for Option<T> {
+    fn split(&self) -> (Fr, Fr) {
+        match self {
+            None => (Fr::zero(), Fr::zero()),
+            Some(value) => value.split(),
+        }
+    }
+}