fix all the cases except ext-mod node only

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

@@ -671,20 +671,20 @@ func isTxExt(proofEl []byte) bool {
 
 func printProof(ps [][]byte, t, idx []byte) {
 
-	enable := byte(200)
+	enable := byte(150)
 	fmt.Print(" [")
 	for i, p := range ps {
 		if t[i] == extNode {
 			fmt.Print("EXT - ")
-			if idx[0] > enable {
+			if idx[0] >= enable {
 				fmt.Print(" (", p, ") - ")
 			}
 		} else if t[i] == branchNode {
 			fmt.Print("BRANCH - ")
 			// fmt.Print(" (", p, ") - ")
 		} else if t[i] == leafNode {
 			fmt.Print("LEAF - ")
-			if idx[0] > enable {
+			if idx[0] >= enable {
 				fmt.Print(" (", p, ") - ")
 			}
 		} else if t[i] == hashedNode {
@@ -706,7 +706,6 @@ func (st *StackTrie) UpdateAndGetProof(db ethdb.KeyValueReader, indexBuf, value
 	if err != nil {
 		return StackProof{}, err
 	}
-	len1 := len(proofS)
 	printProof(proofS, typesS, indexBuf)
 
 	st.Update(indexBuf, value)
@@ -715,14 +714,8 @@ func (st *StackTrie) UpdateAndGetProof(db ethdb.KeyValueReader, indexBuf, value
 	if err != nil {
 		return StackProof{}, err
 	}
-	len2 := len(proofC)
 	printProof(proofC, typesC, indexBuf)
 
-	// fmt.Println(len1, len2)
-	if len1 > len2 {
-		fmt.Println(KeybytesToHex(indexBuf))
-	}
-
 	return StackProof{proofS, proofC, nibblesS, nibblesC, typesS, typesC}, nil
 }
 
@@ -806,7 +799,7 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, []
 		// fmt.Print(" ", k[i], "/", c.nodeType, " | ")
 		proofType = append(proofType, c.nodeType)
 		if c.nodeType == extNode {
-			fmt.Print(c.key, " ")
+			// fmt.Print(c.key, " ")
 			i += len(c.key) - 1
 			nodes = append(nodes, c)
 			c = c.children[0]
@@ -825,19 +818,18 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, []
 				panic(error)
 			}
 
-			if c_rlp[0] > 192 && c_rlp[0] < 248 {
-				numNibbles := c_rlp[0] - 225
-				var nibble = make([]byte, numNibbles)
-				for i := 0; i < int(numNibbles); i++ {
-					nibble[i] = c_rlp[i+1] - 16
-				}
-				// fmt.Println(" HASHED Ext nibble:", nibble, c_rlp)
-				nibbles = append(nibbles, nibble)
-			}
+			// if c_rlp[0] > 192 && c_rlp[0] < 248 {
+			// 	numNibbles := c_rlp[0] - 225
+			// 	var nibble = make([]byte, numNibbles)
+			// 	for i := 0; i < int(numNibbles); i++ {
+			// 		nibble[i] = c_rlp[i+1] - 16
+			// 	}
+			// 	// fmt.Println(" HASHED Ext nibble:", nibble, c_rlp)
+			// }
 
-			fmt.Print(" hashed node key:", c.key, " ")
 			// fmt.Println(" c_rlp:", c_rlp)
 			proof = append(proof, c_rlp)
+			nibbles = append(nibbles, c.key)
 			branchChild := st.getNodeFromBranchRLP(c_rlp, int(k[i]))
 			// branchChild is of length 1 when there is no child at this position in the branch
 			// (`branchChild = [128]` in this case), but it is also of length 1 when `c_rlp` is a leaf.
@@ -886,11 +878,11 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, []
 				// if it's an ext node, it contains 1.)nibbles and 2.) 32bytes hashed value
 				// 2.) 32 bytes long data plus rlp flag, it becomes 33 bytes long data
 				// 192 + 33 = 225, and the left bytes are for nibbles.
-				numNibbles := raw_rlp[0] - 225
-				var nibble = make([]byte, numNibbles)
-				for i := 0; i < int(numNibbles); i++ {
-					nibble[i] = raw_rlp[i+1] - 16
-				}
+				// numNibbles := raw_rlp[0] - 225
+				// var nibble = make([]byte, numNibbles)
+				// for i := 0; i < int(numNibbles); i++ {
+				// 	nibble[i] = raw_rlp[i+1] - 16
+				// }
 				// fmt.Println(" Ext nibble:", numNibbles, node.key)
 				nibbles = append(nibbles, node.key)
 			} else {

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

@@ -1,7 +1,6 @@
 package witness
 
 import (
-	"fmt"
 	"log"
 
 	"github.com/ethereum/go-ethereum/rlp"
@@ -237,11 +236,12 @@ 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 [][]byte, extNibblesS, extNibblesC []byte,
-	leafRow0, key []byte, keyIndex int, isShorterProofLastLeaf bool,
+	proofTx, leafRow0, key []byte, keyIndex int, isShorterProofLastLeaf bool,
 ) (bool, bool, int, Node) {
 
 	len1 := len(proof1)
 	len2 := len(proof2)
+	isTxProof := len(proofTx) != 0
 
 	var node Node
 
@@ -252,22 +252,14 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC [
 		extValues = append(extValues, make([]byte, valueLen))
 	}
 
-	// For stack trie
-	// if 1 st node of proof2 is a branch node and 1st node of Proof1 is an ext node
-	need_placeholder_ext := isBranch(proof2[0]) && isTxExt(proof1[0])
-	if need_placeholder_ext {
-		fmt.Println("need_placeholder_ext", isTxLeaf(proof1[0]), isBranch(proof1[0]), proof1[0])
-		fmt.Println("need_placeholder_ext", isBranch(proof2[0]), proof2[0])
-		fmt.Println("nibble:", extNibblesS)
-	}
-	isExtension := (len1 == len2+2) || (len2 == len1+2)
-	if isExtension || need_placeholder_ext {
+	isExtension := (len1 == len2+2) || (len2 == len1+2) || (isTxProof && !isBranch(proofTx))
+	if isExtension {
 		var numNibbles byte
 		var proof []byte
 		var extNibbles []byte
-		if need_placeholder_ext {
+		if isTxProof {
 			extNibbles = extNibblesS
-			proof = proof1[0]
+			proof = proofTx
 		} else {
 			if len1 > len2 {
 				extNibbles = extNibblesS
@@ -306,18 +298,19 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC [
 		in the rows before we add a leaf.
 	*/
 	var longExtNode []byte
-	if len1 > len2 {
-		longExtNode = proof2[len2-1]
+	if isTxProof {
+		longExtNode = proofTx
 	} else {
-		longExtNode = proof1[len1-1]
+		if len1 > len2 {
+			longExtNode = proof2[len2-1]
+		} else {
+			longExtNode = proof1[len1-1]
+		}
 	}
 
-	// TODO: fix
 	var extNode []byte
-	if need_placeholder_ext {
-		extNode = proof1[0]
-		// FIXME should move to above and need to avoid above [len-3] operation
-		isExtension = need_placeholder_ext
+	if isTxProof {
+		extNode = proofTx
 	} else {
 		if isExtension {
 			if len1 > len2 {
@@ -329,7 +322,7 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC [
 	}
 
 	// Note that isModifiedExtNode happens also when we have a branch instead of shortExtNode
-	isModifiedExtNode := (!isBranch(longExtNode) && !isShorterProofLastLeaf) || need_placeholder_ext
+	isModifiedExtNode := (!isBranch(longExtNode) && !isShorterProofLastLeaf)
 
 	// 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

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

@@ -121,8 +121,9 @@ func transactionsStackTrieInsertionTemplate(t *testing.T, n int) {
 }
 
 func TestTransactionInsertion(t *testing.T) {
-	// TODO: check drifted_index at 128
-	txs := makeTransactions(130)
+	// potential ext. mod nodes
+	// 145, 513, 529, 785, 1041, 1297, 1553, 1809
+	txs := makeTransactions(2015)
 	prepareStackTrieWitness("TransactionInsertion", types.Transactions(txs))
 }