feat: working (only for some cases) but ugly version

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

@@ -230,6 +230,7 @@ func (st *StackTrie) getDiffIndex(key []byte) int {
 
 // Helper function to that inserts a (key, value) pair into
 // the trie.
+// https://github.dev/ethereum/go-ethereum/blob/00905f7dc406cfb67f64cd74113777044fb886d8/core/types/hashing.go#L105-L134
 func (st *StackTrie) insert(key, value []byte) {
 	switch st.nodeType {
 	case branchNode: /* Branch */
@@ -619,8 +620,10 @@ func (st *StackTrie) getNodeFromBranchRLP(branch []byte, idx int) []byte {
 }
 
 type StackProof struct {
-	proofS [][]byte
-	proofC [][]byte
+	proofS   [][]byte
+	proofC   [][]byte
+	nibblesS [][]byte
+	nibblesC [][]byte
 }
 
 func (sp *StackProof) GetProofS() [][]byte {
@@ -630,21 +633,61 @@ func (sp *StackProof) GetProofS() [][]byte {
 func (sp *StackProof) GetProofC() [][]byte {
 	return sp.proofC
 }
+func (sp *StackProof) GetNibblesS() [][]byte {
+	return sp.nibblesS
+}
+func (sp *StackProof) GetNibblesC() [][]byte {
+	return sp.nibblesC
+}
+
+func printProof(ps [][]byte) {
+
+	fmt.Print(" ")
+	for _, p := range ps {
+		if p[0] == 226 && p[1] == 16 && p[2] == 160 {
+			fmt.Print("EXT - ")
+		} else if p[0] == 248 || p[0] == 249 {
+			offset := p[0] - 248 + 1
+			if ((p[offset]-81)%32 == 0) && (p[offset+1] == 128 || p[offset+1] == 160) {
+				fmt.Print("BRANCH - ")
+			} else {
+				fmt.Print("LEAF - ")
+				// fmt.Println("  ", p)
+			}
+		} else {
+			fmt.Print("LEAF -")
+			fmt.Println("  ", p)
+		}
+	}
+	fmt.Println("  ")
+
+}
 
 func (st *StackTrie) UpdateAndGetProof(db ethdb.KeyValueReader, indexBuf, value []byte) (StackProof, error) {
-	proofS, err := st.GetProof(db, indexBuf)
+	fmt.Println("====")
+	proofS, nibblesS, err := st.GetProof(db, indexBuf)
 	if err != nil {
 		return StackProof{}, err
 	}
+	len1 := len(proofS)
+	printProof(proofS)
 
 	st.Update(indexBuf, value)
 
-	proofC, err := st.GetProof(db, indexBuf)
+	proofC, nibblesC, err := st.GetProof(db, indexBuf)
 	if err != nil {
 		return StackProof{}, err
 	}
+	len2 := len(proofC)
+	fmt.Println(" Proof S C ", len1, len2)
+	printProof(proofC)
 
-	return StackProof{proofS, proofC}, nil
+	// fmt.Println(len1, len2)
+	if len1 >= len2 {
+		fmt.Println(KeybytesToHex(indexBuf))
+	}
+
+	return StackProof{proofS, proofC, nibblesS, nibblesC}, nil
 }
 
 // We refer to the link below for this function.
@@ -690,14 +733,16 @@ func (st *StackTrie) UpdateAndGetProofs(db ethdb.KeyValueReader, list types.Deri
 		}
 		proofs = append(proofs, proof)
 	}
+	// fmt.Println("* ROOT", root)
 
 	return proofs, nil
 }
 
-func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, error) {
+func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, [][]byte, error) {
 	k := KeybytesToHex(key)
+	fmt.Println("k", k)
 	if st.nodeType == emptyNode {
-		return [][]byte{}, nil
+		return [][]byte{}, nil, nil
 	}
 
 	// Note that when root is a leaf, this leaf should be returned even if you ask for a different key (than the key of
@@ -709,9 +754,10 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, er
 	// in the S proof (another reason is that the S proof with a placeholder leaf would be an empty trie and thus with
 	// a root of an empty trie - which is not the case in S proof).
 	if st.nodeType == leafNode {
-		return [][]byte{st.val}, nil
+		return [][]byte{st.val}, nil, nil
 	}
 
+	var nibbles [][]byte
 	var proof [][]byte
 	var nodes []*StackTrie
 	c := st
@@ -760,7 +806,6 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, er
 		lNodes := len(nodes)
 		for i := lNodes - 1; i >= 0; i-- {
 			node := nodes[i]
-
 			if node.nodeType == leafNode {
 				rlp, error := db.Get(node.val)
 				if error != nil { // TODO: avoid error when RLP
@@ -771,11 +816,26 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, er
 			} else if node.nodeType == branchNode || node.nodeType == extNode {
 				node.hash(false)
 
-				rlp, error := db.Get(node.val)
+				raw_rlp, error := db.Get(node.val)
 				if error != nil {
-					return nil, error
+					return nil, nil, error
+				}
+				proof = append(proof, raw_rlp)
+				if node.nodeType == extNode {
+
+					rlp_flag := uint(raw_rlp[0])
+					if rlp_flag < 192 || rlp_flag >= 248 {
+						panic("should not happen!")
+					}
+
+					element, _, _ := rlp.SplitList(raw_rlp)
+					// fmt.Println("** ", element)
+
+					// FIXME only one nibble case
+					nibble := element[0]
+					// fmt.Println(" Ext nibble:", nibble)
+					nibbles = append(nibbles, []byte{nibble})
 				}
-				proof = append(proof, rlp)
 			}
 
 		}
@@ -784,5 +844,9 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, er
 		slices.Reverse(proof)
 	}
 
-	return proof, nil
+	// given a default value
+	if len(nibbles) == 0 {
+		nibbles = append(nibbles, []byte{0})
+	}
+	return proof, nibbles, nil
 }

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

@@ -1,23 +1,41 @@
 package witness
 
 import (
+	"fmt"
 	"log"
 
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
+func PrepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []byte, extValues [][]byte, key, driftedInd byte,
+	isBranchSPlaceholder, isBranchCPlaceholder, isExtension bool) Node {
+	return prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes, extValues, key, driftedInd,
+		isBranchSPlaceholder, isBranchCPlaceholder, isExtension)
+}
+
 // isBranch takes GetProof element and returns whether the element is a branch.
 func isBranch(proofEl []byte) bool {
 	elems, _, err := rlp.SplitList(proofEl)
 	check(err)
 	c, err1 := rlp.CountValues(elems)
 	check(err1)
-	if c != 2 && c != 17 {
-		log.Fatal("Proof element is neither leaf or branch")
+	// 9: for tx (Nonce, Gas, GasPrice, Value, To, Data, r, s, v)
+	if c != 2 && c != 17 && c != 9 {
+		log.Fatal("Proof element is neither leaf or branch ", c, proofEl, elems)
 	}
 	return c == 17
 }
 
+func isTxLeaf(proofEl []byte) bool {
+	elems, _, err := rlp.SplitList(proofEl)
+	check(err)
+	c, err1 := rlp.CountValues(elems)
+	check(err1)
+	fmt.Println("ISLEAF:", c)
+	// 9: for tx (Nonce, Gas, GasPrice, Value, To, Data, r, s, v)
+	return c == 9
+}
+
 // prepareBranchWitness takes the rows that are to be filled with branch data and it takes
 // a branch as returned by GetProof. There are 19 rows for a branch and prepareBranchWitness
 // fills the rows from index 1 to index 16 (index 0 is init, index 17 and 18 are for extension
@@ -29,6 +47,7 @@ func prepareBranchWitness(rows [][]byte, branch []byte, branchStart int, branchR
 	rowInd := 1
 	colInd := branchNodeRLPLen - 1
 
+	// TODO: if input branch is a leaf node, it'll work abnormally
 	i := 0
 	insideInd := -1
 	for {
@@ -154,12 +173,23 @@ func prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []b
 func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
 	var nibbles []byte
 	if leafKeyRow[0] != 248 {
-		keyLen := int(leafKeyRow[1] - 128)
-		if (leafKeyRow[2] != 32) && (leafKeyRow[2] != 0) { // second term is for extension node
-			if leafKeyRow[2] < 32 { // extension node
-				nibbles = append(nibbles, leafKeyRow[2]-16)
-			} else { // leaf
-				nibbles = append(nibbles, leafKeyRow[2]-48)
+		keyLen := 0
+		if leafKeyRow[0] == 226 {
+			// In this case, we only have 1 nibble
+			// Prove: 226 - 192 = 34, the payload is 34 bytes and the 1st byte is RLP byte (aka 226)
+			// So, 33 bytes left, hash occupies 32 bytes in the end of the payload.
+			// 33 - 32 = 1, which is the nibble.
+			keyLen = 1
+			nibbles = append(nibbles, leafKeyRow[1]-16)
+			numberOfNibbles = 0
+		} else {
+			keyLen = int(leafKeyRow[1] - 128)
+			if (leafKeyRow[2] != 32) && (leafKeyRow[2] != 0) { // second term is for extension node
+				if leafKeyRow[2] < 32 { // extension node
+					nibbles = append(nibbles, leafKeyRow[2]-16)
+				} else { // leaf
+					nibbles = append(nibbles, leafKeyRow[2]-48)
+				}
 			}
 		}
 		for i := 0; i < keyLen-1; i++ { // -1 because the first byte doesn't have any nibbles
@@ -192,12 +222,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,
-	extNibblesS, extNibblesC [][]byte,
+func addBranchAndPlaceholder(proof1, proof2, extNibblesS, extNibblesC [][]byte,
 	leafRow0, key, neighbourNode []byte,
 	keyIndex, extensionNodeInd int,
-	additionalBranch, isAccountProof, nonExistingAccountProof,
-	isShorterProofLastLeaf bool, toBeHashed *[][]byte) (bool, bool, int, Node) {
+	additionalBranch, isAccountProof, nonExistingAccountProof, isShorterProofLastLeaf bool,
+	toBeHashed *[][]byte) (bool, bool, int, Node) {
+
 	len1 := len(proof1)
 	len2 := len(proof2)
 
@@ -210,14 +240,29 @@ func addBranchAndPlaceholder(proof1, proof2,
 		extValues = append(extValues, make([]byte, valueLen))
 	}
 
+	// FIXME have a 'isExtension' function
+	// 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]) && (!isTxLeaf(proof1[0]) && !isBranch(proof1[0]))
+
 	isExtension := (len1 == len2+2) || (len2 == len1+2)
-	if isExtension {
+	if isExtension || need_placeholder_ext {
 		var numNibbles byte
-		if len1 > len2 {
-			numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesS, extensionNodeInd, proof1[len1-3], proof1[len1-3])
+		var proof []byte
+		var extNibbles [][]byte
+		if need_placeholder_ext {
+			extNibbles = extNibblesS
+			proof = proof1[0]
 		} else {
-			numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesC, extensionNodeInd, proof2[len2-3], proof2[len2-3])
+			if len1 > len2 {
+				extNibbles = extNibblesS
+				proof = proof1[len1-3]
+			} else {
+				extNibbles = extNibblesC
+				proof = proof2[len2-3]
+			}
 		}
+		numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibbles, extensionNodeInd, proof, proof)
 		numberOfNibbles = int(numNibbles)
 	}
 
@@ -254,35 +299,33 @@ func addBranchAndPlaceholder(proof1, proof2,
 
 	// TODO: fix
 	var extNode []byte
-	if isExtension {
-		if len1 > len2 {
-			extNode = proof1[len1-3]
-		} else {
-			extNode = proof2[len2-3]
+	if need_placeholder_ext {
+		extNode = proof1[0]
+	} else {
+		if isExtension {
+			if len1 > len2 {
+				extNode = proof1[len1-3]
+			} else {
+				extNode = proof2[len2-3]
+			}
 		}
 	}
 
+	// FIXME should move to above and need to avoid above [len-3] operation
+	isExtension = need_placeholder_ext
+
 	// Note that isModifiedExtNode happens also when we have a branch instead of shortExtNode
-	isModifiedExtNode := !isBranch(longExtNode) && !isShorterProofLastLeaf
+	isModifiedExtNode := (!isBranch(longExtNode) && !isShorterProofLastLeaf) || need_placeholder_ext
 
+	// 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
+	// into the new branch.
+	driftedInd := getDriftedPosition(leafRow0, numberOfNibbles)
 	if len1 > len2 {
-		// 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
-		// into the new branch.
-		driftedInd := getDriftedPosition(leafRow0, numberOfNibbles)
-
 		node = prepareBranchNode(proof1[len1-2], proof1[len1-2], extNode, extNode, extListRlpBytes, extValues,
 			key[keyIndex+numberOfNibbles], driftedInd, false, true, isExtension)
 
-		// 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
-		// into the new branch.
 	} else {
-		// 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
-		// into the new branch.
-		driftedInd := getDriftedPosition(leafRow0, numberOfNibbles)
-
 		node = prepareBranchNode(proof2[len2-2], proof2[len2-2], extNode, extNode, extListRlpBytes, extValues,
 			key[keyIndex+numberOfNibbles], driftedInd, true, false, isExtension)
 	}

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

@@ -120,13 +120,10 @@ func transactionsStackTrieInsertionTemplate(t *testing.T, n int) {
 	}
 }
 
-// func TestJson(t *testing.T, n int) {
-// 	txs := makeTransactions(n)
-// 	db := rawdb.NewMemoryDatabase()
-// 	stackTrie := trie.NewStackTrie(db)
-
-// 	proofs, _ := stackTrie.UpdateAndGetProofs(db, types.Transactions(txs))
-// }
+func TestTransactionInsertion(t *testing.T) {
+	txs := makeTransactions(4)
+	prepareStackTrieWitness("TransactionInsertion", types.Transactions(txs))
+}
 
 func TestStackTrieInsertion_1Tx(t *testing.T) {
 	// Only one leaf
@@ -166,6 +163,12 @@ func TestStackTrieInsertion_33Txs(t *testing.T) {
 	transactionsStackTrieInsertionTemplate(t, 33)
 }
 
+func TestStackTrieInsertion_129Txs(t *testing.T) {
+	// The first tx (index 0) is inserted into position 8 of the top branch
+	// Th 129th tx is the neighbor of the first tx
+	transactionsStackTrieInsertionTemplate(t, 129)
+}
+
 func TestStackTrieInsertion_ManyTxs(t *testing.T) {
 	// Just randomly picking a large number.
 	// The cap of block gas limit is 30M, the minimum gas cost of a tx is 21k
@@ -190,7 +193,7 @@ func batchedTransactionsStackTrieProofTemplate(n int) {
 	var indexBuf []byte
 	indexBuf = rlp.AppendUint64(indexBuf[:0], uint64(1))
 
-	proofS, err := stackTrie.GetProof(db, indexBuf)
+	proofS, _, err := stackTrie.GetProof(db, indexBuf)
 	if err != nil {
 		fmt.Println(err)
 		return