Skip to content
This repository was archived by the owner on Jul 5, 2024. It is now read-only.

Commit ea2f647

Browse files
KimiWu123KimiWu123
committed
refine for mod-ext, but still not working
1 parent 85cfe2d commit ea2f647

File tree

4 files changed

+93
-84
lines changed

4 files changed

+93
-84
lines changed

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -701,7 +701,8 @@ func printProof(ps [][]byte, t, idx []byte) {
701701
}
702702

703703
func (st *StackTrie) UpdateAndGetProof(db ethdb.KeyValueReader, indexBuf, value []byte) (StackProof, error) {
704-
fmt.Println(" ====", indexBuf)
704+
fmt.Println(" ====", indexBuf, "-->", KeybytesToHex(indexBuf))
705+
705706
proofS, nibblesS, typesS, err := st.GetProof(db, indexBuf)
706707
if err != nil {
707708
return StackProof{}, err
@@ -772,7 +773,7 @@ func (st *StackTrie) UpdateAndGetProofs(db ethdb.KeyValueReader, list types.Deri
772773

773774
func (st *StackTrie) GetProof(db ethdb.KeyValueReader, key []byte) ([][]byte, [][]byte, []uint8, error) {
774775
k := KeybytesToHex(key)
775-
fmt.Println(" k", k)
776+
// fmt.Println(" k", k)
776777
if st.nodeType == emptyNode {
777778
return [][]byte{}, nil, []uint8{emptyNode}, nil
778779
}

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

+6-6
Original file line numberDiff line numberDiff line change
@@ -239,12 +239,12 @@ func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
239239
// addBranchAndPlaceholder adds to the rows a branch and its placeholder counterpart
240240
// (used when one of the proofs have one branch more than the other).
241241
func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC []byte,
242-
proofTx, leafRow0, key []byte, keyIndex int, isShorterProofLastLeaf bool,
242+
extProofTx, leafRow0, key []byte, keyIndex int, isShorterProofLastLeaf bool,
243243
) (bool, bool, int, Node) {
244244

245245
len1 := len(proof1)
246246
len2 := len(proof2)
247-
isTxProof := len(proofTx) != 0
247+
isTxProof := len(extProofTx) != 0
248248

249249
var node Node
250250

@@ -255,14 +255,14 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC [
255255
extValues = append(extValues, make([]byte, valueLen))
256256
}
257257

258-
isExtension := (len1 == len2+2) || (len2 == len1+2) || (isTxProof && !isBranch(proofTx))
258+
isExtension := (len1 == len2+2) || (len2 == len1+2) || (isTxProof && !isBranch(extProofTx))
259259
if isExtension {
260260
var numNibbles byte
261261
var proof []byte
262262
var extNibbles []byte
263263
if isTxProof {
264264
extNibbles = extNibblesS
265-
proof = proofTx
265+
proof = extProofTx
266266
} else {
267267
if len1 > len2 {
268268
extNibbles = extNibblesS
@@ -302,7 +302,7 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC [
302302
*/
303303
var longExtNode []byte
304304
if isTxProof {
305-
longExtNode = proofTx
305+
longExtNode = extProofTx
306306
} else {
307307
if len1 > len2 {
308308
longExtNode = proof2[len2-1]
@@ -313,7 +313,7 @@ func addBranchAndPlaceholder(proof1, proof2 [][]byte, extNibblesS, extNibblesC [
313313

314314
var extNode []byte
315315
if isTxProof {
316-
extNode = proofTx
316+
extNode = extProofTx
317317
} else {
318318
if isExtension {
319319
if len1 > len2 {

‎geth-utils/gethutil/mpt/witness/modified_extension_node.go

+12-12
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,19 @@ import (
1111
// These rows are added only when an existing extension node gets shortened or elongated (in terms
1212
// of the extension node nibbles) because of another extension node being added or deleted.
1313
// The rows added are somewhat exceptional as otherwise they do not appear.
14-
func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr common.Address, proof1, proof2, proofTx,
15-
extNibblesS, extNibblesC [][]byte,
16-
key []byte, keyIndex, numberOfNibbles int, isAccountProof bool) Node {
14+
func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr common.Address,
15+
proof1, proof2, extNibblesS, extNibblesC [][]byte,
16+
proofTx, key []byte,
17+
keyIndex, numberOfNibbles int, isAccountProof bool) Node {
18+
1719
len1 := len(proof1)
1820
len2 := len(proof2)
21+
isTxProof := len(proofTx) != 0
1922

2023
var longExtNode []byte
2124
// FIXME this is a workaround to get ext node
22-
if len(proofTx) != 0 {
23-
longExtNode = proof1[0]
25+
if isTxProof {
26+
longExtNode = proofTx
2427
} else {
2528
if len1 > len2 {
2629
longExtNode = proof2[len2-1]
@@ -38,7 +41,7 @@ func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr c
3841

3942
_, extListRlpBytesS, extValuesS := prepareExtensions(extNibbles[len(extNibbles)-1], longExtNode, longExtNode)
4043

41-
// Get nibbles of the extension node that gets shortened because of the newly insertd
44+
// Get nibbles of the extension node that gets shortened because of the newly inserted
4245
// extension node:
4346
longNibbles := getExtensionNodeNibbles(longExtNode)
4447

@@ -57,10 +60,7 @@ func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr c
5760
ky := common.BytesToHash(k)
5861
var proof [][]byte
5962
var err error
60-
// FIXME refactor for 3 different kinds of proofs
61-
if len(proofTx) != 0 {
62-
proof = proofTx
63-
} else {
63+
if !isTxProof {
6464
if isAccountProof {
6565
proof, _, _, _, _, err = statedb.GetProof(addr)
6666
} else {
@@ -72,7 +72,7 @@ func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr c
7272
// There is no short extension node when `len(longNibbles) - numberOfNibbles = 1`, in this case there
7373
// is simply a branch instead.
7474
// stack trie is always a short ext node.
75-
shortExtNodeIsBranch := (len(longNibbles)-numberOfNibbles == 1) || (len(proofTx) != 0)
75+
shortExtNodeIsBranch := (len(longNibbles)-numberOfNibbles == 1) || isTxProof
7676

7777
var shortExtNode []byte
7878
var extListRlpBytesC []byte
@@ -144,7 +144,7 @@ func equipLeafWithModExtensionNode(statedb *state.StateDB, leafNode Node, addr c
144144
keccakData = append(keccakData, longExtNode)
145145
keccakData = append(keccakData, shortExtNode)
146146

147-
if len(proofTx) != 0 {
147+
if isTxProof {
148148
leafNode.Transaction.ModListRlpBytes = listRlpBytes
149149
} else {
150150
if leafNode.Account == nil {

‎geth-utils/gethutil/mpt/witness/prepare_witness.go

+72-64
Original file line numberDiff line numberDiff line change
@@ -292,7 +292,7 @@ func obtainTwoProofsAndConvertToWitness(trieModifications []TrieModification, st
292292
// prepareStackTrieWitness obtains the GetProof proof before and after the modification for each
293293
// of the modification. It then converts the two proofs into an MPT circuit witness for each of
294294
// the modifications and stores it into a file.
295-
func prepareStackTrieWitness(testName string, list types.DerivableList) {
295+
func prepareStackTrieWitness(testName string, list types.DerivableList, loop bool) {
296296
db := rawdb.NewMemoryDatabase()
297297
stackTrie := trie.NewStackTrie(db)
298298
proofs, _ := stackTrie.UpdateAndGetProofs(db, list)
@@ -303,14 +303,21 @@ func prepareStackTrieWitness(testName string, list types.DerivableList) {
303303
for i, proof := range proofs {
304304
idx := i + 1
305305

306-
// debug section
307-
// i := len(proofs) - 1
308-
// if len(proofs) > 128 {
309-
// i = len(proofs) - 1
310-
// }
311-
// proof := proofs[i]
312-
// idx := i
306+
// ==== debug section
307+
if !loop {
308+
i := len(proofs) - 2
309+
if len(proofs) > 128 {
310+
i = len(proofs) - 1
311+
}
312+
proof = proofs[i]
313+
idx = i
314+
315+
// for _, p := range proof.GetProofC() {
316+
// fmt.Println("C: ", p)
317+
// }
318+
}
313319
// =====
320+
314321
var subNodes []Node
315322
subNodes = append(subNodes, GetStartNode(testName, common.Hash{}, root, 0))
316323
var node []Node
@@ -326,6 +333,10 @@ func prepareStackTrieWitness(testName string, list types.DerivableList) {
326333
verifyNodeNumber(subNodes, proof)
327334

328335
nodes = append(nodes, subNodes...)
336+
337+
if !loop {
338+
break
339+
}
329340
}
330341
StoreNodes(testName, nodes)
331342

@@ -411,12 +422,19 @@ func prepareWitnessSpecial(testName string, trieModifications []TrieModification
411422
// -[o] [BRANCH] -> [BRANCH - LEAF] --> 0
412423
// -[o] [BRANCH - BRANCH - LEAF] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF] --> 129
413424
// -[o] [BRANCH - BRANCH - EXT - BRANCH] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF] --> 130
414-
// -[M] [BRANCH - BRANCH - EXT - BRANCH - HASHED] -> [BRANCH - BRANCH - BRANCH - LEAF] --> 144(hashed node)
415-
// -[M] [BRANCH - BRANCH - EXT - BRANCH - BRANCH - HASHED] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF] --> 512 (hashed node)
425+
// -[M] [BRANCH - BRANCH - EXT - BRANCH - HASHED] -> [BRANCH - BRANCH - BRANCH - LEAF] --> 144
426+
// -[M] [BRANCH - BRANCH - EXT - BRANCH - BRANCH - HASHED] -> [BRANCH - BRANCH - EXT - BRANCH - LEAF] --> 512
416427
// -[o] [BRANCH - BRANCH - (...BRANCH)] -> [BRANCH - BRANCH - (...BRANCH) - LEAF] --> 146 ~ 176
417428
// -[o] [BRANCH - BRANCH - EXT - (BRANCH..)] -> [BRANCH - BRANCH - EXT - (BRANCH..) - LEAF] --> 258~
418429
// -[o] [BRANCH - BRANCH - EXT - BRANCH - LEAF] -> [BRANCH - BRANCH - EXT - BRANCH - EXT - BRANCH - LEAF] --> 513
419430
// -[o] [BRANCH - BRANCH - EXT - BRANCH - EXT - BRANCH] -> [BRANCH - BRANCH - EXT - BRANCH - EXT - BRANCH - LEAF] --> 514~
431+
//
432+
// TODO modified extension node
433+
// Take tx144 as example, the proof is
434+
// [BRANCH_S1 - BRANCH_S2 - EXT_S - BRANCH_S3 - HASHED] -> [BRANCH_C1 - BRANCH_C2 - BRANCH_C3 - LEAF]
435+
// We need to generate a json with nodes
436+
// [{BRANCH_S1-BRANCH_C1}, {BRANCH_S2-BRANCH_C2}, {EXT_S, BRANCH_S3-placeholder}, {placeholder-BRANCH_C3}, {placeholder-LEAF}]
437+
// We didn't have the 4th node, {placeholder-BRANCH_C3} now.
420438
func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []Node {
421439
k := trie.KeybytesToHex(key)
422440
k = k[:len(k)-1]
@@ -436,7 +454,8 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
436454
len1 := len(proofS)
437455
len2 := len(proofC)
438456
var nodes []Node
439-
// Empty stack trie
457+
458+
// Special case for the 1st tx, an empty stack trie
440459
if len1 == 0 {
441460
leafNode := prepareTxLeafAndPlaceholderNode(1, proofC[0], k, false)
442461
return append(nodes, leafNode)
@@ -450,26 +469,24 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
450469
// FIXME: using enum(branch, leaf...) to replace magic numbers
451470
upTo := minLen
452471
additionalBranch := true
453-
if len1 > 0 {
454-
// If both of proofs end with a leaf or a hashed node, e.g. [BRANCH - LEAF] --> [BRANCH - BRANCH - LEAF]
455-
// upTo minus 1 means we handle branch placeholder and leaf in `additionalBranch`
456-
// if len1 != len2 && (lastProofTypeS == lastProofTypeC) && (lastProofTypeC == 3) {
457-
if len1 != len2 && (lastProofTypeS == 3 || lastProofTypeS == 4) && (lastProofTypeC == 3 || lastProofTypeC == 4) {
458-
upTo--
459-
}
460472

461-
// The length of proofS and proofC is equal and
462-
// the last element of proofS is a hashed node or a leaf
463-
if len1 == len2 && (lastProofTypeS == 3 || lastProofTypeS == 4) {
464-
additionalBranch = false
465-
}
473+
// If both of proofs end with either a leaf or a hashed node, e.g. [BRANCH - LEAF] --> [BRANCH - BRANCH - LEAF]
474+
// The 2nd BRANCH in above proofC should have a branch placeholder for it.
475+
// We handle branch placeholder in `additionalBranch` and that's why we need to minus `upTo` by 1 here.
476+
if len1 != len2 && (lastProofTypeS == 3 || lastProofTypeS == 4) && (lastProofTypeC == 3 || lastProofTypeC == 4) {
477+
upTo--
478+
}
466479

467-
// Special case for the 2nd tx.
468-
// In this case, proofS only contains a leaf node and proofC are [EXT - BRANCH - LEAF].
469-
// `additionalBranch` can handle the mismatched the order of the type.
470-
if len1 == 1 && lastProofTypeS == 3 {
471-
upTo = 0
472-
}
480+
// The length of proofS and proofC is equal and the last element of proofS is a hashed node or a leaf
481+
if len1 == len2 && (lastProofTypeS == 3 || lastProofTypeS == 4) {
482+
additionalBranch = false
483+
}
484+
485+
// Special case for the 2nd tx.
486+
// In this case, proofS only contains a leaf node and proofC is [EXT - BRANCH - LEAF].
487+
// `additionalBranch` can handle the mismatched the order of the type.
488+
if len1 == 1 && lastProofTypeS == 3 {
489+
upTo = 0
473490
}
474491

475492
var extListRlpBytes []byte
@@ -478,14 +495,16 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
478495
extValues = append(extValues, make([]byte, 34))
479496
}
480497

498+
var numberOfNibbles byte
481499
isExtension := false
482500
mismatchedIdx := -1
483501
fmt.Println("upto", upTo, additionalBranch, proofTypeS)
484502
for i := 0; i < upTo; i++ {
485503
if proofTypeS[i] != 1 {
486504
// fmt.Println("extNibbleS/C", extNibblesS, "|", extNibblesC)
487505

488-
// This is designed for ext-mod case due to the order of the types will be mismatched. See this example,
506+
// This is for the case of extension modified node due to the order of the types mismatched.
507+
// See this example,
489508
// [BRANCH - BRANCH - EXT - BRANCH - HASHED] -> [BRANCH - BRANCH - BRANCH - LEAF]
490509
// In this case, `mismatchedIdx`` is 2 and stops at `EXT` node
491510
if proofTypeS[i] != proofTypeC[i] {
@@ -495,12 +514,9 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
495514

496515
areThereNibbles := len(extNibblesS[i]) != 0 || len(extNibblesC[i]) != 0
497516
if areThereNibbles { // extension node
498-
var numberOfNibbles byte
499517
isExtension = true
500-
501518
numberOfNibbles, extListRlpBytes, extValues = prepareExtensions(extNibblesS[i], proofS[i], proofC[i])
502519
keyIndex += int(numberOfNibbles)
503-
// fmt.Println("Increase keyIdx", keyIndex)
504520
continue
505521
}
506522

@@ -514,7 +530,6 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
514530
extNode2 = proofC[i-1]
515531
}
516532

517-
// fmt.Print(i, "/", proofTypeS[i], " - ")
518533
bNode := prepareBranchNode(
519534
proofS[i], proofC[i], extNode1, extNode2, extListRlpBytes,
520535
extValues, k[keyIndex], k[keyIndex],
@@ -528,66 +543,59 @@ func GenerateWitness(txIdx uint, key, value []byte, proof *trie.StackProof) []No
528543

529544
// To address the length of proofS and proofC is not equal or the order of the type is matched.
530545
if additionalBranch {
531-
// fmt.Println("additionalBranch: extNibbleS/C", extNibblesS, "|", extNibblesC)
546+
lastProofType := lastProofTypeS
532547
leafRow0 := proofS[len1-1] // To compute the drifted position.
533548
if len1 > len2 {
534549
leafRow0 = proofC[len2-1]
550+
lastProofType = lastProofTypeC
535551
}
536552

537553
// In most of cases, proofs are like this [BRANCH - (BRANCH, EXT)] -> [BRANCH - (BRANCH, EXT) - LEAF]
538-
// That means proofC only appends a leaf node (or a hashed node in some cases) to proofS.
539-
// In these cases, we don't need to add a placeholder branch
540-
skip_branch_placeholder := len1 == len2-1 && (lastProofTypeS != lastProofTypeC) && (lastProofTypeC == 3)
541-
542-
isModifiedExtNode := false
543-
if !skip_branch_placeholder {
544-
var curProofS []byte
554+
// That means proofC only appends a leaf node to proofS.
555+
// In such cases, we don't need to add a placeholder branch
556+
need_branch_placeholder := !(len1 == len2-1 && (lastProofTypeS != lastProofTypeC) && (lastProofTypeC == 3))
557+
if need_branch_placeholder {
558+
var extProofS []byte
545559
if mismatchedIdx != -1 {
546-
curProofS = proofS[mismatchedIdx]
560+
extProofS = proofS[mismatchedIdx]
547561
}
548562

549563
// This is a special case when the number of txs is 2.
550-
// In this case, proofS is a leaf and len1 is 1 but there is no nibbles
564+
// In this case, proofS is a leaf and len1 is 1, but there is no nibbles
551565
var lastExtNibbleS []byte
552566
if len(extNibblesS) != 0 {
553567
lastExtNibbleS = extNibblesS[len1-1]
554568
}
555569

556570
var branchNode Node
557-
isModifiedExtNode, _, _, branchNode =
558-
addBranchAndPlaceholder(
559-
proofS, proofC, lastExtNibbleS, extNibblesC[len2-1], curProofS,
560-
leafRow0, k, keyIndex, lastProofTypeS == 3)
571+
_, _, _, branchNode = addBranchAndPlaceholder(
572+
proofS, proofC, lastExtNibbleS, extNibblesC[len2-1], extProofS,
573+
leafRow0, k, keyIndex, lastProofType == 3)
561574
nodes = append(nodes, branchNode)
562575
}
563576

564577
var leafNode Node
565-
// Add a tx leaf after branch placeholder
566-
if !isModifiedExtNode {
567-
if skip_branch_placeholder && (lastProofTypeS == 3 || lastProofTypeS == 4) {
578+
// In stack trie proofs, the order of the type is the same except the case of modification extension node
579+
// So, we use `mismatchedIdx` to represent the case.
580+
if mismatchedIdx == -1 {
581+
// Add a tx leaf after branch placeholder
582+
if lastProofTypeS == 3 {
568583
leafNode = prepareTxLeafNode(txIdx, proofS[len1-1], proofC[len2-1], k, nil, isBranch(proofS[len1-1]), false, false)
569584
} else {
570585
leafNode = prepareTxLeafAndPlaceholderNode(txIdx, proofC[len2-1], k, false)
571586
}
587+
572588
} else {
573-
fmt.Println("MODIFIED EXT CASE, IGNORE NOW!!")
574-
// TODO might not fit our case bcs we have [EXT - BRANCH] --> [BRANCH - LEAF]
575-
// isSModExtension := false
576-
// if len2 > len1 {
577-
// isSModExtension = true
578-
// }
579-
// leafNode = prepareTxLeafAndPlaceholderNode(txIdx, proofC[len2-1], k, isSModExtension)
589+
fmt.Println("MODIFIED EXT CASE!!")
590+
leafNode = prepareTxLeafAndPlaceholderNode(txIdx, proofC[len2-1], k, true)
580591

581592
// When a proof element is a modified extension node (new extension node appears at the position
582593
// of the existing extension node), additional rows are added (extension node before and after
583594
// modification).
584-
// leafNode = equipLeafWithModExtensionNode(nil, leafNode, common.Address{0}, proofS, proofC, proofC,
585-
// extNibblesS, extNibblesC, k, nil,
586-
// keyIndex, numberOfNibbles, false, &toBeHashed)
595+
leafNode = equipLeafWithModExtensionNode(nil, leafNode, common.Address{byte(txIdx)}, proofS, proofC,
596+
extNibblesS, extNibblesC, proofS[mismatchedIdx], k, keyIndex, int(numberOfNibbles), false)
587597
}
588-
589598
nodes = append(nodes, leafNode)
590-
591599
}
592600

593601
return nodes
@@ -759,7 +767,7 @@ func convertProofToWitness(statedb *state.StateDB, addr common.Address, addrh []
759767
// of the existing extension node), additional rows are added (extension node before and after
760768
// modification).
761769
if isModifiedExtNode {
762-
leafNode = equipLeafWithModExtensionNode(statedb, leafNode, addr, proof1, proof2, nil, extNibblesS, extNibblesC,
770+
leafNode = equipLeafWithModExtensionNode(statedb, leafNode, addr, proof1, proof2, extNibblesS, extNibblesC, nil,
763771
key, keyIndex, numberOfNibbles, isAccountProof)
764772
}
765773
nodes = append(nodes, leafNode)

0 commit comments

Comments
 (0)