perf: Share proposal between hash and commit

graft/coreth/core/blockchain.go

@@ -1851,6 +1851,9 @@ func (bc *BlockChain) reprocessState(current *types.Block, reexec uint64) error
 
 	// If the state is already available and the acceptor tip is up to date, skip re-processing.
 	if bc.HasState(current.Root()) && acceptorTipUpToDate {
+		if t, ok := bc.triedb.Backend().(*firewood.TrieDB); ok {
+			t.SetHashAndHeight(current.Hash(), current.NumberU64())
+		}
 		log.Info("Skipping state reprocessing", "root", current.Root())
 		return nil
 	}
@@ -1902,6 +1905,9 @@ func (bc *BlockChain) reprocessState(current *types.Block, reexec uint64) error
 	)
 	// Note: we add 1 since in each iteration, we attempt to re-execute the next block.
 	log.Info("Re-executing blocks to generate state for last accepted block", "from", current.NumberU64()+1, "to", origin)
+	if t, ok := bc.triedb.Backend().(*firewood.TrieDB); ok {
+		t.SetHashAndHeight(current.Hash(), current.NumberU64())
+	}
 	var roots []common.Hash
 	for current.NumberU64() < origin {
 		// TODO: handle canceled context

graft/coreth/core/genesis.go

@@ -38,6 +38,7 @@ import (
 	"github.com/ava-labs/avalanchego/graft/coreth/params"
 	"github.com/ava-labs/avalanchego/graft/coreth/plugin/evm/customtypes"
 	"github.com/ava-labs/avalanchego/graft/coreth/plugin/evm/upgrade/ap3"
+	"github.com/ava-labs/avalanchego/graft/evm/firewood"
 	"github.com/ava-labs/avalanchego/graft/evm/triedb/pathdb"
 	"github.com/ava-labs/avalanchego/vms/evm/acp226"
 	"github.com/ava-labs/libevm/common"
@@ -343,6 +344,10 @@ func (g *Genesis) toBlock(db ethdb.Database, triedb *triedb.Database) *types.Blo
 		if err := triedb.Commit(root, true); err != nil {
 			panic(fmt.Sprintf("unable to commit genesis block: %v", err))
 		}
+	} else {
+		if t, ok := triedb.Backend().(*firewood.TrieDB); ok {
+			t.SetHashAndHeight(block.Hash(), 0)
+		}
 	}
 	return block
 }

graft/evm/firewood/account_trie.go

