Adding info whether ext. node is in the last level

Author: miha-stopar

geth-utils/gethutil/mpt/witness/branch.go

@@ -66,7 +66,7 @@ func prepareBranchWitness(rows [][]byte, branch []byte, branchStart int, branchR
 }
 
 func prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []byte, extValues [][]byte, key, driftedInd byte,
-	isBranchSPlaceholder, isBranchCPlaceholder, isExtension bool) Node {
+	isBranchSPlaceholder, isBranchCPlaceholder, isExtension, isLastLevel bool) Node {
 	extensionNode := ExtensionNode{
 		ListRlpBytes: extListRlpBytes,
 	}
@@ -115,6 +115,7 @@ func prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []b
 	extensionBranch := ExtensionBranchNode{
 		IsExtension:   isExtension,
 		IsPlaceholder: [2]bool{isBranchSPlaceholder, isBranchCPlaceholder},
+		IsLastLevel:   isLastLevel,
 		Extension:     extensionNode,
 		Branch:        branchNode,
 	}
@@ -206,10 +207,8 @@ func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
 // (used when one of the proofs have one branch more than the other).
 func addBranchAndPlaceholder(proof1, proof2 [][]byte,
 	extNibblesS, extNibblesC []byte,
-	leafRow0, key, neighbourNode []byte,
-	keyIndex int,
-	additionalBranch, isAccountProof, nonExistingAccountProof,
-	isShorterProofLastLeaf bool) (bool, bool, int, Node) {
+	leafRow0, key []byte,
+	keyIndex int, isShorterProofLastLeaf bool) (bool, bool, int, Node) {
 	len1 := len(proof1)
 	len2 := len(proof2)
 
@@ -284,7 +283,7 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte,
 		driftedInd := getDriftedPosition(leafRow0, numberOfNibbles)
 
 		node = prepareBranchNode(proof1[len1-2], proof1[len1-2], extNode, extNode, extListRlpBytes, extValues,
-			key[keyIndex+numberOfNibbles], driftedInd, false, true, isExtension)
+			key[keyIndex+numberOfNibbles], driftedInd, false, true, isExtension, true)
 
 		// We now get the first nibble of the leaf that was turned into branch.
 		// This first nibble presents the position of the leaf once it moved
@@ -296,7 +295,7 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte,
 		driftedInd := getDriftedPosition(leafRow0, numberOfNibbles)
 
 		node = prepareBranchNode(proof2[len2-2], proof2[len2-2], extNode, extNode, extListRlpBytes, extValues,
-			key[keyIndex+numberOfNibbles], driftedInd, true, false, isExtension)
+			key[keyIndex+numberOfNibbles], driftedInd, true, false, isExtension, true)
 	}
 
 	return isModifiedExtNode, isExtension, numberOfNibbles, node

geth-utils/gethutil/mpt/witness/nodes.go

@@ -70,6 +70,7 @@ type ExtensionBranchNode struct {
 	// nibbles of the original extension node).
 	IsModExtension [2]bool       `json:"is_mod_extension"`
 	IsPlaceholder  [2]bool       `json:"is_placeholder"`
+	IsLastLevel    bool          `json:"is_last_level"`
 	Extension      ExtensionNode `json:"extension"`
 	Branch         BranchNode    `json:"branch"`
 }

geth-utils/gethutil/mpt/witness/prepare_witness.go

@@ -412,7 +412,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 			}
 
 			bNode := prepareBranchNode(proof1[i], proof2[i], extNode1, extNode2, extListRlpBytes, extValues,
-				key[keyIndex], key[keyIndex], false, false, isExtension)
+				key[keyIndex], key[keyIndex], false, false, isExtension, false)
 			nodes = append(nodes, bNode)
 
 			keyIndex += 1
@@ -430,9 +430,8 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 
 			isModifiedExtNode, _, numberOfNibbles, bNode := addBranchAndPlaceholder(proof1, proof2,
 				extNibblesS[len1-1], extNibblesC[len2-1],
-				leafRow0, key, neighbourNode,
-				keyIndex, additionalBranch,
-				isAccountProof, nonExistingAccountProof, isShorterProofLastLeaf)
+				leafRow0, key,
+				keyIndex, isShorterProofLastLeaf)
 
 			nodes = append(nodes, bNode)
 
@@ -531,7 +530,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 
 				extNode := proof2[len(proof2)-1] // Let's name it E1
 				bNode := prepareBranchNode(branchRlp, branchRlp, extNode, extNode, extListRlpBytes, extValues,
-					key[keyIndex], key[keyIndex], false, false, isExtension)
+					key[keyIndex], key[keyIndex], false, false, isExtension, true)
 				nodes = append(nodes, bNode)
 
 				// Let's construct the leaf L1 that will have the correct key (the queried one)

