core, trie, eth: refactor stacktrie constructor

This change enhances the stacktrie constructor by introducing an option struct. It also simplifies the `Hash` and `Commit` operations, getting rid of the special handling round root node.
axieinfinity · Nov 4, 2024 · 097f17b · 097f17b
1 parent a28c4a1
commit 097f17b
Show file tree

Hide file tree

Showing 5 changed files with 112 additions and 121 deletions.
diff --git a/core/state/snapshot/conversion.go b/core/state/snapshot/conversion.go
@@ -363,22 +363,16 @@ func generateTrieRoot(db ethdb.KeyValueWriter, scheme string, it Iterator, accou
 
 func stackTrieGenerate(db ethdb.KeyValueWriter, scheme string, owner common.Hash, in chan trieKV, out chan common.Hash) {
 
-	var nodeWriter trie.NodeWriteFunc
+	options := trie.NewStackTrieOptions()
 	// Implement nodeWriter in case db is existed otherwise let it be nil.
 	if db != nil {
-		nodeWriter = func(path []byte, hash common.Hash, blob []byte) {
+		options = options.WithWriter(func(path []byte, hash common.Hash, blob []byte) {
 			rawdb.WriteTrieNode(db, owner, path, hash, blob, scheme)
-		}
+		})
 	}
-	t := trie.NewStackTrie(nodeWriter)
+	t := trie.NewStackTrie(options)
 	for leaf := range in {
 		t.TryUpdate(leaf.key[:], leaf.value)
 	}
-	var root common.Hash
-	if db == nil {
-		root = t.Hash()
-	} else {
-		root, _ = t.Commit()
-	}
-	out <- root
+	out <- t.Commit()
 }
diff --git a/eth/protocols/snap/sync.go b/eth/protocols/snap/sync.go
@@ -721,9 +721,11 @@ func (s *Syncer) loadSyncStatus() {
 						s.accountBytes += common.StorageSize(len(key) + len(value))
 					},
 				}
-				task.genTrie = trie.NewStackTrie(func(path []byte, hash common.Hash, val []byte) {
-					rawdb.WriteTrieNode(task.genBatch, common.Hash{}, path, hash, val, s.scheme)
+				options := trie.NewStackTrieOptions()
+				options = options.WithWriter(func(path []byte, hash common.Hash, blob []byte) {
+					rawdb.WriteTrieNode(task.genBatch, common.Hash{}, path, hash, blob, s.scheme)
 				})
+				task.genTrie = trie.NewStackTrie(options)
 
 				for accountHash, subtasks := range task.SubTasks {
 					for _, subtask := range subtasks {
@@ -736,9 +738,12 @@ func (s *Syncer) loadSyncStatus() {
 							},
 						}
 						owner := accountHash // local assignment for stacktrie writer closure
-						subtask.genTrie = trie.NewStackTrie(func(path []byte, hash common.Hash, val []byte) {
-							rawdb.WriteTrieNode(subtask.genBatch, owner, path, hash, val, s.scheme)
+						options := trie.NewStackTrieOptions()
+
+						options = options.WithWriter(func(path []byte, hash common.Hash, blob []byte) {
+							rawdb.WriteTrieNode(subtask.genBatch, owner, path, hash, blob, s.scheme)
 						})
+						subtask.genTrie = trie.NewStackTrie(options)
 					}
 				}
 			}
@@ -787,14 +792,16 @@ func (s *Syncer) loadSyncStatus() {
 				s.accountBytes += common.StorageSize(len(key) + len(value))
 			},
 		}
+		options := trie.NewStackTrieOptions()
+		options = options.WithWriter(func(path []byte, hash common.Hash, blob []byte) {
+			rawdb.WriteTrieNode(batch, common.Hash{}, path, hash, blob, s.scheme)
+		})
 		s.tasks = append(s.tasks, &accountTask{
 			Next:     next,
 			Last:     last,
 			SubTasks: make(map[common.Hash][]*storageTask),
 			genBatch: batch,
-			genTrie: trie.NewStackTrie(func(path []byte, hash common.Hash, val []byte) {
-				rawdb.WriteTrieNode(batch, common.Hash{}, path, hash, val, s.scheme)
-			}),
+			genTrie:  trie.NewStackTrie(options),
 		})
 		log.Debug("Created account sync task", "from", next, "last", last)
 		next = common.BigToHash(new(big.Int).Add(last.Big(), common.Big1))
@@ -1975,14 +1982,16 @@ func (s *Syncer) processStorageResponse(res *storageResponse) {
 						},
 					}
 					owner := account // local assignment for stacktrie writer closure