@@ -20,13 +20,14 @@ import (
 
 var _ state.Trie = (*accountTrie)(nil)
 
-// accountTrie implements state.Trie for managing account states.
-// There are a couple caveats to the current implementation:
-//  1. `Commit` is not used as expected in the state package. The `StorageTrie` doesn't return
-//     values, and we thus rely on the `accountTrie`.
-//  2. The `Hash` method actually creates the proposal, since Firewood cannot calculate
-//     the hash of the trie without committing it. It is immediately dropped, and this
-//     can likely be optimized.
+// accountTrie implements [state.Trie] for managing account states.
+// Although it fulfills the [state.Trie] interface, it has some important differences:
+//  1. [accountTrie.Commit] is not used as expected in the state package. The `StorageTrie` doesn't return
+//     values, and we thus rely on the `accountTrie`. Additionally, no [trienode.NodeSet] is
+//     actually constructed, since Firewood manages nodes internally and the list of changes
+//     is not needed externally.
+//  2. The [accountTrie.Hash] method actually creates the [ffi.Proposal], since Firewood cannot calculate
+//     the hash of the trie without committing it.
 //
 // Note this is not concurrent safe.
 type accountTrie struct {
@@ -172,7 +173,7 @@ func (a *accountTrie) DeleteAccount(addr common.Address) error {
 	// Queue the key for deletion
 	a.dirtyKeys[string(key)] = nil
 	a.updateKeys = append(a.updateKeys, key)
-	a.updateValues = append(a.updateValues, nil) // Nil value indicates deletion
+	a.updateValues = append(a.updateValues, nil) // Must use nil to indicate deletion
 	a.hasChanges = true                          // Mark that there are changes to commit
 	return nil
 }
@@ -188,14 +189,16 @@ func (a *accountTrie) DeleteStorage(addr common.Address, key []byte) error {
 	// Queue the key for deletion
 	a.dirtyKeys[string(combinedKey[:])] = nil
 	a.updateKeys = append(a.updateKeys, combinedKey[:])
-	a.updateValues = append(a.updateValues, nil) // Nil value indicates deletion
+	a.updateValues = append(a.updateValues, nil) // Must use nil to indicate deletion
 	a.hasChanges = true                          // Mark that there are changes to commit
 	return nil
 }
 
 // Hash returns the current hash of the state trie.
-// This will create a proposal and drop it, so it is not efficient to call for each transaction.
+// This will create the necessary proposals to guarantee that the changes can
+// later be committed. All new proposals will be tracked by the [TrieDB].
 // If there are no changes since the last call, the cached root is returned.
+// On error, the zero hash is returned.
 func (a *accountTrie) Hash() common.Hash {
 	hash, err := a.hash()
 	if err != nil {
@@ -208,7 +211,7 @@ func (a *accountTrie) Hash() common.Hash {
 func (a *accountTrie) hash() (common.Hash, error) {
 	// If we haven't already hashed, we need to do so.
 	if a.hasChanges {
-		root, err := a.fw.getProposalHash(a.parentRoot, a.updateKeys, a.updateValues)
+		root, err := a.fw.createProposals(a.parentRoot, a.updateKeys, a.updateValues)
 		if err != nil {
 			return common.Hash{}, err
 		}
@@ -218,51 +221,51 @@ func (a *accountTrie) hash() (common.Hash, error) {
 	return a.root, nil
 }
 
-// Commit returns the new root hash of the trie and a NodeSet containing all modified accounts and storage slots.
-// The format of the NodeSet is different than in go-ethereum's trie implementation due to Firewood's design.
-// This boolean is ignored, as it is a relic of the StateTrie implementation.
+// Commit returns the new root hash of the trie and an empty [trienode.NodeSet].
+// The boolean input is ignored, as it is a relic of the StateTrie implementation.
+// If the changes are not yet already tracked by the [TrieDB], they are created.
 func (a *accountTrie) Commit(bool) (common.Hash, *trienode.NodeSet, error) {
 	// Get the hash of the trie.
+	// Ensures all changes are tracked by the Database.
 	hash, err := a.hash()
 	if err != nil {
 		return common.Hash{}, nil, err
 	}
 
-	// Create the NodeSet. This will be sent to `triedb.Update` later.
-	nodeset := trienode.NewNodeSet(common.Hash{})
-	for i, key := range a.updateKeys {
-		nodeset.AddNode(key, &trienode.Node{
-			Blob: a.updateValues[i],
-		})
-	}
-
-	return hash, nodeset, nil
+	set := trienode.NewNodeSet(common.Hash{})
+	return hash, set, nil
 }
 
 // UpdateContractCode implements state.Trie.
-// Contract code is controlled by rawdb, so we don't need to do anything here.
+// Contract code is controlled by `rawdb`, so we don't need to do anything here.
+// This always returns nil.
 func (*accountTrie) UpdateContractCode(common.Address, common.Hash, []byte) error {
 	return nil
 }
 
 // GetKey implements state.Trie.
-// This should not be used, since any user should not be accessing by raw key.
+// Preimages are not yet supported in Firewood.
+// It always returns nil.
 func (*accountTrie) GetKey([]byte) []byte {
 	return nil
 }
 
 // NodeIterator implements state.Trie.
 // Firewood does not support iterating over internal nodes.
+// This always returns an error.
 func (*accountTrie) NodeIterator([]byte) (trie.NodeIterator, error) {
 	return nil, errors.New("NodeIterator not implemented for Firewood")
 }
 
 // Prove implements state.Trie.
-// Firewood does not yet support providing key proofs.
+// Firewood does not support providing key proofs.
+// This always returns an error.
 func (*accountTrie) Prove([]byte, ethdb.KeyValueWriter) error {
 	return errors.New("Prove not implemented for Firewood")
 }
 
+// Copy creates a deep copy of the [accountTrie].
+// The [database.Reader] is shared, since it is read-only.
 func (a *accountTrie) Copy() *accountTrie {
 	// Create a new AccountTrie with the same root and reader
 	newTrie := &accountTrie{

graft/evm/firewood/storage_trie.go

@@ -23,22 +23,19 @@ func newStorageTrie(accountTrie *accountTrie) *storageTrie {
 	}
 }
 
-// Actual commit is handled by the account trie.
-// Return the old storage root as if there was no change since Firewood
-// will manage the hash calculations without it.
-// All changes are managed by the account trie.
+// Commit is a no-op for storage tries, as all changes are managed by the account trie.
+// It always returns a nil NodeSet and zero hash.
 func (*storageTrie) Commit(bool) (common.Hash, *trienode.NodeSet, error) {
 	return common.Hash{}, nil, nil
 }
 
-// Firewood doesn't require tracking storage roots inside of an account.
-// They will be updated in place when hashing of the proposal takes place.
+// Hash returns an empty hash, as the storage roots are managed internally to Firewood.
 func (*storageTrie) Hash() common.Hash {
 	return common.Hash{}
 }
 
-// Copy should never be called on a storage trie, as it is just a wrapper around the account trie.
-// Each storage trie should be re-opened with the account trie separately.
+// Copy returns nil, as storage tries do not need to be copied separately.
+// All usage of a copied storage trie should first ensure it is non-nil.
 func (*storageTrie) Copy() *storageTrie {
 	return nil
 }