is_extension added to ParentData (for wrong extension data handling);…

… constructedLeaf in witness generator used now only for wrong row

geth-utils/gethutil/mpt/state/statedb.go

@@ -457,7 +457,7 @@ func (s *StateDB) SetStateObjectIfExists(addr common.Address) {
 			}
 
 			/*
-				When an account that does not exist is tried to be fetched by PrefetchAccount and when the some other account
+				When an account that does not exist is being fetched by PrefetchAccount and when some other account
 				exist at the overlapping address (the beginning of it), this (wrong) account is obtained by PrefetchAccount
 				and needs to be ignored.
 			*/

geth-utils/gethutil/mpt/trie/proof.go

@@ -102,7 +102,6 @@ func (t *Trie) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) (
 			fromLevel--
 			continue
 		}
-		// var hn Node
 
 		// We need nibbles in witness for extension keys.
 		// copy n.Key before it gets changed in ProofHash
@@ -115,20 +114,11 @@ func (t *Trie) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) (
 			}
 		}
 
-		// n, hn = hasher.ProofHash(n)
 		n, _ = hasher.ProofHash(n)
-		// if hash, ok := hn.(HashNode); ok || i == 0 {
 		// If the node's database encoding is a hash (or is the
 		// root node), it becomes a proof element.
 		enc, _ := rlp.EncodeToBytes(n)
-		/*
-			if !ok {
-				hash = hasher.HashData(enc)
-			}
-		*/
-		// proofDb.Put(hash, enc)
 		proofDb.Put([]byte{1, 1, 1}, enc)
-		// }
 	}
 
 	isNeighbourNodeHashed := false

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

@@ -341,7 +341,7 @@ func prepareLeafAndPlaceholderNode(addr common.Address, addrh []byte, proof1, pr
 			isSPlaceholder = true
 		}
 
-		return prepareStorageLeafNode(leaf, leaf, nil, storage_key, key, false, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension)
+		return prepareStorageLeafNode(leaf, leaf, nil, nil, storage_key, key, false, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension)
 	}
 }
 
