From 76ab4ecbb6e971dbce3f43a963cf5e49f72644f5 Mon Sep 17 00:00:00 2001 From: KimiWu Date: Mon, 8 Apr 2024 18:18:10 +0800 Subject: [PATCH] cleanup and add a test --- geth-utils/gethutil/mpt/trie/stacktrie.go | 3 + geth-utils/gethutil/mpt/witness/branch.go | 56 ++++++++++--------- .../gethutil/mpt/witness/extension_node.go | 13 ++--- .../witness/gen_witness_transactions_test.go | 7 ++- .../mpt/witness/modified_extension_node.go | 8 +-- .../gethutil/mpt/witness/prepare_witness.go | 50 +++++++---------- 6 files changed, 68 insertions(+), 69 deletions(-) diff --git a/geth-utils/gethutil/mpt/trie/stacktrie.go b/geth-utils/gethutil/mpt/trie/stacktrie.go index e78172004c..79c2888b28 100644 --- 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 index c9caff2975..869d99bada 100644 --- a/geth-utils/gethutil/mpt/witness/branch.go +++ b/geth-utils/gethutil/mpt/witness/branch.go @@ -178,21 +178,12 @@ func prepareBranchNode( return node } -// getDriftedPosition returns the position in branch to which the leaf drifted because another -// leaf has been added to the same slot. This information is stored into a branch init row. -func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte { +// getNibbles returns the nibbles of the leaf or extension node. +func getNibbles(leafKeyRow []byte) []byte { var nibbles []byte if leafKeyRow[0] != 248 { - 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 { + var keyLen int + if leafKeyRow[1] > 128 { keyLen = int(leafKeyRow[1] - 128) if (leafKeyRow[2] != 32) && (leafKeyRow[2] != 0) { // second term is for extension node if leafKeyRow[2] < 32 { // extension node @@ -201,13 +192,16 @@ func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte { nibbles = append(nibbles, leafKeyRow[2]-48) } } - } - for i := 0; i < keyLen-1; i++ { // -1 because the first byte doesn't have any nibbles - b := leafKeyRow[3+i] - n1 := b / 16 - n2 := b - n1*16 - nibbles = append(nibbles, n1) - nibbles = append(nibbles, n2) + for i := 0; i < keyLen-1; i++ { // -1 because the first byte doesn't have any nibbles + b := leafKeyRow[3+i] + n1 := b / 16 + n2 := b - n1*16 + nibbles = append(nibbles, n1) + nibbles = append(nibbles, n2) + } + } else { + keyLen = 1 + nibbles = append(nibbles, leafKeyRow[1]-16) } } else { keyLen := int(leafKeyRow[2] - 128) @@ -227,16 +221,24 @@ func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte { } } + return nibbles +} + +// getDriftedPosition returns the position in branch to which the leaf drifted because another +// leaf has been added to the same slot. This information is stored into a branch init row. +func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte { + nibbles := getNibbles(leafKeyRow) + if len(nibbles) == 1 { + return nibbles[0] + } return nibbles[numberOfNibbles] } // 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 +264,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 +277,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 index e52a06f1e6..fee54ac26a 100644 --- 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 index a2f1ab0ad3..c7d77cab72 100644 --- 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 index 7852d0ba75..828f790401 100644 --- 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 index b1d72bf33e..03bb292d38 100644 --- 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 {