adding more testings and refactor getNodeFromBranchRLP

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

@@ -551,50 +551,49 @@ func (st *StackTrie) Commit() (common.Hash, error) {
 	return common.BytesToHash(st.val), nil
 }
 
-func (st *StackTrie) getNodeFromBranchRLP(branch []byte, ind byte) []byte {
-	start := 2 // when branch[0] == 248
-	if branch[0] == 249 {
-		start = 3
-	}
-
-	i := 0
-	insideInd := -1
-	cInd := byte(0)
-	for {
-		if start+i == len(branch)-1 { // -1 because of the last 128 (branch value)
-			return []byte{0}
-		}
-		b := branch[start+i]
-		if insideInd == -1 && b == 128 {
-			if cInd == ind {
+const RLP_SHORT_STR_FLAG = 128
+const RLP_LONG_LIST_FLAG = 248
+const LEN_OF_HASH = 32
+
+// Note:
+// In RLP encoding, if the value is between [0x80, 0xb7] ([128, 183]),
+// it means following data is a short string (0 - 55bytes).
+// Which implies if the value is 128, it's an empty string.
+func (st *StackTrie) getNodeFromBranchRLP(branch []byte, idx int) []byte {
+	// (branch[0] - RLP_LONG_LIST_FLAG) means the length in bytes of the length of the payload
+	// and the payload is right after the length.
+	// That's why we add `2` here
+	// e.g. [248 81 128 160 ...]
+	// `81` is the length of the payload and payload starts from `128`
+	start := int(branch[0]) - RLP_LONG_LIST_FLAG + 2
+
+	// If 1st node is not 128(empty node) or 160, it should be a leaf
+	b := int(branch[start])
+	if b != RLP_SHORT_STR_FLAG || b != (RLP_SHORT_STR_FLAG+LEN_OF_HASH) {
+		return []byte{0}
+	}
+
+	current_idx := 0
+	for i := start; i < len(branch); i++ {
+		b = int(branch[i])
+		switch b {
+		case RLP_SHORT_STR_FLAG: // 128
+			// if the current index is we're looking for, return an empty node directly
+			if current_idx == idx {
 				return []byte{128}
-			} else {
-				cInd += 1
 			}
-		} else if insideInd == -1 && b != 128 {
-			if b == 160 {
-				if cInd == ind {
-					return branch[start+i+1 : start+i+1+32]
-				}
-				insideInd = 32
-			} else {
-				// non-hashed node
-				if cInd == ind {
-					return branch[start+i+1 : start+i+1+int(b)-192]
-				}
-				insideInd = int(b) - 192
-			}
-			cInd += 1
-		} else {
-			if insideInd == 1 {
-				insideInd = -1
-			} else {
-				insideInd--
+			current_idx++
+		case RLP_SHORT_STR_FLAG + LEN_OF_HASH: // 160
+			if current_idx == idx {
+				return branch[i+1 : i+1+LEN_OF_HASH]
 			}
+			// jump to next encoded element
+			i += LEN_OF_HASH
+			current_idx++
 		}
 
-		i++
 	}
+	return []byte{0}
 }
 
 type StackProof struct {
@@ -700,7 +699,7 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, er
 			}
 
 			proof = append(proof, c_rlp)
-			branchChild := st.getNodeFromBranchRLP(c_rlp, k[i])
+			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.

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

@@ -8,16 +8,40 @@ import (
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
+func makeTransactions(n int) []*types.Transaction {
+	txs := make([]*types.Transaction, n)
+	key, _ := crypto.GenerateKey()
+	signer := types.LatestSigner(params.TestChainConfig)
+
+	for i := range txs {
+		amount := big.NewInt(int64(i)*10 ^ 9)
+		gas_price := big.NewInt(300_000)
+		data := make([]byte, 100)
+		// randomly assigned for debugging
+		data[3] = 3
+		data[4] = 3
+		data[5] = 3
+		data[6] = 3
+		data[7] = 3
+		tx := types.NewTransaction(uint64(i), common.Address{}, amount, 10*10^6, gas_price, data)
+		signedTx, err := types.SignTx(tx, signer, key)
+		if err != nil {
+			panic(err)
+		}
+		txs[i] = signedTx
+	}
+	return txs
+}
+
 /*
-TestNonHashedTransactionsStackTrieGetProof inserts 70 transactions into a stacktrie.
-For each of the 70 modifications of the trie it asks for a proof - GetProof is called before
+transactionsStackTrieInsertionTemplate inserts n transactions into a stacktrie.
+For each of the n modifications of the trie it asks for a proof - GetProof is called before
 and after the modification. The transactions in the trie are not hashed and thus GetProof
 does not require to query a database to get the preimages.
 
@@ -62,23 +86,9 @@ The first proof element is a branch with children at position 0 (branch B) and 1
 The second element is the 16-th transaction. For example, the third byte (16) represents
 the transaction index.
 */
-func TestNonHashedTransactionsStackTrieGetProof(t *testing.T) {
-	txs := make([]*types.Transaction, 70)
-	key, _ := crypto.GenerateKey()
-	signer := types.LatestSigner(params.TestChainConfig)
-
-	for i := range txs {
-		amount := math.BigPow(2, int64(i))
-		price := big.NewInt(300000)
-		data := make([]byte, 100)
-		tx := types.NewTransaction(uint64(i), common.Address{}, amount, 123457, price, data)
-		signedTx, err := types.SignTx(tx, signer, key)
-		if err != nil {
-			panic(err)
-		}
-		txs[i] = signedTx
-	}
 
+func transactionsStackTrieInsertionTemplate(n int) {
+	txs := makeTransactions(n)
 	db := rawdb.NewMemoryDatabase()
 	stackTrie := trie.NewStackTrie(db)
 
@@ -87,34 +97,59 @@ func TestNonHashedTransactionsStackTrieGetProof(t *testing.T) {
 	fmt.Println("===")
 }
 
+func TestStackTrieInsertion_1Tx(t *testing.T) {
+	// Only one leaf
+	transactionsStackTrieInsertionTemplate(1)
+}
+
+func TestStackTrieInsertion_2Txs(t *testing.T) {
+	// One ext. node and one leaf
+	transactionsStackTrieInsertionTemplate(2)
+}
+
+func TestStackTrieInsertion_3Txs(t *testing.T) {
+	// One ext. node, one branch and one leaf
+	transactionsStackTrieInsertionTemplate(3)
+}
+
+func TestStackTrieInsertion_4Txs(t *testing.T) {
+	// One ext. node, one branch and two leaves
+	transactionsStackTrieInsertionTemplate(4)
+}
+
+func TestStackTrieInsertion_16Txs(t *testing.T) {
+	// One ext. node and one branch with full leaves (16 leaves)
+	transactionsStackTrieInsertionTemplate(16)
+}
+
+func TestStackTrieInsertion_17Txs(t *testing.T) {
+	// One ext. node, 3 branches and 17 leaves.
+	// The original ext. node turns into a branch (B1) which has children at position 0 and 1.
+	// At position 0 of B1, it has a branch with full leaves
+	// At position 1 of B1, it has a newly leaf
+	transactionsStackTrieInsertionTemplate(17)
+}
+
+func TestStackTrieInsertion_33Txs(t *testing.T) {
+	// Follow above test and have one more branch generated
+	transactionsStackTrieInsertionTemplate(33)
+}
+
+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
+	// 30M / 21k ~= 1429
+	transactionsStackTrieInsertionTemplate(2000)
+}
+
 /*
-TestHashedTransactionsStackTrieGetProof inserts 2 transactions into a stacktrie,
+batchedTransactionsStackTrieProofTemplate inserts n transactions into a stacktrie,
 the trie is then hashed (DeriveSha call).
-The proof is asked for one of the two transactions. The transactions in the trie are hashed and thus
+The proof is asked for one of the n transactions. The transactions in the trie are hashed and thus
 GetProof requires to query a database to get the preimages.
 */
-func TestHashedTransactionsStackTrieGetProof(t *testing.T) {
-	txs := make([]*types.Transaction, 2)
-	key, _ := crypto.GenerateKey()
-	signer := types.LatestSigner(params.TestChainConfig)
-
-	for i := range txs {
-		amount := math.BigPow(2, int64(i))
-		price := big.NewInt(300000)
-		data := make([]byte, 100)
-		data[3] = 3
-		data[4] = 3
-		data[5] = 3
-		data[6] = 3
-		data[7] = 3
-		tx := types.NewTransaction(uint64(i), common.Address{}, amount, 123457, price, data)
-		signedTx, err := types.SignTx(tx, signer, key)
-		if err != nil {
-			panic(err)
-		}
-		txs[i] = signedTx
-	}
-
+func batchedTransactionsStackTrieProofTemplate(n int) {
+	txs := makeTransactions(n)
 	db := rawdb.NewMemoryDatabase()
 	stackTrie := trie.NewStackTrie(db)
 
@@ -131,6 +166,34 @@ func TestHashedTransactionsStackTrieGetProof(t *testing.T) {
 	}
 
 	fmt.Println(proofS)
-
 	fmt.Println("===")
 }
+
+func TestBatchedTxsProof_1Tx(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(1)
+}
+
+func TestBatchedTxsProof_2Txs(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(2)
+}
+
+func TestBatchedTxsProof_3Txs(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(3)
+}
+func TestBatchedTxsProof_4Txs(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(4)
+}
+
+func TestBatchedTxsProof_16Txs(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(16)
+}
+
+func TestBatchedTxsProof_17Txs(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(17)
+}
+func TestBatchedTxsProof_33Txs(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(33)
+}
+func TestBatchedTxsProof_ManyTxs(t *testing.T) {
+	batchedTransactionsStackTrieProofTemplate(2000)
+}