cleanup and add a test

privacy-scaling-explorations · Apr 8, 2024 · 66b7f0f · 66b7f0f
1 parent 59ac6a4
commit 66b7f0f
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 49 deletions.
diff --git a/geth-utils/gethutil/mpt/trie/stacktrie.go b/geth-utils/gethutil/mpt/trie/stacktrie.go
@@ -850,6 +850,7 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, []
 				} else {
 					proof = append(proof, rlp)
 				}
+				nibbles = append(nibbles, nil)
 			} else if node.nodeType == branchNode || node.nodeType == extNode {
 				node.hash(false)
 
@@ -876,6 +877,8 @@ func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, []
 					}
 					fmt.Println(" Ext nibble:", numNibbles, nibble, raw_rlp)
 					nibbles = append(nibbles, nibble)
+				} else {
+					nibbles = append(nibbles, nil)
 				}
 			}
 

diff --git a/geth-utils/gethutil/mpt/witness/branch.go b/geth-utils/gethutil/mpt/witness/branch.go
@@ -232,11 +232,9 @@ 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,
-	leafRow0, key, neighbourNode []byte,
-	keyIndex, extensionNodeInd int,
-	additionalBranch, isAccountProof, nonExistingAccountProof, isShorterProofLastLeaf bool,
-	toBeHashed *[][]byte) (bool, bool, int, Node) {
+func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC []byte,
+	leafRow0, key []byte, keyIndex int, isShorterProofLastLeaf bool,
+) (bool, bool, int, Node) {
 
 	len1 := len(proof1)
 	len2 := len(proof2)
@@ -262,7 +260,7 @@ func addBranchAndPlaceholder(proof1, proof2, extNibblesS, extNibblesC [][]byte,
 	if isExtension || need_placeholder_ext {
 		var numNibbles byte
 		var proof []byte
-		var extNibbles [][]byte
+		var extNibbles []byte
 		if need_placeholder_ext {
 			extNibbles = extNibblesS
 			proof = proof1[0]
@@ -275,7 +273,7 @@ func addBranchAndPlaceholder(proof1, proof2, extNibblesS, extNibblesC [][]byte,
 				proof = proof2[len2-3]
 			}
 		}
-		numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibbles, extensionNodeInd, proof, proof)
+		numNibbles, extListRlpBytes, extValues = prepareExtensions(extNibbles, proof, proof)
 		numberOfNibbles = int(numNibbles)
 	}
 

diff --git a/geth-utils/gethutil/mpt/witness/extension_node.go b/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)
@@ -15,11 +15,11 @@ func prepareExtensions(extNibbles [][]byte, extensionNodeInd int, proofEl1, proo
 	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 {
@@ -30,9 +30,8 @@ func prepareExtensions(extNibbles [][]byte, extensionNodeInd int, proofEl1, proo
 		}
 	}
 	ind := 0
-	for j := startNibblePos; j < len(extNibbles[extensionNodeInd]); j += 2 {
-		v3[2+ind] = // TODO: check 2 + ind
-			extNibbles[extensionNodeInd][j]
+	for j := startNibblePos; j < len(extNibbles); j += 2 {
+		v3[2+ind] = extNibbles[j]
 		ind++
 	}
 	values = append(values, v1)

diff --git a/geth-utils/gethutil/mpt/witness/gen_witness_transactions_test.go b/geth-utils/gethutil/mpt/witness/gen_witness_transactions_test.go
@@ -120,8 +120,13 @@ func transactionsStackTrieInsertionTemplate(t *testing.T, n int) {
 	}
 }
 
+func TestTransaction144(t *testing.T) {
+	txs := makeTransactions(146)
+	prepareStackTrieWitness("TransactionInsertion", types.Transactions(txs))
+}
+
 func TestTransactionInsertion(t *testing.T) {
-	txs := makeTransactions(256)
+	txs := makeTransactions(145)
 	prepareStackTrieWitness("TransactionInsertion", types.Transactions(txs))
 }
 

diff --git a/geth-utils/gethutil/mpt/witness/modified_extension_node.go b/geth-utils/gethutil/mpt/witness/modified_extension_node.go
@@ -13,9 +13,7 @@ import (
 // The rows added are somewhat exceptional as otherwise they do not appear.
 func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr common.Address, proof1, proof2, proofTx,
 	extNibblesS, extNibblesC [][]byte,
-	key, neighbourNode []byte,
-	keyIndex, extensionNodeInd, numberOfNibbles int,
-	isAccountProof bool, toBeHashed *[][]byte) Node {
+	key []byte, keyIndex, numberOfNibbles int, isAccountProof bool) Node {
 	len1 := len(proof1)
 	len2 := len(proof2)
 
@@ -38,7 +36,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:
@@ -117,7 +115,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

diff --git a/geth-utils/gethutil/mpt/witness/prepare_witness.go b/geth-utils/gethutil/mpt/witness/prepare_witness.go
@@ -339,21 +339,21 @@ func prepareWitnessSpecial(testName string, trieModifications []TrieModification
 
 // For stack trie, we have the following combinations ([proofS] -> [proofC])
 //
-//	-[] [(empty)] -> [LEAF]
+//	-[x] [(empty)] -> [LEAF]
 //	-[] [LEAF] -> [EXT - BRANCH - LEAF]
-//	-[] [EXT - BRANCH] -> [EXT - BRANCH - LEAF]
+//	-[x] [EXT - BRANCH] -> [EXT - BRANCH - LEAF]
 //	-[] [EXT - BRANCH] -> [BRANCH - LEAF] (modified extension)
 //	-[] [EXT] -> [BRANCH - BRANCH - BRANCH - LEAF] --> 144
 //	-[] [BRANCH - LEAF] -> [BRANCH - BRANCH - LEAF] (modified extension)
-//	-[] [BRANCH - BRANCH - (...BRANCH)] -> [BRANCH - BRANCH - (...BRANCH) - LEAF]
-//	-[] [BRANCH - BRANCH - LEAF ] -> [BRANCH - BRANCH - LEAF - BRANCH - LEAF]
-//	-[] [LEAF] -> [LEAF]
+//	-[x] [BRANCH - BRANCH - (...BRANCH)] -> [BRANCH - BRANCH - (...BRANCH) - LEAF]
+//	-[] [BRANCH - BRANCH - LEAF] -> [BRANCH - BRANCH - LEAF - BRANCH - LEAF]
+//	-[x] [LEAF] -> [LEAF]  --> 510
 //	-[] [EXT] -> [EXT]
 //	-[] [EXT - EXT] -> [EXT - EXT]
 //	-[] [EXT - LEAF] -> [LEAF]
 //	-[] [BRANCH - LEAF] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF]
-//	-[] [BRANCH - BRANCH - EXT] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF]
-//	-[] [BRANCH - BRANCH - EXT - BRANCH - (LEAF)] -> [BRANCH - BRANCH - EXT - BRANCH - EXT - BRANCH - LEAF]
+//	-[] [BRANCH - BRANCH - EXT] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF] --> 512
+//	-[x] [BRANCH - BRANCH - EXT - BRANCH - (LEAF)] -> [BRANCH - BRANCH - EXT - BRANCH - EXT - BRANCH - LEAF]
 //	-[] [BRANCH - BRANCH - EXT - BRANCH - (...BRANCH)] -> [BRANCH - BRANCH - EXT - BRANCH - (...BRANCH) - LEAF]
 //	-[] [LEAF] -> [BRANCH - BRANCH - EXT - BRANCH - BRANCH - LEAF]
 func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []Node {
@@ -364,7 +364,6 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
 	k = append(k, kk...)
 	fmt.Println("txIdx", txIdx)
 
-	toBeHashed := make([][]byte, 0)
 	proofS := proof.GetProofS()
 	proofC := proof.GetProofC()
 	extNibblesS := proof.GetNibblesS()
@@ -394,18 +393,18 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
 		}
 
 		// The length of proofS and proofC is equal and
-		// the last element of proofC is an extension
-		if len1 == len2 && isTxExt(proofC[len2-1]) {
+		// the last element of proofC is an extension or a leaf
+		if len1 == len2 && (isTxExt(proofC[len2-1]) || isTxLeaf(proofC[len2-1])) {
 			additionalBranch = false
 		}
 
-		if isTxLeaf(proofS[len1-1]) && isTxLeaf(proofC[len2-1]) {
+		// isStartWithExt && len1 > 1 is for "[EXT - BRANCH] -> [EXT - BRANCH - LEAF]""
+		if (isStartWithExt && len1 > 1) || (len1 == len2 && isTxLeaf(proofS[len1-1]) && isTxLeaf(proofC[len2-1])) {
 			upTo = minLen
 		}
 	}
 
 	isExtension := false
-	extensionNodeInd := 0
 
 	var extListRlpBytes []byte
 	var extValues [][]byte
@@ -417,17 +416,15 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
 	for i := 0; i < upTo; i++ {
 		if !isBranch(proofS[i]) {
 			fmt.Println("extNibbleS/C", extNibblesS, extNibblesC)
-			areThereNibbles := len(extNibblesS) != 0
+			areThereNibbles := len(extNibblesS[i]) != 0
 			if areThereNibbles { // extension node
 				var numberOfNibbles byte
 				isExtension = true
 
 				// FIXME: handle the case of proofS(ext), proofC(branch), and it needs to add a branch placeholder at proofS?
-				numberOfNibbles, extListRlpBytes, extValues =
-					prepareExtensions(extNibblesS, extensionNodeInd, proofS[i], proofC[i])
+				numberOfNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesS[i], proofS[i], proofC[i])
 
 				keyIndex += int(numberOfNibbles)
-				// extensionNodeInd++
 				fmt.Println("Increase keyIdx", keyIndex)
 				continue
 			}
@@ -460,10 +457,8 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
 		leafRow0 := proofS[0] // To compute the drifted position.
 		isModifiedExtNode, _, _, bNode :=
 			addBranchAndPlaceholder(
-				proofS, proofC, extNibblesS, extNibblesC,
-				leafRow0, k, nil,
-				keyIndex, extensionNodeInd, additionalBranch,
-				false, false, is_end_with_leaf, &toBeHashed)
+				proofS, proofC, extNibblesS[len1-1], extNibblesC[len2-1],
+				leafRow0, k, keyIndex, is_end_with_leaf)
 
 		if !isStartWithExt {
 			nodes = append(nodes, bNode)
@@ -491,7 +486,7 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
 			// modification).
 			// leafNode = equipLeafWithModExtensionNode(nil, leafNode, common.Address{0}, proofS, proofC, proofC,
 			// 	extNibblesS, extNibblesC, k, nil,
-			// 	keyIndex, extensionNodeInd, numberOfNibbles, false, &toBeHashed)
+			// 	keyIndex, numberOfNibbles, false, &toBeHashed)
 		}
 
 		nodes = append(nodes, leafNode)
@@ -530,7 +525,6 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 	proof1, proof2, extNibblesS, extNibblesC [][]byte,
 	storage_key common.Hash, key []byte, neighbourNode []byte,
 	isAccountProof, nonExistingAccountProof, nonExistingStorageProof, isShorterProofLastLeaf bool) []Node {
-	toBeHashed := make([][]byte, 0)
 
 	minLen := len(proof1)
 	if len(proof2) < minLen {
@@ -587,7 +581,7 @@ 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++
@@ -628,10 +622,8 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 				leafRow0 = proof2[len2-1]
 			}
 
-			isModifiedExtNode, _, numberOfNibbles, bNode := addBranchAndPlaceholder(proof1, proof2, extNibblesS, extNibblesC,
-				leafRow0, key, neighbourNode,
-				keyIndex, extensionNodeInd, additionalBranch,
-				isAccountProof, nonExistingAccountProof, isShorterProofLastLeaf, &toBeHashed)
+			isModifiedExtNode, _, numberOfNibbles, bNode :=
+				addBranchAndPlaceholder(proof1, proof2, extNibblesS[len1-1], extNibblesC[len2-1], leafRow0, key, keyIndex, isShorterProofLastLeaf)
 
 			nodes = append(nodes, bNode)
 
@@ -670,8 +662,8 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
 			// of the existing extension node), additional rows are added (extension node before and after
 			// modification).
 			if isModifiedExtNode {
-				leafNode = equipLeafWithModExtensionNode(statedb, leafNode, addr, proof1, proof2, nil, extNibblesS, extNibblesC, key, neighbourNode,
-					keyIndex, extensionNodeInd, numberOfNibbles, isAccountProof, &toBeHashed)
+				leafNode = equipLeafWithModExtensionNode(statedb, leafNode, addr, proof1, proof2, nil, extNibblesS, extNibblesC,
+					key, keyIndex, numberOfNibbles, isAccountProof)
 			}
 			nodes = append(nodes, leafNode)
 		} else {