ModExtensionGadget: is_account

privacy-scaling-explorations · Apr 22, 2024 · 646601d · 646601d
1 parent 96ee425
commit 646601d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 10 deletions.
diff --git a/zkevm-circuits/src/mpt_circuit/account_leaf.rs b/zkevm-circuits/src/mpt_circuit/account_leaf.rs
@@ -359,6 +359,7 @@ impl<F: Field> AccountLeafConfig<F> {
                     ctx.clone(),
                     parent_data,
                     key_data,
+                    true,
                 );
             }};
 
@@ -836,7 +837,7 @@ impl<F: Field> AccountLeafConfig<F> {
                 &account.mod_list_rlp_bytes[1],
             ];
             self.mod_extension
-                .assign(region, offset, rlp_values, mod_list_rlp_bytes, true)?;
+                .assign(region, offset, rlp_values, mod_list_rlp_bytes)?;
         }
 
         let mut new_value = value[false.idx()];

diff --git a/zkevm-circuits/src/mpt_circuit/mod_extension.rs b/zkevm-circuits/src/mpt_circuit/mod_extension.rs
@@ -34,6 +34,7 @@ pub(crate) struct ModExtensionGadget<F> {
     is_short_branch: IsEqualGadget<F>,
     is_key_part_odd: [Cell<F>; 2], // Whether the number of nibbles is odd or not.
     mult_key: Cell<F>,
+    is_account: bool,
 }
 
 impl<F: Field> ModExtensionGadget<F> {
@@ -43,49 +44,67 @@ impl<F: Field> ModExtensionGadget<F> {
         ctx: MPTContext<F>,
         parent_data: &mut [ParentData<F>; 2],
         key_data: &mut [KeyData<F>; 2],
+        is_account: bool,
     ) -> Self {
         let mut config = ModExtensionGadget::default();
+        config.is_account = is_account;
+
+        let mut long_ext_node_key = StorageRowType::LongExtNodeKey as usize;
+        let mut short_ext_node_key = StorageRowType::ShortExtNodeKey as usize;
+        let mut long_ext_node_nibbles = StorageRowType::LongExtNodeNibbles as usize;
+        let mut short_ext_node_nibbles = StorageRowType::ShortExtNodeNibbles as usize;
+        let mut long_ext_node_value = StorageRowType::LongExtNodeValue as usize;
+        let mut short_ext_node_value = StorageRowType::ShortExtNodeValue as usize;
+
+        if is_account {
+            long_ext_node_key = AccountRowType::LongExtNodeKey as usize;
+            short_ext_node_key = AccountRowType::ShortExtNodeKey as usize;
+            long_ext_node_nibbles = AccountRowType::LongExtNodeNibbles as usize;
+            short_ext_node_nibbles = AccountRowType::ShortExtNodeNibbles as usize;
+            long_ext_node_value = AccountRowType::LongExtNodeValue as usize;
+            short_ext_node_value = AccountRowType::ShortExtNodeValue as usize;
+        }
 
         circuit!([meta, cb], {
             let key_items = [
                 ctx.rlp_item(
                     meta,
                     cb,
-                    StorageRowType::LongExtNodeKey as usize,
+                    long_ext_node_key,
                     RlpItemType::Key,
                 ),
                 ctx.rlp_item(
                     meta,
                     cb,
-                    StorageRowType::ShortExtNodeKey as usize,
+                    short_ext_node_key,
                     RlpItemType::Key,
                 ),
             ];
             let key_nibbles = [
                 ctx.rlp_item(
                     meta,
                     cb,
-                    StorageRowType::LongExtNodeNibbles as usize,
+                    long_ext_node_nibbles,
                     RlpItemType::Nibbles,
                 ),
                 ctx.rlp_item(
                     meta,
                     cb,
-                    StorageRowType::ShortExtNodeNibbles as usize,
+                    short_ext_node_nibbles,
                     RlpItemType::Nibbles,
                 ),
             ];
             let rlp_value = [
                 ctx.rlp_item(
                     meta,
                     cb,
-                    StorageRowType::LongExtNodeValue as usize,
+                    long_ext_node_value,
                     RlpItemType::Value,
                 ),
                 ctx.rlp_item(
                     meta,
                     cb,
-                    StorageRowType::ShortExtNodeValue as usize,
+                    short_ext_node_value,
                     RlpItemType::Value,
                 ),
             ];
@@ -269,7 +288,6 @@ impl<F: Field> ModExtensionGadget<F> {
         offset: usize,
         rlp_values: &[RLPItemWitness],
         list_rlp_bytes: [&[u8]; 2],
-        is_account: bool,
     ) -> Result<(), Error> {
         let mut key_items = [
             rlp_values[StorageRowType::LongExtNodeKey as usize].clone(),
@@ -288,7 +306,7 @@ impl<F: Field> ModExtensionGadget<F> {
             rlp_values[StorageRowType::ShortExtNodeNibbles as usize].clone(),
         ];
 
-        if is_account {
+        if self.is_account {
             key_items = [
                 rlp_values[AccountRowType::LongExtNodeKey as usize].clone(),
                 rlp_values[AccountRowType::ShortExtNodeKey as usize].clone(),

diff --git a/zkevm-circuits/src/mpt_circuit/storage_leaf.rs b/zkevm-circuits/src/mpt_circuit/storage_leaf.rs
@@ -275,6 +275,7 @@ impl<F: Field> StorageLeafConfig<F> {
                     ctx.clone(),
                     parent_data,
                     key_data,
+                    false,
                 );
             }};
 
@@ -685,7 +686,7 @@ impl<F: Field> StorageLeafConfig<F> {
                 &storage.mod_list_rlp_bytes[1],
             ];
             self.mod_extension
-                .assign(region, offset, rlp_values, mod_list_rlp_bytes, false)?;
+                .assign(region, offset, rlp_values, mod_list_rlp_bytes)?;
         }
 
         let mut new_value = value_word[false.idx()];