Wrong extension node: three parts of the key adds up to key rlc

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

@@ -15,11 +15,11 @@ func prepareExtensions(extNibbles []byte, proofEl1, proofEl2 []byte) (byte, []by
 	numberOfNibbles := getExtensionNumberOfNibbles(proofEl1)
 
 	// We need nibbles as witness to compute key RLC, so we set them
-	// into extensionRowC s_advices (we can do this because both extension
+	// into extensionRowC (we can do this because both extension
 	// nodes have the same key, so we can have this info only in one).
 	// There can be more up to 64 nibbles, but there is only 32 bytes
-	// in extensionRowC s_advices. So we store every second nibble (having
-	// the whole byte and one nibble is enough to compute the other nibble).
+	// in extensionRowC. So we store every second nibble (having
+	// the byte and one nibble is enough to compute the other nibble).
 
 	startNibblePos := 2 // we don't need any nibbles for case keyLen = 1
 	if keyLen > 1 {
@@ -31,8 +31,7 @@ func prepareExtensions(extNibbles []byte, proofEl1, proofEl2 []byte) (byte, []by
 	}
 	ind := 0
 	for j := startNibblePos; j < len(extNibbles); j += 2 {
-		v3[2+ind] = // TODO: check 2 + ind
-			extNibbles[j]
+		v3[2+ind] = extNibbles[j]
 		ind++
 	}
 	values = append(values, v1)

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

@@ -580,13 +580,41 @@ func prepareStorageLeafNode(leafS, leafC, constructedLeaf, neighbourNode []byte,
 	return node
 }
 
-func equipLeafWithWrongExtension(leafNode Node, keyBefore, keyMiddle, keyAfter []byte) Node {
+func equipLeafWithWrongExtension(leafNode Node, keyBefore, keyMiddle, keyAfter, nibblesMiddle, nibblesAfter []byte) Node {
 	l := len(leafNode.Values)
 	// The rows which are used for nibbles (it means there are checks for values being < 16)
 	// in the modified extension node cases are omitted.
 	leafNode.Values[l-modifiedExtensionNodeRowLen - 1] = keyBefore
+
 	leafNode.Values[l-modifiedExtensionNodeRowLen] = keyMiddle
+	startNibblePos := 2 // we don't need any nibbles for case keyLen = 1
+	if len(keyMiddle) > 1 {
+		if len(nibblesMiddle) % 2 == 0 {
+			startNibblePos = 1
+		} else {
+			startNibblePos = 2
+		}
+	}
+	ind := 0
+	for j := startNibblePos; j < len(nibblesMiddle); j += 2 {
+		leafNode.Values[l-modifiedExtensionNodeRowLen + 1][2+ind] = nibblesMiddle[j]
+		ind++
+	}
+
 	leafNode.Values[l-modifiedExtensionNodeRowLen + 3] = keyAfter
+	startNibblePos = 2 // we don't need any nibbles for case keyLen = 1
+	if len(keyAfter) > 1 {
+		if len(nibblesAfter) % 2 == 0 {
+			startNibblePos = 1
+		} else {
+			startNibblePos = 2
+		}
+	}
+	ind = 0
+	for j := startNibblePos; j < len(nibblesAfter); j += 2 {
+		leafNode.Values[l-modifiedExtensionNodeRowLen + 4][2+ind] = nibblesAfter[j]
+		ind++
+	}
 
 	return leafNode
 }

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

@@ -1,7 +1,6 @@
 package witness
 
 import (
-	"fmt"
 	"math/big"
 
 	"main/gethutil/mpt/oracle"
@@ -517,22 +516,19 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 
 				start := keyIndex - len(nibbles)
 
-				before := newKey[0:start]
-				if len(before) % 2 == 1 {
-					before = append(before, 0)
+				nibblesBefore := newKey[0:start]
+				if len(nibblesBefore) % 2 == 1 {
+					nibblesBefore = append(nibblesBefore, 0)
 				}
-				keyBefore := trie.HexToKeybytes(before)
-				fmt.Println(keyBefore)
+				keyBefore := trie.HexToKeybytes(nibblesBefore)
 
-				middle := newKey[start:keyIndex]
+				nibblesMiddle := newKey[start:keyIndex]
 				keyMiddle := []byte{160}
-				keyMiddle = append(keyMiddle, trie.HexToCompact(middle)...)
-				fmt.Println(keyMiddle)
+				keyMiddle = append(keyMiddle, trie.HexToCompact(nibblesMiddle)...)
 
-				after := newKey[keyIndex:]
+				nibblesAfter := newKey[keyIndex:len(newKey)-1] // Remove the last flag (16)
 				keyAfter := []byte{160}
-				keyAfter = append(keyAfter, trie.HexToCompact(after)...)
-				fmt.Println(keyAfter)
+				keyAfter = append(keyAfter, trie.HexToCompact(nibblesAfter)...)
 
 				isExtension := true
 				// Dummy branch. The constraint of the branch being in the extension node and the constraint of the
@@ -549,7 +545,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 					dummyLeaf := []byte{248,108,157,52,45,53,199,120,18,165,14,109,22,4,141,198,233,128,219,44,247,218,241,231,2,206,125,246,58,246,15,3,184,76,248,74,4,134,85,156,208,108,8,0,160,86,232,31,23,27,204,85,166,255,131,69,230,146,192,248,110,91,72,224,27,153,108,173,192,1,98,47,181,227,99,180,33,160,197,210,70,1,134,247,35,60,146,126,125,178,220,199,3,192,229,0,182,83,202,130,39,59,123,250,216,4,93,133,164,112}
 					node := prepareAccountLeafNode(addr, addrh, dummyLeaf, dummyLeaf, dummyLeaf, nil, addr_nibbles, false, false, false)
 
-					node = equipLeafWithWrongExtension(node, keyBefore, keyMiddle, keyAfter)
+					node = equipLeafWithWrongExtension(node, keyBefore, keyMiddle, keyAfter, nibblesMiddle, nibblesAfter)
 
 					nodes = append(nodes, node)
 				} else {

zkevm-circuits/src/mpt_circuit/account_leaf.rs

@@ -1,5 +1,5 @@
 use eth_types::{Field, OpsIdentity, U256};
-use gadgets::util::{pow, Scalar};
+use gadgets::util::{pow, xor, Scalar};
 use halo2_proofs::{
     circuit::Value,
     plonk::{Error, Expression, VirtualCells},
@@ -16,17 +16,15 @@ use crate::{
     circuit_tools::{
         cached_region::CachedRegion,
         cell_manager::Cell,
-        constraint_builder::{RLCChainableRev, RLCable},
+        constraint_builder::{RLCChainable, RLCChainableRev, RLCable},
         gadgets::IsEqualGadget,
     },
     evm_circuit::util::from_bytes,
     mpt_circuit::{
         helpers::{
-            key_memory, main_memory, num_nibbles, parent_memory, DriftedGadget, Indexable,
-            IsPlaceholderLeafGadget, KeyData, MPTConstraintBuilder, ParentData, WrongGadget,
-            KECCAK,
+            ext_key_rlc_expr, key_memory, leaf_key_rlc, main_memory, num_nibbles, parent_memory, DriftedGadget, Indexable, IsPlaceholderLeafGadget, KeyData, MPTConstraintBuilder, ParentData, WrongGadget, KECCAK
         },
-        param::{EMPTY_TRIE_HASH, KEY_LEN_IN_NIBBLES, RLP_LIST_LONG, RLP_LONG},
+        param::{EMPTY_TRIE_HASH, KEY_LEN_IN_NIBBLES, KEY_PREFIX_ODD, RLP_LIST_LONG, RLP_LONG},
         MPTConfig, MPTContext, MptMemory, RlpItemType,
     },
     table::MPTProofType,
@@ -368,7 +366,7 @@ impl<F: Field> AccountLeafConfig<F> {
                 &cb.key_r.expr(),
             );
 
-            // Wrong leaf / extension node handling
+            // Wrong leaf handling
             config.wrong = WrongGadget::construct(
                 cb,
                 key_item.hash_rlc(),
@@ -383,30 +381,72 @@ impl<F: Field> AccountLeafConfig<F> {
                 &cb.key_r.expr(),
             );
 
-            // TODO: wrong extension node
-            /*
-            let wrong_bytes =
-                ctx.rlp_item(meta, cb, AccountRowType::Wrong as usize, RlpItemType::Key);
-            */
+            // TODO: wrong extension node gadget
             let wrong_ext_middle =
                 ctx.rlp_item(meta, cb, AccountRowType::LongExtNodeKey as usize, RlpItemType::Key);
+            let wrong_ext_middle_nibbles =
+                ctx.rlp_item(meta, cb, AccountRowType::LongExtNodeNibbles as usize, RlpItemType::Nibbles);
             let wrong_ext_after =
                 ctx.rlp_item(meta, cb, AccountRowType::ShortExtNodeKey as usize, RlpItemType::Key);
+            let wrong_ext_after_nibbles =
+                ctx.rlp_item(meta, cb, AccountRowType::ShortExtNodeNibbles as usize, RlpItemType::Nibbles);
+
+            // In the wrong extension node, AccountRowType::Wrong stores the bytes of the key nibbles
+            // up until the extension node.
+            let (mut rlc, _) = wrong_bytes.rlc_chain_data(); 
+            // The nibbles stored in the Wrong row (up until the extension node)
+            // need to be the same as the nibbles in the path.
+            require!(rlc => config.key_data_prev.rlc.expr());
+
+            // We have a key split into three parts in
+            // the wrong extension node case, meaning that there the first part parity doesn't
+            // tell us about the parity of the second part (depends on the third part as well).
+
+            let data0 = [wrong_ext_middle.clone(), wrong_ext_middle_nibbles.clone()];
+            rlc = rlc
+                + ext_key_rlc_expr(
+                    cb,
+                    wrong_ext_middle,
+                    config.key_data_prev.mult.expr(),
+                    config.key_data[1].is_odd.expr(),
+                    config.key_data_prev.is_odd.expr(),
+                    data0
+                        .iter()
+                        .map(|item| item.bytes_be())
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
+                    &cb.key_r.expr(),
+                );
+
+            // odd odd -> even
+            // odd even -> odd
+            // even odd -> odd
+            // even even -> even
+            let after_two_parts_is_odd =
+                xor::expr(config.key_data_prev.is_odd.expr(), config.key_data[1].is_odd.expr());
 
-            let (before, before_mult) = wrong_bytes.rlc_chain_data(); 
-            require!(before => config.key_data_prev.rlc.expr());
+            // The total number of nibbles is odd, thus:
+            let third_part_is_odd = after_two_parts_is_odd.clone();
 
-            /*
-            wrong_ext_before = key_data[is_s.idx()].rlc.expr()
-                + rlp_key.key.expr(
+            let data1 = [wrong_ext_after.clone(), wrong_ext_after_nibbles.clone()];
+            rlc = rlc
+                + ext_key_rlc_expr(
                     cb,
-                    rlp_key.key_value.clone(),
-                    key_data[is_s.idx()].mult.expr(),
-                    key_data[is_s.idx()].is_odd.expr(),
+                    wrong_ext_after,
+                    config.key_data[1].mult.expr(),
+                    third_part_is_odd,
+                    after_two_parts_is_odd,
+                    data1
+                        .iter()
+                        .map(|item| item.bytes_be())
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
                     &cb.key_r.expr(),
                 );
-            */
 
+            require!(key_item.hash_rlc() => rlc);
 
             // Anything following this node is below the account
             // TODO(Brecht): For non-existing accounts it should be impossible to prove