@@ -416,7 +416,7 @@ func prepareStorageLeafPlaceholderNode(storage_key common.Hash, key []byte, keyI
 	keyLen := getLeafKeyLen(keyIndex)
 	leaf[0] = 192 + 1 + byte(keyLen) + 1
 
-	return prepareStorageLeafNode(leaf, leaf, nil, storage_key, key, false, true, true, false, false)
+	return prepareStorageLeafNode(leaf, leaf, nil, nil, storage_key, key, false, true, true, false, false)
 }
 
 func prepareStorageLeafInfo(row []byte, valueIsZero, isPlaceholder bool) ([]byte, []byte, []byte, []byte) {
@@ -502,7 +502,7 @@ func prepareStorageLeafInfo(row []byte, valueIsZero, isPlaceholder bool) ([]byte
 	return key, value, keyRlp, valueRlp
 }
 
-func prepareStorageLeafNode(leafS, leafC, neighbourNode []byte, storage_key common.Hash, key []byte, nonExistingStorageProof, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension bool) Node {
+func prepareStorageLeafNode(leafS, leafC, constructedLeaf, neighbourNode []byte, storage_key common.Hash, key []byte, nonExistingStorageProof, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension bool) Node {
 	var rows [][]byte
 
 	keyS, valueS, listRlpBytes1, valueRlpBytes1 := prepareStorageLeafInfo(leafS, false, isSPlaceholder)
@@ -533,7 +533,11 @@ func prepareStorageLeafNode(leafS, leafC, neighbourNode []byte, storage_key comm
 	var nonExistingStorageRow []byte
 	var wrongRlpBytes []byte
 	if nonExistingStorageProof {
-		wrongRlpBytes, nonExistingStorageRow = prepareNonExistingStorageRow(leafC, key)
+		if constructedLeaf != nil {
+			wrongRlpBytes, nonExistingStorageRow = prepareNonExistingStorageRow(constructedLeaf, key)
+		} else {
+			wrongRlpBytes, nonExistingStorageRow = prepareNonExistingStorageRow(leafC, key)
+		}
 	} else {
 		nonExistingStorageRow = prepareEmptyNonExistingStorageRow()
 	}
@@ -555,6 +559,7 @@ func prepareStorageLeafNode(leafS, leafC, neighbourNode []byte, storage_key comm
 		ValueRlpBytes:   valueRlpBytes,
 		IsModExtension:  [2]bool{isSModExtension, isCModExtension},
 	}
+
 	keccakData := [][]byte{leafS, leafC, storage_key.Bytes()}
 	if neighbourNode != nil {
 		keccakData = append(keccakData, neighbourNode)

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

@@ -412,7 +412,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 			if isAccountProof {
 				node = prepareAccountLeafNode(addr, addrh, proof1[l-1], proof2[l-1], nil, key, false, false, false)
 			} else {
-				node = prepareStorageLeafNode(proof1[l-1], proof2[l-1], nil, storage_key, key, nonExistingStorageProof, false, false, false, false)
+				node = prepareStorageLeafNode(proof1[l-1], proof2[l-1], nil, nil, storage_key, key, nonExistingStorageProof, false, false, false, false)
 			}
 
 			nodes = append(nodes, node)
@@ -466,7 +466,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 			} else {
 				// Add storage leaf after branch placeholder
 				if !isModifiedExtNode {
-					leafNode = prepareStorageLeafNode(proof1[len1-1], proof2[len2-1], neighbourNode, storage_key, key, nonExistingStorageProof, false, false, false, false)
+					leafNode = prepareStorageLeafNode(proof1[len1-1], proof2[len2-1], nil, neighbourNode, storage_key, key, nonExistingStorageProof, false, false, false, false)
 				} else {
 					isSModExtension := false
 					isCModExtension := false
@@ -589,7 +589,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 					// Add dummy value:
 					constructedLeaf = append(constructedLeaf, 1)
 
-					node := prepareStorageLeafNode(proof[len(proof)-1], constructedLeaf, nil, storage_key, key, nonExistingStorageProof, false, false, false, false)
+					node := prepareStorageLeafNode(proof[len(proof)-1], proof[len(proof)-1], constructedLeaf, nil, storage_key, key, nonExistingStorageProof, false, false, false, false)
 					nodes = append(nodes, node)
 				}
 			}

zkevm-circuits/src/mpt_circuit/account_leaf.rs

@@ -335,6 +335,7 @@ impl<F: Field> AccountLeafConfig<F> {
                     0.expr(),
                     true.expr(),
                     false.expr(),
+                    false.expr(),
                     storage_items[is_s.idx()].word(),
                 );
             }
@@ -372,6 +373,7 @@ impl<F: Field> AccountLeafConfig<F> {
                 &key_rlc[true.idx()],
                 &wrong_bytes,
                 config.is_placeholder_leaf[true.idx()].expr(),
+                config.parent_data[true.idx()].is_extension.expr(),
                 config.key_data[true.idx()].clone(),
                 &cb.key_r.expr(),
             );
@@ -641,6 +643,7 @@ impl<F: Field> AccountLeafConfig<F> {
                 0.scalar(),
                 true,
                 false,
+                false,
                 storage_items[is_s.idx()].word(),
             )?;
         }

zkevm-circuits/src/mpt_circuit/extension_branch.rs

@@ -158,6 +158,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                         branch.mod_rlc[is_s.idx()].expr(),
                         false.expr(),
                         false.expr(),
+                        config.is_extension.expr(),
                         WordLoHi::zero(),
                     );
                  } elsex {
@@ -184,6 +185,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                         config.parent_data[is_s.idx()].rlc.expr(),
                         config.parent_data[is_s.idx()].is_root.expr(),
                         true.expr(),
+                        config.is_extension.expr(),
                         branch.mod_word[is_s.idx()].clone(),
                     );
                 }}
@@ -205,8 +207,8 @@ impl<F: Field> ExtensionBranchConfig<F> {
     ) -> Result<(), Error> {
         let extension_branch = &node.extension_branch.clone().unwrap();
 
-        self.is_extension
-            .assign(region, offset, extension_branch.is_extension.scalar())?;
+        let is_extension = extension_branch.is_extension.scalar();
+        self.is_extension.assign(region, offset, is_extension)?;
 
         let key_data =
             self.key_data
@@ -291,6 +293,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                     mod_node_hash_rlc[is_s.idx()],
                     false,
                     false,
+                    is_extension == 1.into(),
                     WordLoHi::zero(),
                 )?;
             } else {
@@ -313,6 +316,7 @@ impl<F: Field> ExtensionBranchConfig<F> {
                     parent_data[is_s.idx()].rlc,
                     parent_data[is_s.idx()].is_root,
                     true,
+                    is_extension == 1.into(),
                     mod_node_hash_word[is_s.idx()],
                 )?;
             }

zkevm-circuits/src/mpt_circuit/helpers.rs

@@ -558,6 +558,7 @@ pub(crate) struct ParentData<F> {
     pub(crate) rlc: Cell<F>,
     pub(crate) is_root: Cell<F>,
     pub(crate) is_placeholder: Cell<F>,
+    pub(crate) is_extension: Cell<F>,
     pub(crate) drifted_parent_hash: WordLoHiCell<F>,
 }
 
@@ -567,6 +568,7 @@ pub(crate) struct ParentDataWitness<F> {
     pub(crate) rlc: F,
     pub(crate) is_root: bool,
     pub(crate) is_placeholder: bool,
+    pub(crate) is_extension: bool,
     pub(crate) drifted_parent_hash: WordLoHi<F>,
 }
 
