feat: triedb.Database.Update options via statedb.Commit

core/state/statedb.go

@@ -1152,7 +1152,7 @@ func (s *StateDB) handleDestruction(nodes *trienode.MergedNodeSet) (map[common.A
 //
 // The associated block number of the state transition is also provided
 // for more chain context.
-func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool, opts ...stateconf.SnapshotUpdateOption) (common.Hash, error) {
+func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool, opts ...stateconf.StateDBCommitOption) (common.Hash, error) {
 	// Short circuit in case any database failure occurred earlier.
 	if s.dbErr != nil {
 		return common.Hash{}, fmt.Errorf("commit aborted due to earlier error: %v", s.dbErr)
@@ -1242,7 +1242,7 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool, opts ...statecon
 		start := time.Now()
 		// Only update if there's a state transition (skip empty Clique blocks)
 		if parent := s.snap.Root(); parent != root {
-			if err := s.snaps.Update(root, parent, s.convertAccountSet(s.stateObjectsDestruct), s.accounts, s.storages, opts...); err != nil {
+			if err := s.snaps.Update(root, parent, s.convertAccountSet(s.stateObjectsDestruct), s.accounts, s.storages, stateconf.ExtractSnapshotUpdateOpts(opts...)...); err != nil {
 				log.Warn("Failed to update snapshot tree", "from", parent, "to", root, "err", err)
 			}
 			// Keep 128 diff layers in the memory, persistent layer is 129th.
@@ -1268,7 +1268,7 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool, opts ...statecon
 	if root != origin {
 		start := time.Now()
 		set := triestate.New(s.accountsOrigin, s.storagesOrigin, incomplete)
-		if err := s.db.TrieDB().Update(root, origin, block, nodes, set); err != nil {
+		if err := s.db.TrieDB().Update(root, origin, block, nodes, set, stateconf.ExtractTrieDBUpdateOpts(opts...)...); err != nil {
 			return common.Hash{}, err
 		}
 		s.originalRoot = root

core/state/statedb.libevm_test.go

@@ -26,23 +26,46 @@ import (
 	"github.com/ava-labs/libevm/core/rawdb"
 	"github.com/ava-labs/libevm/core/state/snapshot"
 	"github.com/ava-labs/libevm/core/types"
+	"github.com/ava-labs/libevm/ethdb"
 	"github.com/ava-labs/libevm/libevm/stateconf"
+	"github.com/ava-labs/libevm/trie"
+	"github.com/ava-labs/libevm/trie/trienode"
+	"github.com/ava-labs/libevm/trie/triestate"
+	"github.com/ava-labs/libevm/triedb"
+	"github.com/ava-labs/libevm/triedb/database"
+	"github.com/ava-labs/libevm/triedb/hashdb"
 )
 
 func TestStateDBCommitPropagatesOptions(t *testing.T) {
+	memdb := rawdb.NewMemoryDatabase()
+	triedb := triedb.NewDatabase(
+		memdb,
+		&triedb.Config{
+			DBOverride: func(_ ethdb.Database) triedb.DBOverride {
+				return &triedbRecorder{Database: hashdb.New(memdb, nil, &trie.MerkleResolver{})}
+			},
+		},
+	)
 	var rec snapTreeRecorder
-	sdb, err := New(types.EmptyRootHash, NewDatabase(rawdb.NewMemoryDatabase()), &rec)
+	sdb, err := New(types.EmptyRootHash, NewDatabaseWithNodeDB(memdb, triedb), &rec)
 	require.NoError(t, err, "New()")
 
 	// Ensures that rec.Update() will be called.
 	sdb.SetNonce(common.Address{}, 42)
 
-	const payload = "hello world"
-	opt := stateconf.WithUpdatePayload(payload)
-	_, err = sdb.Commit(0, false, opt)
-	require.NoErrorf(t, err, "%T.Commit(..., %T)", sdb, opt)
+	const snapshotPayload = "hello world"
+	const trieDBPayload = "goodbye world"
+	snapshotOpt := stateconf.WithSnapshotUpdatePayload(snapshotPayload)
+	triedbOpt := stateconf.WithTrieDBUpdatePayload(trieDBPayload)
+	_, err = sdb.Commit(0, false, stateconf.WithSnapshotUpdateOpts(snapshotOpt), stateconf.WithTrieDBUpdateOpts(triedbOpt))
+	require.NoErrorf(t, err, "%T.Commit(..., %T, %T)", sdb, snapshotOpt, triedbOpt)
 
-	assert.Equalf(t, payload, rec.gotPayload, "%T payload propagated via %T.Commit() to %T.Update()", opt, sdb, rec)
+	assert.Equalf(t, snapshotPayload, rec.gotPayload, "%T payload propagated via %T.Commit() to %T.Update()", snapshotOpt, sdb, rec)
+	innerTrieDB, ok := triedb.Backend().(*triedbRecorder)
+	if !ok {
+		t.Fatalf("expected %T to be a *triedbRecorder", triedb.Backend())
+	}
+	assert.Equalf(t, trieDBPayload, innerTrieDB.gotPayload, "%T payload propagated via %T.Commit() to %T.Update()", triedbOpt, sdb, rec)
 }
 
 type snapTreeRecorder struct {
@@ -59,7 +82,7 @@ func (r *snapTreeRecorder) Update(
 	_ map[common.Hash]struct{}, _ map[common.Hash][]byte, _ map[common.Hash]map[common.Hash][]byte,
 	opts ...stateconf.SnapshotUpdateOption,
 ) error {
-	r.gotPayload = stateconf.ExtractUpdatePayload(opts...)
+	r.gotPayload = stateconf.ExtractSnapshotUpdatePayload(opts...)
 	return nil
 }
 
@@ -78,3 +101,24 @@ func (snapshotStub) Account(common.Hash) (*types.SlimAccount, error) {
 func (snapshotStub) Root() common.Hash {
 	return common.Hash{}
 }
+
+type triedbRecorder struct {
+	*hashdb.Database
+	gotPayload any
+}
+
+func (r *triedbRecorder) Update(
+	root common.Hash,
+	parent common.Hash,
+	block uint64,
+	nodes *trienode.MergedNodeSet,
+	states *triestate.Set,
+	opts ...stateconf.TrieDBUpdateOption,
+) error {
+	r.gotPayload = stateconf.ExtractTrieDBUpdatePayload(opts...)
+	return r.Database.Update(root, parent, block, nodes, states)
+}
+
+func (r *triedbRecorder) Reader(_ common.Hash) (database.Reader, error) {
+	return r.Database.Reader(common.Hash{})
+}

libevm/stateconf/conf.go

@@ -17,7 +17,52 @@
 // Package stateconf configures state management.
 package stateconf
 
-import "github.com/ava-labs/libevm/libevm/options"
+import (
+	"github.com/ava-labs/libevm/libevm/options"
+)
+
+// A StateDBCommitOption configures the behaviour of all update calls within the
+// state.StateDB.Commit() implementations.
+// This is provided in two distinct types to allow customizability to both
+// snapshot and trie database updates, but are ignored in libevm implementations.
+type StateDBCommitOption = options.Option[stateDBCommitConfig]
+
+type stateDBCommitConfig struct {
+	snapshotOpts []SnapshotUpdateOption
+	triedbOpts   []TrieDBUpdateOption
+}
+
+// WithSnapshotUpdateOpts returns a StateDBCommitOption carrying a list of
+// SnapshotUpdateOptions.
+// If the list is not of length 1, the last option in the list is used.
+func WithSnapshotUpdateOpts(opts ...SnapshotUpdateOption) StateDBCommitOption {
+	return options.Func[stateDBCommitConfig](func(c *stateDBCommitConfig) {
+		// I don't like append() because there's no way to remove options, but that's a weakly held opinion
+		c.snapshotOpts = opts
+	})
+}
+
+// ExtractSnapshotUpdateOpts returns the list of SnapshotUpdateOptions carried
+// by the provided slice of StateDBCommitOption.
+func ExtractSnapshotUpdateOpts(opts ...StateDBCommitOption) []SnapshotUpdateOption {
+	return options.As(opts...).snapshotOpts
+}
+
+// WithTrieDBUpdateOpts returns a StateDBCommitOption carrying a list of
+// TrieDBUpdateOptions. If the list is not of length 1, the last option in the
+// list is used.
+func WithTrieDBUpdateOpts(opts ...TrieDBUpdateOption) StateDBCommitOption {
+	return options.Func[stateDBCommitConfig](func(c *stateDBCommitConfig) {
+		// I don't like append() because there's no way to remove options, but that's a weakly held opinion
+		c.triedbOpts = opts
+	})
+}
+
+// ExtractTrieDBUpdateOpts returns the list of TrieDBUpdateOptions carried by
+// the provided slice of StateDBCommitOption.
+func ExtractTrieDBUpdateOpts(opts ...StateDBCommitOption) []TrieDBUpdateOption {
+	return options.As(opts...).triedbOpts
+}
 
 // A SnapshotUpdateOption configures the behaviour of
 // state.SnapshotTree.Update() implementations. This will be removed along with
@@ -28,18 +73,41 @@ type snapshotUpdateConfig struct {
 	payload any
 }
 
-// WithUpdatePayload returns a SnapshotUpdateOption carrying an arbitrary
+// WithSnapshotUpdatePayload returns a SnapshotUpdateOption carrying an arbitrary
 // payload. It acts only as a carrier to exploit existing function plumbing and
 // the effect on behaviour is left to the implementation receiving it.
-func WithUpdatePayload(p any) SnapshotUpdateOption {
+func WithSnapshotUpdatePayload(p any) SnapshotUpdateOption {
 	return options.Func[snapshotUpdateConfig](func(c *snapshotUpdateConfig) {
 		c.payload = p
 	})
 }
 
-// ExtractUpdatePayload returns the payload carried by a [WithUpdatePayload]
+// ExtractSnapshotUpdatePayload returns the payload carried by a [WithSnapshotUpdatePayload]
+// option. Only one such option can be used at once; behaviour is otherwise
+// undefined.
+func ExtractSnapshotUpdatePayload(opts ...SnapshotUpdateOption) any {
+	return options.As(opts...).payload
+}
+
+// A TrieDBUpdateOption configures the behaviour of triedb.Database.Update() implementations.
+type TrieDBUpdateOption = options.Option[triedbUpdateConfig]
+
+type triedbUpdateConfig struct {
+	payload any
+}
+
+// WithTrieDBUpdatePayload returns a TrieDBUpdateOption carrying an arbitrary
+// payload. It acts only as a carrier to exploit existing function plumbing and
+// the effect on behaviour is left to the implementation receiving it.
+func WithTrieDBUpdatePayload(p any) TrieDBUpdateOption {
+	return options.Func[triedbUpdateConfig](func(c *triedbUpdateConfig) {
+		c.payload = p
+	})
+}
+
+// ExtractTrieDBUpdatePayload returns the payload carried by a [WithSnapshotUpdatePayload]
 // option. Only one such option can be used at once; behaviour is otherwise
 // undefined.
-func ExtractUpdatePayload(opts ...SnapshotUpdateOption) any {
+func ExtractTrieDBUpdatePayload(opts ...TrieDBUpdateOption) any {
 	return options.As(opts...).payload
 }

triedb/database.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/ava-labs/libevm/common"
 	"github.com/ava-labs/libevm/ethdb"
+	"github.com/ava-labs/libevm/libevm/stateconf"
 	"github.com/ava-labs/libevm/log"
 	"github.com/ava-labs/libevm/trie"
 	"github.com/ava-labs/libevm/trie/trienode"
@@ -70,7 +71,7 @@ type backend interface {
 	//
 	// The passed in maps(nodes, states) will be retained to avoid copying
 	// everything. Therefore, these maps must not be changed afterwards.
-	Update(root common.Hash, parent common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set) error
+	Update(root common.Hash, parent common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set, opts ...stateconf.TrieDBUpdateOption) error
 
 	// Commit writes all relevant trie nodes belonging to the specified state
 	// to disk. Report specifies whether logs will be displayed in info level.
@@ -148,11 +149,11 @@ func (db *Database) Reader(blockRoot common.Hash) (database.Reader, error) {
 //
 // The passed in maps(nodes, states) will be retained to avoid copying everything.
 // Therefore, these maps must not be changed afterwards.
-func (db *Database) Update(root common.Hash, parent common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set) error {
+func (db *Database) Update(root common.Hash, parent common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set, opts ...stateconf.TrieDBUpdateOption) error {
 	if db.preimages != nil {
 		db.preimages.commit(false)
 	}
-	return db.backend.Update(root, parent, block, nodes, states)
+	return db.backend.Update(root, parent, block, nodes, states, opts...)
 }
 
 // Commit iterates over all the children of a particular node, writes them out

triedb/hashdb/database.go

@@ -28,6 +28,7 @@ import (
 	"github.com/ava-labs/libevm/core/rawdb"
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/ethdb"
+	"github.com/ava-labs/libevm/libevm/stateconf"
 	"github.com/ava-labs/libevm/log"
 	"github.com/ava-labs/libevm/metrics"
 	"github.com/ava-labs/libevm/rlp"
@@ -548,7 +549,7 @@ func (db *Database) Initialized(genesisRoot common.Hash) bool {
 
 // Update inserts the dirty nodes in provided nodeset into database and link the
 // account trie with multiple storage tries if necessary.
-func (db *Database) Update(root common.Hash, parent common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set) error {
+func (db *Database) Update(root common.Hash, parent common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set, opts ...stateconf.TrieDBUpdateOption) error {
 	// Ensure the parent state is present and signal a warning if not.
 	if parent != types.EmptyRootHash {
 		if blob, _ := db.node(parent); len(blob) == 0 {

triedb/pathdb/database.go

@@ -27,6 +27,7 @@ import (
 	"github.com/ava-labs/libevm/core/rawdb"
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/ethdb"
+	"github.com/ava-labs/libevm/libevm/stateconf"
 	"github.com/ava-labs/libevm/log"
 	"github.com/ava-labs/libevm/params"
 	"github.com/ava-labs/libevm/trie/trienode"
@@ -223,7 +224,7 @@ func (db *Database) Reader(root common.Hash) (layer, error) {
 //
 // The passed in maps(nodes, states) will be retained to avoid copying everything.
 // Therefore, these maps must not be changed afterwards.
-func (db *Database) Update(root common.Hash, parentRoot common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set) error {
+func (db *Database) Update(root common.Hash, parentRoot common.Hash, block uint64, nodes *trienode.MergedNodeSet, states *triestate.Set, opts ...stateconf.TrieDBUpdateOption) error {
 	// Hold the lock to prevent concurrent mutations.
 	db.lock.Lock()
 	defer db.lock.Unlock()