+					options := trie.NewStackTrieOptions()
+					options = options.WithWriter(func(path []byte, hash common.Hash, blob []byte) {
+						rawdb.WriteTrieNode(batch, owner, path, hash, blob, s.scheme)
+					})
 					tasks = append(tasks, &storageTask{
 						Next:     common.Hash{},
 						Last:     r.End(),
 						root:     acc.Root,
 						genBatch: batch,
-						genTrie: trie.NewStackTrie(func(path []byte, hash common.Hash, val []byte) {
-							rawdb.WriteTrieNode(batch, owner, path, hash, val, s.scheme)
-						}),
+						genTrie:  trie.NewStackTrie(options),
 					})
 					for r.Next() {
 						batch := ethdb.HookedBatch{
@@ -1991,14 +2000,16 @@ func (s *Syncer) processStorageResponse(res *storageResponse) {
 								s.storageBytes += common.StorageSize(len(key) + len(value))
 							},
 						}
+						options := trie.NewStackTrieOptions()
+						options = options.WithWriter(func(path []byte, hash common.Hash, blob []byte) {
+							rawdb.WriteTrieNode(batch, owner, path, hash, blob, s.scheme)
+						})
 						tasks = append(tasks, &storageTask{
 							Next:     r.Start(),
 							Last:     r.End(),
 							root:     acc.Root,
 							genBatch: batch,
-							genTrie: trie.NewStackTrie(func(path []byte, hash common.Hash, val []byte) {
-								rawdb.WriteTrieNode(batch, owner, path, hash, val, s.scheme)
-							}),
+							genTrie:  trie.NewStackTrie(options),
 						})
 					}
 					for _, task := range tasks {
@@ -2044,9 +2055,11 @@ func (s *Syncer) processStorageResponse(res *storageResponse) {
 
 		if i < len(res.hashes)-1 || res.subTask == nil {
 			// no need to make local reassignment of account: this closure does not outlive the loop
-			tr := trie.NewStackTrie(func(path []byte, hash common.Hash, val []byte) {
-				rawdb.WriteTrieNode(batch, account, path, hash, val, s.scheme)
+			options := trie.NewStackTrieOptions()
+			options = options.WithWriter(func(path []byte, hash common.Hash, blob []byte) {
+				rawdb.WriteTrieNode(batch, account, path, hash, blob, s.scheme)
 			})
+			tr := trie.NewStackTrie(options)
 			for j := 0; j < len(res.hashes[i]); j++ {
 				tr.Update(res.hashes[i][j][:], res.slots[i][j])
 			}
@@ -2068,9 +2081,8 @@ func (s *Syncer) processStorageResponse(res *storageResponse) {
 	// Large contracts could have generated new trie nodes, flush them to disk
 	if res.subTask != nil {
 		if res.subTask.done {
-			if root, err := res.subTask.genTrie.Commit(); err != nil {
-				log.Error("Failed to commit stack slots", "err", err)
-			} else if root == res.subTask.root {
+			root := res.subTask.genTrie.Commit()
+			if root == res.subTask.root {
 				// If the chunk's root is an overflown but full delivery, clear the heal request
 				for i, account := range res.mainTask.res.hashes {
 					if account == res.accounts[len(res.accounts)-1] {
@@ -2286,9 +2298,7 @@ func (s *Syncer) forwardAccountTask(task *accountTask) {
 	// flush after finalizing task.done. It's fine even if we crash and lose this
 	// write as it will only cause more data to be downloaded during heal.
 	if task.done {
-		if _, err := task.genTrie.Commit(); err != nil {
-			log.Error("Failed to commit stack account", "err", err)
-		}
+		task.genTrie.Commit()
 	}
 	if task.genBatch.ValueSize() > ethdb.IdealBatchSize || task.done {
 		if err := task.genBatch.Write(); err != nil {

diff --git a/tests/fuzzers/stacktrie/trie_fuzzer.go b/tests/fuzzers/stacktrie/trie_fuzzer.go
@@ -153,9 +153,10 @@ func (f *fuzzer) fuzz() int {
 		trieA   = trie.NewEmpty(dbA)
 		spongeB = &spongeDb{sponge: sha3.NewLegacyKeccak256()}
 		dbB     = trie.NewDatabase(rawdb.NewDatabase(spongeB), nil)
-		trieB   = trie.NewStackTrie(func(path []byte, hash common.Hash, blob []byte) {
+		options = trie.NewStackTrieOptions().WithWriter(func(path []byte, hash common.Hash, blob []byte) {
 			rawdb.WriteTrieNode(spongeB, common.Hash{}, path, hash, blob, dbB.Scheme())
 		})
+		trieB       = trie.NewStackTrie(options)
 		vals        kvs
 		useful      bool
 		maxElements = 10000
@@ -203,9 +204,7 @@ func (f *fuzzer) fuzz() int {
 		trieB.Update(kv.k, kv.v)
 	}
 	rootB := trieB.Hash()
-	if _, err := trieB.Commit(); err != nil {
-		panic(err)
-	}
+	trieB.Commit()
 	if rootA != rootB {
 		panic(fmt.Sprintf("roots differ: (trie) %x != %x (stacktrie)", rootA, rootB))
 	}
@@ -217,19 +216,20 @@ func (f *fuzzer) fuzz() int {
 	// Ensure all the nodes are persisted correctly
 	// Need tracked deleted nodes.
 	var (
-		nodeset = make(map[string][]byte) // path -> blob
-		trieC   = trie.NewStackTrie(func(path []byte, hash common.Hash, blob []byte) {
+		nodeset  = make(map[string][]byte) // path -> blob
+		optionsC = trie.NewStackTrieOptions().WithWriter(func(path []byte, hash common.Hash, blob []byte) {
 			if crypto.Keccak256Hash(blob) != hash {
 				panic("invalid node blob")
 			}
 			nodeset[string(path)] = common.CopyBytes(blob)
 		})
+		trieC   = trie.NewStackTrie(optionsC)
 		checked int
 	)
 	for _, kv := range vals {
 		trieC.Update(kv.k, kv.v)
 	}
-	rootC, _ := trieC.Commit()
+	rootC := trieC.Commit()
 	if rootA != rootC {
 		panic(fmt.Sprintf("roots differ: (trie) %x != %x (stacktrie)", rootA, rootC))
 	}