@@ -581,6 +583,7 @@ impl<F: Field> ParentData<F> {
             rlc: cb.query_cell_with_type(MptCellType::StoragePhase2),
             is_root: cb.query_cell(),
             is_placeholder: cb.query_cell(),
+            is_extension: cb.query_cell(),
             drifted_parent_hash: cb.query_word_unchecked(),
         };
         circuit!([meta, cb.base], {
@@ -593,6 +596,7 @@ impl<F: Field> ParentData<F> {
                     parent_data.rlc.expr(),
                     parent_data.is_root.expr(),
                     parent_data.is_placeholder.expr(),
+                    parent_data.is_extension.expr(),
                     parent_data.drifted_parent_hash.lo().expr(),
                     parent_data.drifted_parent_hash.hi().expr(),
                 ],
@@ -608,6 +612,7 @@ impl<F: Field> ParentData<F> {
         rlc: Expression<F>,
         is_root: Expression<F>,
         is_placeholder: Expression<F>,
+        is_extension: Expression<F>,
         drifted_parent_hash: WordLoHi<Expression<F>>,
     ) {
         memory.store(
@@ -618,6 +623,7 @@ impl<F: Field> ParentData<F> {
                 rlc,
                 is_root,
                 is_placeholder,
+                is_extension,
                 drifted_parent_hash.lo(),
                 drifted_parent_hash.hi(),
             ],
@@ -633,6 +639,7 @@ impl<F: Field> ParentData<F> {
         rlc: F,
         force_hashed: bool,
         is_placeholder: bool,
+        is_extension: bool,
         drifted_parent_hash: WordLoHi<F>,
     ) -> Result<(), Error> {
         memory.witness_store(
@@ -643,6 +650,7 @@ impl<F: Field> ParentData<F> {
                 rlc,
                 force_hashed.scalar(),
                 is_placeholder.scalar(),
+                is_extension.scalar(),
                 drifted_parent_hash.lo(),
                 drifted_parent_hash.hi(),
             ],
@@ -664,19 +672,21 @@ impl<F: Field> ParentData<F> {
         self.rlc.assign(region, offset, values[2])?;
         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.drifted_parent_hash
             .lo()
-            .assign(region, offset, values[5])?;
+            .assign(region, offset, values[6])?;
         self.drifted_parent_hash
             .hi()
-            .assign(region, offset, values[6])?;
+            .assign(region, offset, values[7])?;
 
         Ok(ParentDataWitness {
             hash: WordLoHi::new([values[0], values[1]]),
             rlc: values[2],
             is_root: values[3] == 1.scalar(),
             is_placeholder: values[4] == 1.scalar(),
-            drifted_parent_hash: WordLoHi::new([values[5], values[6]]),
+            is_extension: values[5] == 1.scalar(),
+            drifted_parent_hash: WordLoHi::new([values[6], values[7]]),
         })
     }
 }
@@ -1220,13 +1230,14 @@ impl<F: Field> WrongGadget<F> {
         key_rlc: &Expression<F>,
         expected_item: &RLPItemView<F>,
         is_placeholder: Expression<F>,
+        is_parent_extension: Expression<F>,
         key_data: KeyData<F>,
         r: &Expression<F>,
     ) -> Self {
         let mut config = WrongGadget::default();
         circuit!([meta, cb.base], {
             // Get the previous key data
-            ifx! {(is_non_existing, not!(is_placeholder)) => {
+            ifx! {and::expr(&[is_non_existing, not!(is_placeholder), not!(is_parent_extension)]) => {
                 // Calculate the key
                 config.wrong_rlp_key = ListKeyGadget::construct(cb, expected_item);
                 let key_rlc_wrong = key_data.rlc.expr() + config.wrong_rlp_key.key.expr(

zkevm-circuits/src/mpt_circuit/start.rs

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

zkevm-circuits/src/mpt_circuit/storage_leaf.rs

@@ -250,6 +250,7 @@ impl<F: Field> StorageLeafConfig<F> {
                     0.expr(),
                     true.expr(),
                     false.expr(),
+                    false.expr(),
                     WordLoHi::zero(),
                 );
             }
@@ -291,6 +292,7 @@ impl<F: Field> StorageLeafConfig<F> {
                 &key_rlc[true.idx()],
                 &expected_item,
                 config.is_placeholder_leaf[true.idx()].expr(),
+                config.parent_data[true.idx()].is_extension.expr(),
                 config.key_data[true.idx()].clone(),
                 &cb.key_r.expr(),
             );
@@ -517,6 +519,7 @@ impl<F: Field> StorageLeafConfig<F> {
                 F::ZERO,
                 true,
                 false,
+                false,
                 WordLoHi::<F>::new([F::ZERO, F::ZERO]),
             )?;