extNibbles returned by getProof has the same number as the proof now

geth-utils/gethutil/mpt/README.md

@@ -0,0 +1,33 @@
+# MPT witness generator
+
+## Generate witnesses
+
+To generate witnesses for the MPT circuit, execute
+
+```
+go test -v ./...
+```
+
+To generate the tests that use a local blockchain you need a local `geth`. You would
+need to run something like:
+```
+geth --dev --http --ipcpath ~/Library/Ethereum/geth.ipc
+
+```
+The local `geth` is used to generate some tests that have a small number of accounts so that
+these accounts appear in the first or second level of the trie. You might need to remove the
+database if you already have some accounts:
+
+```
+geth removedb
+```
+
+The witness files will appear in generated_witnesses folder.
+
+## Format the code
+
+To format the code use:
+
+```
+gofmt -w ./*
+```

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

@@ -107,11 +107,15 @@ func (t *Trie) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) (
 		// copy n.Key before it gets changed in ProofHash
 		var nCopy []byte
 		if short, ok := n.(*ShortNode); ok {
-			if !hasTerm(short.Key) { // only for extension keys
+			if !hasTerm(short.Key) { // extension keys
 				nCopy = make([]byte, len(short.Key))
 				copy(nCopy, short.Key)
 				extNibbles = append(extNibbles, nCopy)
+			} else {
+				extNibbles = append(extNibbles, []byte{})
 			}
+		} else {
+			extNibbles = append(extNibbles, []byte{})
 		}
 
 		n, _ = hasher.ProofHash(n)

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

@@ -204,10 +204,10 @@ func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
 
 // addBranchAndPlaceholder adds to the rows a branch and its placeholder counterpart
 // (used when one of the proofs have one branch more than the other).
-func addBranchAndPlaceholder(proof1, proof2,
-	extNibblesS, extNibblesC [][]byte,
+func addBranchAndPlaceholder(proof1, proof2 [][]byte,
+	extNibblesS, extNibblesC []byte,
 	leafRow0, key, neighbourNode []byte,
-	keyIndex, extensionNodeInd int,
+	keyIndex int,
 	additionalBranch, isAccountProof, nonExistingAccountProof,
 	isShorterProofLastLeaf bool, toBeHashed *[][]byte) (bool, bool, int, Node) {
 	len1 := len(proof1)
@@ -226,9 +226,9 @@ func addBranchAndPlaceholder(proof1, proof2,
 	if isExtension {
 		var numNibbles byte
 		if len1 > len2 {
-			numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesS, extensionNodeInd, proof1[len1-3], proof1[len1-3])
+			numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesS, proof1[len1-3], proof1[len1-3])
 		} else {
-			numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesC, extensionNodeInd, proof2[len2-3], proof2[len2-3])
+			numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesC, proof2[len2-3], proof2[len2-3])
 		}
 		numberOfNibbles = int(numNibbles)
 	}

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

@@ -1,6 +1,6 @@
 package witness
 
-func prepareExtensions(extNibbles [][]byte, extensionNodeInd int, proofEl1, proofEl2 []byte) (byte, []byte, [][]byte) {
+func prepareExtensions(extNibbles []byte, proofEl1, proofEl2 []byte) (byte, []byte, [][]byte) {
 	var values [][]byte
 	v1 := make([]byte, valueLen)
 	v2 := make([]byte, valueLen)
@@ -30,9 +30,9 @@ func prepareExtensions(extNibbles [][]byte, extensionNodeInd int, proofEl1, proo
 		}
 	}
 	ind := 0
-	for j := startNibblePos; j < len(extNibbles[extensionNodeInd]); j += 2 {
+	for j := startNibblePos; j < len(extNibbles); j += 2 {
 		v3[2+ind] = // TODO: check 2 + ind
-			extNibbles[extensionNodeInd][j]
+			extNibbles[j]
 		ind++
 	}
 	values = append(values, v1)

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

@@ -14,7 +14,7 @@ import (
 func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr common.Address, proof1, proof2,
 	extNibblesS, extNibblesC [][]byte,
 	key, neighbourNode []byte,
-	keyIndex, extensionNodeInd, numberOfNibbles int,
+	keyIndex, numberOfNibbles int,
 	additionalBranch, isAccountProof, nonExistingAccountProof,
 	isShorterProofLastLeaf bool, toBeHashed *[][]byte) Node {
 	len1 := len(proof1)
@@ -34,7 +34,7 @@ func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr c
 		extNibbles = extNibblesS
 	}
 
-	_, extListRlpBytesS, extValuesS := prepareExtensions(extNibbles, extensionNodeInd, longExtNode, longExtNode)
+	_, extListRlpBytesS, extValuesS := prepareExtensions(extNibbles[len(extNibbles)-1], longExtNode, longExtNode)
 
 	// Get nibbles of the extension node that gets shortened because of the newly insertd
 	// extension node:
@@ -107,7 +107,7 @@ func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr c
 		// Enable `prepareExtensionRows` call:
 		extNibbles = append(extNibbles, nibbles)
 
-		_, extListRlpBytesC, extValuesC = prepareExtensions(extNibbles, extensionNodeInd+1, shortExtNode, shortExtNode)
+		_, extListRlpBytesC, extValuesC = prepareExtensions(extNibbles[len(extNibbles)-1], shortExtNode, shortExtNode)
 	} else {
 		// When the short node is a branch (and not an extension node), we have nothing to be put in
 		// the C extension node witness (as a short node). We copy the long node (S extension node) to let

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

@@ -267,22 +267,6 @@ func obtainTwoProofsAndConvertToWitness(trieModifications []TrieModification, st
 			// Needs to be after `specialTest == 1` preparation:
 			nodes = append(nodes, GetStartNode(proofType, sRoot, cRoot, specialTest))
 
-			/*
-			if tMod.Type == StorageDoesNotExist {
-				fmt.Println("===================");
-				fmt.Println(addr);
-				fmt.Println(tMod.Key);
-				fmt.Println("");
-
-				for i := 0; i < len(storageProof); i++ {
-					fmt.Println(storageProof[i])
-					fmt.Println("")
-				}
-				fmt.Println("========");
-				fmt.Println("");
-			}
-			*/
-
 			// In convertProofToWitness, we can't use account address in its original form (non-hashed), because
 			// of the "special" test for which we manually manipulate the "hashed" address and we don't have a preimage.
 			// TODO: addr is used for calling GetProof for modified extension node only, might be done in a different way
@@ -380,8 +364,6 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 	}
 
 	var isExtension bool
-	extensionNodeInd := 0
-
 	var extListRlpBytes []byte
 	var extValues [][]byte
 	for i := 0; i < 4; i++ {
@@ -393,7 +375,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 	for i := 0; i < upTo; i++ {
 		if !isBranch(proof1[i]) {
 			isNonExistingProof := (isAccountProof && nonExistingAccountProof) || (!isAccountProof && nonExistingStorageProof)
-			areThereNibbles := len(extNibblesS) != 0 || len(extNibblesC) != 0
+			areThereNibbles := len(extNibblesS[i]) != 0 || len(extNibblesC[i]) != 0
 			// If i < upTo-1, it means it's not a leaf, so it's an extension node.
 			// There is no any special relation between isNonExistingProof and isExtension,
 			// except that in the non-existing proof the extension node can appear in `i == upTo-1`.
@@ -404,10 +386,9 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 			if (i != upTo-1) || (areThereNibbles && isNonExistingProof) { // extension node
 				var numberOfNibbles byte
 				isExtension = true
-				numberOfNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesS, extensionNodeInd, proof1[i], proof2[i])
+				numberOfNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesS[i], proof1[i], proof2[i])
 
 				keyIndex += int(numberOfNibbles)
-				extensionNodeInd++
 				continue
 			}
 
@@ -445,9 +426,10 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 				leafRow0 = proof2[len2-1]
 			}
 
-			isModifiedExtNode, _, numberOfNibbles, bNode := addBranchAndPlaceholder(proof1, proof2, extNibblesS, extNibblesC,
+			isModifiedExtNode, _, numberOfNibbles, bNode := addBranchAndPlaceholder(proof1, proof2,
+				extNibblesS[len1-1], extNibblesC[len2-1],
 				leafRow0, key, neighbourNode,
-				keyIndex, extensionNodeInd, additionalBranch,
+				keyIndex, additionalBranch,
 				isAccountProof, nonExistingAccountProof, isShorterProofLastLeaf, &toBeHashed)
 
 			nodes = append(nodes, bNode)
@@ -488,7 +470,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 			// modification).
 			if isModifiedExtNode {
 				leafNode = equipLeafWithModExtensionNode(statedb, leafNode, addr, proof1, proof2, extNibblesS, extNibblesC, key, neighbourNode,
-					keyIndex, extensionNodeInd, numberOfNibbles, additionalBranch,
+					keyIndex, numberOfNibbles, additionalBranch,
 					isAccountProof, nonExistingAccountProof, isShorterProofLastLeaf, &toBeHashed)
 			}
 			nodes = append(nodes, leafNode)
@@ -513,7 +495,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 				node := prepareStorageLeafPlaceholderNode(storage_key, key, keyIndex)
 				nodes = append(nodes, node)
 			}
-		} else if len(extNibblesC) > len(proof2)-1 {
+		} else {
 			isLastExtNode := len(extNibblesC[len(proof2)-1]) != 0
 			if isLastExtNode {
 				// We need to add a placeholder branch and a placeholder leaf.
@@ -564,7 +546,11 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 					newKey[keyIndex] = byte(i)
 					k := trie.HexToKeybytes(newKey)
 					ky := common.BytesToHash(k)
-					proof, _, _, _, _, err = statedb.GetStorageProof(addr, ky)
+					if isAccountProof {
+						proof, _, _, _, _, err = statedb.GetProof(addr)
+					} else {
+						proof, _, _, _, _, err = statedb.GetStorageProof(addr, ky)
+					}
 					check(err)
 					if !isBranch(proof[len(proof)-1]) {
 						break
@@ -586,6 +572,8 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 				copy(path, key[:l])
 				// The remaining `key` nibbles are to be stored in the constructed leaf - in our example [1 2 4 ...]
 
+				// TODO: construct for account proof
+
 				compact := trie.HexToCompact(key[l:])
 				// Add RLP:
 				compactLen := byte(len(compact))