zkevm-circuits/src/mpt_circuit/account_leaf.rs

@@ -339,6 +339,7 @@ impl<F: Field> AccountLeafConfig<F> {
                     true.expr(),
                     false.expr(),
                     false.expr(),
+                    false.expr(),
                     storage_items[is_s.idx()].word(),
                 );
             }
@@ -671,6 +672,7 @@ impl<F: Field> AccountLeafConfig<F> {
                 true,
                 false,
                 false,
+                false,
                 storage_items[is_s.idx()].word(),
             )?;
         }
@@ -740,7 +742,6 @@ impl<F: Field> AccountLeafConfig<F> {
             &account.wrong_rlp_bytes,
             &expected_item,
             true,
-            parent_data[1].is_extension,
             key_data[true.idx()].clone(),
             region.key_r,
         )?; 

zkevm-circuits/src/mpt_circuit/extension_branch.rs

@@ -26,6 +26,7 @@ pub(crate) struct ExtensionBranchConfig<F> {
     parent_data: [ParentData<F>; 2],
     is_placeholder: [Cell<F>; 2],
     is_extension: Cell<F>,
+    is_last_level: Cell<F>,
     extension: ExtensionGadget<F>,
     branch: BranchGadget<F>,
 }
@@ -41,6 +42,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
         circuit!([meta, cb], {
             // General inputs
             config.is_extension = cb.query_bool();
+            config.is_last_level = cb.query_bool();
             // If we're in a placeholder, both the extension and the branch parts are
             // placeholders
             for is_s in [true, false] {
@@ -159,6 +161,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                         false.expr(),
                         false.expr(),
                         config.is_extension.expr(),
+                        config.is_last_level.expr(),
                         WordLoHi::zero(),
                     );
                  } elsex {
@@ -186,6 +189,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                         config.parent_data[is_s.idx()].is_root.expr(),
                         true.expr(),
                         config.is_extension.expr(),
+                        config.is_last_level.expr(),
                         branch.mod_word[is_s.idx()].clone(),
                     );
                 }}
@@ -210,6 +214,9 @@ impl<F: Field> ExtensionBranchConfig<F> {
         let is_extension = extension_branch.is_extension.scalar();
         self.is_extension.assign(region, offset, is_extension)?;
 
+        let is_last_level = extension_branch.is_last_level.scalar();
+        self.is_last_level.assign(region, offset, is_last_level)?;
+
         let key_data =
             self.key_data
                 .witness_load(region, offset, &mut memory[key_memory(true)], 0)?;
@@ -294,6 +301,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                     false,
                     false,
                     is_extension == 1.into(),
+                    is_last_level == 1.into(),
                     WordLoHi::zero(),
                 )?;
             } else {
@@ -317,6 +325,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                     parent_data[is_s.idx()].is_root,
                     true,
                     is_extension == 1.into(),
+                    is_last_level == 1.into(),
                     mod_node_hash_word[is_s.idx()],
                 )?;
             }

zkevm-circuits/src/mpt_circuit/helpers.rs

@@ -558,7 +558,13 @@ pub(crate) struct ParentData<F> {
     pub(crate) rlc: Cell<F>,
     pub(crate) is_root: Cell<F>,
     pub(crate) is_placeholder: Cell<F>,
+    // is_extension is used only in a non-existing proof / wrong extension node case -
+    // in account/storage leaf to check whether the parent is an extension node
     pub(crate) is_extension: Cell<F>,
+    // is_ext_last_level is used only in a non-existing proof - in wrong extension node case the branch
+    // is a placeholder and the check for a branch hash being in the parent extension node needs to be ignored,
+    // but it needs to be ignored only in the last branch (the branch above the leaf into which the lookup is made)
+    pub(crate) is_ext_last_level: Cell<F>,
     pub(crate) drifted_parent_hash: WordLoHiCell<F>,
 }
 
@@ -569,6 +575,7 @@ pub(crate) struct ParentDataWitness<F> {
     pub(crate) is_root: bool,
     pub(crate) is_placeholder: bool,
     pub(crate) is_extension: bool,
+    pub(crate) is_ext_last_level: bool,
     pub(crate) drifted_parent_hash: WordLoHi<F>,
 }
 
@@ -584,6 +591,7 @@ impl<F: Field> ParentData<F> {
             is_root: cb.query_cell(),
             is_placeholder: cb.query_cell(),
             is_extension: cb.query_cell(),
+            is_ext_last_level: cb.query_cell(),
             drifted_parent_hash: cb.query_word_unchecked(),
         };
         circuit!([meta, cb.base], {
@@ -597,6 +605,7 @@ impl<F: Field> ParentData<F> {
                     parent_data.is_root.expr(),
                     parent_data.is_placeholder.expr(),
                     parent_data.is_extension.expr(),
+                    parent_data.is_ext_last_level.expr(),
                     parent_data.drifted_parent_hash.lo().expr(),
                     parent_data.drifted_parent_hash.hi().expr(),
                 ],
@@ -613,6 +622,7 @@ impl<F: Field> ParentData<F> {
         is_root: Expression<F>,
         is_placeholder: Expression<F>,
         is_extension: Expression<F>,
+        is_ext_last_level: Expression<F>,
         drifted_parent_hash: WordLoHi<Expression<F>>,
     ) {
         memory.store(
@@ -624,6 +634,7 @@ impl<F: Field> ParentData<F> {
                 is_root,
                 is_placeholder,
                 is_extension,
+                is_ext_last_level,
                 drifted_parent_hash.lo(),
                 drifted_parent_hash.hi(),
             ],
@@ -640,6 +651,7 @@ impl<F: Field> ParentData<F> {
         force_hashed: bool,
         is_placeholder: bool,
         is_extension: bool,
+        is_ext_last_level: bool,
         drifted_parent_hash: WordLoHi<F>,
     ) -> Result<(), Error> {
         memory.witness_store(
@@ -651,6 +663,7 @@ impl<F: Field> ParentData<F> {
                 force_hashed.scalar(),
                 is_placeholder.scalar(),
                 is_extension.scalar(),
+                is_ext_last_level.scalar(),
                 drifted_parent_hash.lo(),
                 drifted_parent_hash.hi(),
             ],
@@ -673,20 +686,22 @@ impl<F: Field> ParentData<F> {
         self.is_root.assign(region, offset, values[3])?;
         self.is_placeholder.assign(region, offset, values[4])?;
         self.is_extension.assign(region, offset, values[5])?;
+        self.is_ext_last_level.assign(region, offset, values[6])?;
         self.drifted_parent_hash
             .lo()
-            .assign(region, offset, values[6])?;
+            .assign(region, offset, values[7])?;
         self.drifted_parent_hash
             .hi()
-            .assign(region, offset, values[7])?;
+            .assign(region, offset, values[8])?;
 
         Ok(ParentDataWitness {
             hash: WordLoHi::new([values[0], values[1]]),
             rlc: values[2],
             is_root: values[3] == 1.scalar(),
             is_placeholder: values[4] == 1.scalar(),
             is_extension: values[5] == 1.scalar(),
-            drifted_parent_hash: WordLoHi::new([values[6], values[7]]),
+            is_ext_last_level: values[6] == 1.scalar(),
+            drifted_parent_hash: WordLoHi::new([values[7], values[8]]),
         })
     }
 }
@@ -1273,7 +1288,6 @@ impl<F: Field> WrongLeafGadget<F> {
         list_bytes: &[u8],
         expected_item: &RLPItemWitness,
         for_placeholder_s: bool,
-        is_parent_extension: bool,
         key_data: KeyDataWitness<F>,
         r: F,
     ) -> Result<(F, F), Error> {

zkevm-circuits/src/mpt_circuit/start.rs

@@ -68,6 +68,7 @@ impl<F: Field> StartConfig<F> {
                     true.expr(),
                     false.expr(),
                     false.expr(),
+                    false.expr(),
                     root[is_s.idx()].clone(),
                 );
                 KeyData::store_defaults(cb, &mut ctx.memory[key_memory(is_s)]);
@@ -123,6 +124,7 @@ impl<F: Field> StartConfig<F> {
                 true,
                 false,
                 false,
+                false,
                 root[is_s.idx()],
             )?;
             KeyData::witness_store(

zkevm-circuits/src/mpt_circuit/storage_leaf.rs

@@ -254,6 +254,7 @@ impl<F: Field> StorageLeafConfig<F> {
                     true.expr(),
                     false.expr(),
                     false.expr(),
+                    false.expr(),
                     WordLoHi::zero(),
                 );
             }
@@ -523,6 +524,7 @@ impl<F: Field> StorageLeafConfig<F> {
                 true,
                 false,
                 false,
+                false,
                 WordLoHi::<F>::new([F::ZERO, F::ZERO]),
             )?;
 
@@ -571,7 +573,6 @@ impl<F: Field> StorageLeafConfig<F> {
             &storage.wrong_rlp_bytes,
             &expected_item,
             false,
-            parent_data[1].is_extension,
             key_data[true.idx()].clone(),
             region.key_r,
         )?;