Merge 51ad188b30 into 485ff4bbff

triedb/pathdb/database_test.go

@@ -27,14 +27,17 @@ import (
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/internal/testrand"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 	"github.com/ethereum/go-ethereum/trie/trienode"
+	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/ethereum/go-ethereum/triedb/database"
 	"github.com/holiman/uint256"
 )
 
-func updateTrie(db *Database, stateRoot common.Hash, addrHash common.Hash, root common.Hash, dirties map[common.Hash][]byte) (common.Hash, *trienode.NodeSet) {
+func updateMerkleTrie(db *Database, stateRoot common.Hash, addrHash common.Hash, root common.Hash, dirties map[common.Hash][]byte) (common.Hash, *trienode.NodeSet) {
 	var id *trie.ID
 	if addrHash == (common.Hash{}) {
 		id = trie.StateTrieID(stateRoot)
@@ -55,11 +58,11 @@ func updateTrie(db *Database, stateRoot common.Hash, addrHash common.Hash, root
 	return tr.Commit(false)
 }
 
-func generateAccount(storageRoot common.Hash) types.StateAccount {
+func generateAccount(storageRoot common.Hash, codeHash []byte) types.StateAccount {
 	return types.StateAccount{
 		Nonce:    uint64(rand.Intn(100)),
 		Balance:  uint256.NewInt(rand.Uint64()),
-		CodeHash: testrand.Bytes(32),
+		CodeHash: codeHash,
 		Root:     storageRoot,
 	}
 }
@@ -72,23 +75,38 @@ const (
 )
 
 type genctx struct {
-	stateRoot     common.Hash
+	stateRoot     common.Hash                               // Parent state root
 	accounts      map[common.Hash][]byte                    // Keyed by the hash of account address
 	storages      map[common.Hash]map[common.Hash][]byte    // Keyed by the hash of account address and the hash of storage key
+	codes         map[common.Hash][]byte                    // Keyed by the contract code hash
 	accountOrigin map[common.Address][]byte                 // Keyed by the account address
 	storageOrigin map[common.Address]map[common.Hash][]byte // Keyed by the account address and the hash of storage key
-	nodes         *trienode.MergedNodeSet
+	nodes         *trienode.MergedNodeSet                   // Aggregated dirty trie nodes
+
+	// Verkle hasher fields
+	isVerkle bool
+	tr       *trie.VerkleTrie
 }
 
-func newCtx(stateRoot common.Hash) *genctx {
+func newCtx(stateRoot common.Hash, db database.NodeDatabase, isVerkle bool) *genctx {
-	return &genctx{
+	ctx := &genctx{
+		isVerkle:      isVerkle,
 		stateRoot:     stateRoot,
 		accounts:      make(map[common.Hash][]byte),
 		storages:      make(map[common.Hash]map[common.Hash][]byte),
+		codes:         make(map[common.Hash][]byte),
 		accountOrigin: make(map[common.Address][]byte),
 		storageOrigin: make(map[common.Address]map[common.Hash][]byte),
 		nodes:         trienode.NewMergedNodeSet(),
 	}
+	if isVerkle {
+		tr, err := trie.NewVerkleTrie(stateRoot, db, utils.NewPointCache(1024))
+		if err != nil {
+			panic(fmt.Errorf("failed to load trie, err: %w", err))
+		}
+		ctx.tr = tr
+	}
+	return ctx
 }
 
 func (ctx *genctx) storageOriginSet(rawStorageKey bool, t *tester) map[common.Address]map[common.Hash][]byte {
@@ -108,6 +126,7 @@ func (ctx *genctx) storageOriginSet(rawStorageKey bool, t *tester) map[common.Ad
 }
 
 type tester struct {
+	disk      ethdb.Database
 	db        *Database
 	roots     []common.Hash
 	preimages map[common.Hash][]byte
@@ -131,6 +150,7 @@ func newTester(t *testing.T, historyLimit uint64, isVerkle bool, layers int) *te
 		}, isVerkle)
 
 		obj = &tester{
+			disk:         disk,
 			db:           db,
 			preimages:    make(map[common.Hash][]byte),
 			accounts:     make(map[common.Hash][]byte),
@@ -141,10 +161,18 @@ func newTester(t *testing.T, historyLimit uint64, isVerkle bool, layers int) *te
 	)
 	for i := 0; i < layers; i++ {
 		var parent = types.EmptyRootHash
+		if isVerkle {
+			parent = types.EmptyVerkleHash
+		}
 		if len(obj.roots) != 0 {
 			parent = obj.roots[len(obj.roots)-1]
 		}
-		root, nodes, states := obj.generate(parent, i > 6)
+		// raw storage key is required in verkle for rollback
+		rawStorageKey := i > 6
+		if isVerkle {
+			rawStorageKey = true
+		}
+		root, nodes, states := obj.generate(parent, rawStorageKey, isVerkle)
 
 		if err := db.Update(root, parent, uint64(i), nodes, states); err != nil {
 			panic(fmt.Errorf("failed to update state changes, err: %w", err))
@@ -167,13 +195,33 @@ func (t *tester) release() {
 	t.db.diskdb.Close()
 }
 
-func (t *tester) randAccount() (common.Address, []byte) {
+func (t *tester) randAccount() (common.Address, []byte, []byte) {
 	for addrHash, account := range t.accounts {
-		return t.accountPreimage(addrHash), account
+		acct, err := types.FullAccount(account)
+		if err != nil {
+			panic(fmt.Errorf("failed to decode account, %v", err))
 		}
-	return common.Address{}, nil
+		return t.accountPreimage(addrHash), account, acct.CodeHash
+	}
+	return common.Address{}, nil, nil
 }
 
+func (t *tester) generateCode(ctx *genctx, addr common.Address) common.Hash {
+	size := rand.Intn(128 * 1024)
+	code := testrand.Bytes(size)
+	hash := crypto.Keccak256Hash(code)
+	ctx.codes[hash] = code
+
+	if ctx.isVerkle {
+		if err := ctx.tr.UpdateContractCode(addr, hash, code); err != nil {
+			panic(fmt.Errorf("failed to update contract code, %v", err))
+		}
+	}
+	return hash
+}
+
+// generateStorage inserts a batch of new storage slots for the specified account.
+// The storage is assumed to be empty before insertion.
 func (t *tester) generateStorage(ctx *genctx, addr common.Address) common.Hash {
 	var (
 		addrHash = crypto.Keccak256Hash(addr.Bytes())
@@ -189,29 +237,52 @@ func (t *tester) generateStorage(ctx *genctx, addr common.Address) common.Hash {
 		storage[hash] = v
 		origin[hash] = nil
 	}
-	root, set := updateTrie(t.db, ctx.stateRoot, addrHash, types.EmptyRootHash, storage)
 
 	ctx.storages[addrHash] = storage
 	ctx.storageOrigin[addr] = origin
+
+	// Generate a merkle storage trie for the newly constructed storage.
+	if !ctx.isVerkle {
+		root, set := updateMerkleTrie(t.db, ctx.stateRoot, addrHash, types.EmptyRootHash, storage)
 		ctx.nodes.Merge(set)
 		return root
 	}
+	// Insert the storage slots in the global verkle trie
+	for key, val := range storage {
+		if err := ctx.tr.UpdateStorage(addr, t.preimages[key], val); err != nil {
+			panic(fmt.Errorf("failed to update storage, %v", err))
+		}
+	}
+	return common.Hash{}
+}
 
+// mutateStorage modifies existing storage slots by deleting/updating some and
+// creating new ones, simulating a typical storage mutation.
 func (t *tester) mutateStorage(ctx *genctx, addr common.Address, root common.Hash) common.Hash {
 	var (
+		deletes int
+		updates int
+
 		addrHash = crypto.Keccak256Hash(addr.Bytes())
 		storage  = make(map[common.Hash][]byte)
 		origin   = make(map[common.Hash][]byte)
 	)
 	for hash, val := range t.storages[addrHash] {
+		if rand.Intn(2) == 0 {
 			origin[hash] = val
 			storage[hash] = nil
-
+			deletes++
-		if len(origin) == 3 {
+		} else {
+			origin[hash] = val
+			v, _ := rlp.EncodeToBytes(common.TrimLeftZeroes(testrand.Bytes(32)))
+			storage[hash] = v
+			updates++
+		}
+		if deletes >= 3 && updates >= 3 {
 			break
 		}
 	}
-	for i := 0; i < 3; i++ {
+	for i := 0; i < deletes; i++ {
 		v, _ := rlp.EncodeToBytes(common.TrimLeftZeroes(testrand.Bytes(32)))
 		key := testrand.Bytes(32)
 		hash := crypto.Keccak256Hash(key)
@@ -220,14 +291,34 @@ func (t *tester) mutateStorage(ctx *genctx, addr common.Address, root common.Has
 		storage[hash] = v
 		origin[hash] = nil
 	}
-	root, set := updateTrie(t.db, ctx.stateRoot, crypto.Keccak256Hash(addr.Bytes()), root, storage)
 
 	ctx.storages[addrHash] = storage
 	ctx.storageOrigin[addr] = origin
+
+	// Update the merkle storage trie with the generated mutation set
+	if !ctx.isVerkle {
+		root, set := updateMerkleTrie(t.db, ctx.stateRoot, crypto.Keccak256Hash(addr.Bytes()), root, storage)
 		ctx.nodes.Merge(set)
 		return root
 	}
+	// Apply the storage mutation into the global verkle trie
+	for key, val := range storage {
+		if len(val) != 0 {
+			if err := ctx.tr.UpdateStorage(addr, t.preimages[key], val); err != nil {
+				panic(fmt.Errorf("failed to update storage, %v", err))
+			}
+		} else {
+			// TODO(rjl493456442) here the zero markers are written instead of
+			// wiping the nodes from the trie.
+			if err := ctx.tr.DeleteStorage(addr, t.preimages[key]); err != nil {
+				panic(fmt.Errorf("failed to delete storage, %v", err))
+			}
+		}
+	}
+	return common.Hash{}
+}
 
+// clearStorage removes the existing storage slots belonging to the specified account.
 func (t *tester) clearStorage(ctx *genctx, addr common.Address, root common.Hash) common.Hash {
 	var (
 		addrHash = crypto.Keccak256Hash(addr.Bytes())
@@ -238,19 +329,32 @@ func (t *tester) clearStorage(ctx *genctx, addr common.Address, root common.Hash
 		origin[hash] = val
 		storage[hash] = nil
 	}
-	root, set := updateTrie(t.db, ctx.stateRoot, addrHash, root, storage)
+	ctx.storages[addrHash] = storage
+	ctx.storageOrigin[addr] = origin
+
+	// Clear the merkle storage trie with the generated deletion set
+	if !ctx.isVerkle {
+		root, set := updateMerkleTrie(t.db, ctx.stateRoot, addrHash, root, storage)
 		if root != types.EmptyRootHash {
 			panic("failed to clear storage trie")
 		}
-	ctx.storages[addrHash] = storage
-	ctx.storageOrigin[addr] = origin
 		ctx.nodes.Merge(set)
 		return root
 	}
+	// Apply the storage deletion into the global verkle trie
+	for key := range storage {
+		// TODO(rjl493456442) here the zero markers are written instead of
+		// wiping the nodes from the trie.
+		if err := ctx.tr.DeleteStorage(addr, t.preimages[key]); err != nil {
+			panic(fmt.Errorf("failed to delete storage, %v", err))
+		}
+	}
+	return common.Hash{}
+}
 
-func (t *tester) generate(parent common.Hash, rawStorageKey bool) (common.Hash, *trienode.MergedNodeSet, *StateSetWithOrigin) {
+func (t *tester) generate(parent common.Hash, rawStorageKey bool, isVerkle bool) (common.Hash, *trienode.MergedNodeSet, *StateSetWithOrigin) {
 	var (
-		ctx     = newCtx(parent)
+		ctx     = newCtx(parent, t.db, isVerkle)
 		dirties = make(map[common.Hash]struct{})
 	)
 	for i := 0; i < 20; i++ {
@@ -276,13 +380,14 @@ func (t *tester) generate(parent common.Hash, rawStorageKey bool) (common.Hash,
 			dirties[addrHash] = struct{}{}
 
 			root := t.generateStorage(ctx, addr)
-			ctx.accounts[addrHash] = types.SlimAccountRLP(generateAccount(root))
+			codeHash := t.generateCode(ctx, addr)
+			ctx.accounts[addrHash] = types.SlimAccountRLP(generateAccount(root, codeHash.Bytes()))
 			ctx.accountOrigin[addr] = nil
 			t.preimages[addrHash] = addr.Bytes()
 
 		case modifyAccountOp:
 			// account mutation
-			addr, account := t.randAccount()
+			addr, account, codeHash := t.randAccount()
 			if addr == (common.Address{}) {
 				continue
 			}
@@ -296,14 +401,14 @@ func (t *tester) generate(parent common.Hash, rawStorageKey bool) (common.Hash,
 
 			acct, _ := types.FullAccount(account)
 			stRoot := t.mutateStorage(ctx, addr, acct.Root)
-			newAccount := types.SlimAccountRLP(generateAccount(stRoot))
+			newAccount := types.SlimAccountRLP(generateAccount(stRoot, codeHash)) // TODO support contract code rewrite
 
 			ctx.accounts[addrHash] = newAccount
 			ctx.accountOrigin[addr] = account
 
 		case deleteAccountOp:
 			// account deletion
-			addr, account := t.randAccount()
+			addr, account, _ := t.randAccount()
 			if addr == (common.Address{}) {
 				continue
 			}
@@ -323,9 +428,6 @@ func (t *tester) generate(parent common.Hash, rawStorageKey bool) (common.Hash,
 			ctx.accountOrigin[addr] = account
 		}
 	}
-	root, set := updateTrie(t.db, parent, common.Hash{}, parent, ctx.accounts)
-	ctx.nodes.Merge(set)
-
 	// Save state snapshot before commit
 	t.snapAccounts[parent] = copyAccounts(t.accounts)
 	t.snapStorages[parent] = copyStorages(t.storages)
@@ -353,9 +455,32 @@ func (t *tester) generate(parent common.Hash, rawStorageKey bool) (common.Hash,
 			delete(t.storages, addrHash)
 		}
 	}
+	// Flush the contract code changes into the key value store
+	for codeHash, code := range ctx.codes {
+		rawdb.WriteCode(t.disk, codeHash, code)
+	}
+	// Flush the pending account changes into the global merkle trie
 	storageOrigin := ctx.storageOriginSet(rawStorageKey, t)
+	if !ctx.isVerkle {
+		root, set := updateMerkleTrie(t.db, parent, common.Hash{}, parent, ctx.accounts)
+		ctx.nodes.Merge(set)
 		return root, ctx.nodes, NewStateSetWithOrigin(ctx.accounts, ctx.storages, ctx.accountOrigin, storageOrigin, rawStorageKey)
 	}
+	// Flush the pending account changes into the global verkle trie
+	for addrHash, account := range ctx.accounts {
+		if len(account) == 0 {
+			// TODO(rjl493456442) here the zero markers are written instead of
+			// wiping the nodes from the trie.
+			ctx.tr.DeleteAccount(common.BytesToAddress(t.preimages[addrHash]))
+		} else {
+			ctx.tr.UpdateAccount(common.BytesToAddress(t.preimages[addrHash]), nil, 0)
+		}
+	}
+	root, nodes := ctx.tr.Commit(false)
+	merged := trienode.NewMergedNodeSet()
+	merged.Merge(nodes)
+	return root, merged, NewStateSetWithOrigin(ctx.accounts, ctx.storages, ctx.accountOrigin, storageOrigin, rawStorageKey)
+}
 
 // lastHash returns the latest root hash, or empty if nothing is cached.
 func (t *tester) lastHash() common.Hash {
@@ -365,7 +490,7 @@ func (t *tester) lastHash() common.Hash {
 	return t.roots[len(t.roots)-1]
 }
 
-func (t *tester) verifyState(root common.Hash) error {
+func (t *tester) verifyMerkleState(root common.Hash) error {
 	tr, err := trie.New(trie.StateTrieID(root), t.db)
 	if err != nil {
 		return err
@@ -465,7 +590,7 @@ func TestDatabaseRollback(t *testing.T) {
 			t.Fatalf("Failed to revert db, err: %v", err)
 		}
 		if i > 0 {
-			if err := tester.verifyState(parent); err != nil {
+			if err := tester.verifyMerkleState(parent); err != nil {
 				t.Fatalf("Failed to verify state, err: %v", err)
 			}
 		}
@@ -583,7 +708,7 @@ func TestCommit(t *testing.T) {
 		t.Fatal("Layer tree structure is invalid")
 	}
 	// Verify states
-	if err := tester.verifyState(tester.lastHash()); err != nil {
+	if err := tester.verifyMerkleState(tester.lastHash()); err != nil {
 		t.Fatalf("State is invalid, err: %v", err)
 	}
 	// Verify state histories
@@ -611,12 +736,12 @@ func TestJournal(t *testing.T) {
 	// Verify states including disk layer and all diff on top.
 	for i := 0; i < len(tester.roots); i++ {
 		if i >= tester.bottomIndex() {
-			if err := tester.verifyState(tester.roots[i]); err != nil {
+			if err := tester.verifyMerkleState(tester.roots[i]); err != nil {
 				t.Fatalf("Invalid state, err: %v", err)
 			}
 			continue
 		}
-		if err := tester.verifyState(tester.roots[i]); err == nil {
+		if err := tester.verifyMerkleState(tester.roots[i]); err == nil {
 			t.Fatal("Unexpected state")
 		}
 	}
@@ -647,12 +772,12 @@ func TestCorruptedJournal(t *testing.T) {
 	tester.db = New(tester.db.diskdb, nil, false)
 	for i := 0; i < len(tester.roots); i++ {
 		if tester.roots[i] == root {
-			if err := tester.verifyState(root); err != nil {
+			if err := tester.verifyMerkleState(root); err != nil {
 				t.Fatalf("Disk state is corrupted, err: %v", err)
 			}
 			continue
 		}
-		if err := tester.verifyState(tester.roots[i]); err == nil {
+		if err := tester.verifyMerkleState(tester.roots[i]); err == nil {
 			t.Fatal("Unexpected state")
 		}
 	}

triedb/pathdb/disklayer.go

@@ -300,7 +300,7 @@ func (dl *diskLayer) revert(h *history) (*diskLayer, error) {
 	// Apply the reverse state changes upon the current state. This must
 	// be done before holding the lock in order to access state in "this"
 	// layer.
-	nodes, err := apply(dl.db, h.meta.parent, h.meta.root, h.meta.version != stateHistoryV0, h.accounts, h.storages)
+	nodes, err := apply(dl.db.isVerkle, dl.db.diskdb, dl.db, h.meta.parent, h.meta.root, h.meta.version != stateHistoryV0, h.accounts, h.storages)
 	if err != nil {
 		return nil, err
 	}

triedb/pathdb/execute.go

@@ -17,180 +17,370 @@
 package pathdb
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 	"github.com/ethereum/go-ethereum/trie/trienode"
+	"github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/ethereum/go-ethereum/triedb/database"
 )
 
+// storageKey wraps a storage key with an additional flag, indicating whether
+// it is the raw storage key or its hash.
+type storageKey struct {
+	raw bool
+	key common.Hash
+}
+
+// hashKey returns the hash of raw storage key.
+func (k storageKey) hashKey(h *hasher) common.Hash {
+	if !k.raw {
+		return k.key
+	}
+	return h.hash(k.key.Bytes())
+}
+
+// stateHasher defines the essential functions needed for state hashing.
+type stateHasher interface {
+	// UpdateAccount abstracts an account write to the hasher.
+	UpdateAccount(address common.Address, blob []byte) error
+
+	// UpdateStorage abstracts a storage slot write to the hasher.
+	UpdateStorage(addr common.Address, key storageKey, value []byte) error
+
+	// DeleteAccount abstracts an account deletion from the hasher.
+	DeleteAccount(address common.Address) error
+
+	// DeleteStorage abstracts a storage slot deletion from the hasher.
+	DeleteStorage(addr common.Address, key storageKey) error
+
+	// Commit applies the pending mutations, rehashes the state, and collects
+	// all modified trie nodes, returning them along with the new state root.
+	Commit() (common.Hash, *trienode.MergedNodeSet, error)
+}
+
+// merkleHasher implements stateHasher in Merkle-Patricia-Trie manner.
+type merkleHasher struct {
+	db          database.NodeDatabase
+	sha256      *hasher
+	root        common.Hash
+	mainTr      *trie.Trie                     // The main account trie
+	subTries    map[common.Address]*trie.Trie  // The set of modified storage tries
+	newSubRoots map[common.Address]common.Hash // Expected storage trie roots after state transition
+}
+
+// newMerkleHasher constructs the merkle hasher with the given state root.
+func newMerkleHasher(db database.NodeDatabase, root common.Hash, sha256 *hasher) (*merkleHasher, error) {
+	tr, err := trie.New(trie.TrieID(root), db)
+	if err != nil {
+		return nil, err
+	}
+	return &merkleHasher{
+		db:          db,
+		sha256:      sha256,
+		root:        root,
+		mainTr:      tr,
+		subTries:    make(map[common.Address]*trie.Trie),
+		newSubRoots: make(map[common.Address]common.Hash),
+	}, nil
+}
+
+// UpdateAccount implements stateHasher, writing the provided account into
+// the trie.
+func (h *merkleHasher) UpdateAccount(address common.Address, blob []byte) error {
+	// The account was encoded in a slim format and is being converted
+	// to full format for the trie update.
+	acct, err := types.FullAccount(blob)
+	if err != nil {
+		return err
+	}
+	h.newSubRoots[address] = acct.Root
+
+	data, err := rlp.EncodeToBytes(acct)
+	if err != nil {
+		return err
+	}
+	return h.mainTr.Update(h.sha256.hash(address.Bytes()).Bytes(), data)
+}
+
+// DeleteAccount implements stateHasher, deleting the specified account from
+// the trie.
+func (h *merkleHasher) DeleteAccount(address common.Address) error {
+	key := h.sha256.hash(address.Bytes()).Bytes()
+	data, err := h.mainTr.Get(key)
+	if err != nil || len(data) == 0 {
+		return fmt.Errorf("the account to be deleted does not exist, %s", address.Hex())
+	}
+	h.newSubRoots[address] = types.EmptyRootHash
+	return h.mainTr.Delete(key)
+}
+
+// openStorageTrie opens the storage trie with the associated account address.
+func (h *merkleHasher) openStorageTrie(address common.Address, emptyAllowed bool) (*trie.Trie, error) {
+	var (
+		account  = types.NewEmptyStateAccount()
+		addrHash = h.sha256.hash(address.Bytes())
+	)
+	blob, err := h.mainTr.Get(addrHash.Bytes())
+	if err != nil {
+		return nil, err
+	}
+	if len(blob) == 0 && !emptyAllowed {
+		return nil, fmt.Errorf("account %x is not found", address)
+	}
+	if len(blob) != 0 {
+		if err := rlp.DecodeBytes(blob, &account); err != nil {
+			return nil, err
+		}
+	}
+	tr, err := trie.New(trie.StorageTrieID(h.root, addrHash, account.Root), h.db)
+	if err != nil {
+		return nil, err
+	}
+	h.subTries[address] = tr
+	return tr, nil
+}
+
+// UpdateStorage implements stateHasher, writing the provided storage slot into
+// the trie.
+func (h *merkleHasher) UpdateStorage(address common.Address, key storageKey, value []byte) error {
+	st, exist := h.subTries[address]
+	if !exist {
+		// Empty storage trie is allowed if the account was removed
+		// before and tries to add it back.
+		tr, err := h.openStorageTrie(address, true)
+		if err != nil {
+			return err
+		}
+		st = tr
+	}
+	return st.Update(key.hashKey(h.sha256).Bytes(), value)
+}
+
+// DeleteStorage implements stateHasher, deleting the specified storage slot from
+// the trie.
+func (h *merkleHasher) DeleteStorage(address common.Address, key storageKey) error {
+	st, exist := h.subTries[address]
+	if !exist {
+		// Empty storage trie is disallowed for storage deletion.
+		tr, err := h.openStorageTrie(address, false)
+		if err != nil {
+			return err
+		}
+		st = tr
+	}
+	return st.Delete(key.hashKey(h.sha256).Bytes())
+}
+
+// Commit implements stateHasher, gathering all modified trie nodes and returns
+// along with the new state root.
+func (h *merkleHasher) Commit() (common.Hash, *trienode.MergedNodeSet, error) {
+	merged := trienode.NewMergedNodeSet()
+	for address, tr := range h.subTries {
+		// Each modified storage trie must have a corresponding post-transition
+		// storage root cached.
+		newRoot, exist := h.newSubRoots[address]
+		if !exist {
+			return common.Hash{}, nil, fmt.Errorf("dangling dirty storage trie: %x", address)
+		}
+		root, nodes := tr.Commit(false)
+		if root != newRoot {
+			return common.Hash{}, nil, fmt.Errorf("unexpected storage root, want: %x, got: %x", newRoot, root)
+		}
+		if nodes != nil {
+			if err := merged.Merge(nodes); err != nil {
+				return common.Hash{}, nil, err
+			}
+		}
+	}
+	root, nodes := h.mainTr.Commit(false)
+	if nodes != nil {
+		if err := merged.Merge(nodes); err != nil {
+			return common.Hash{}, nil, err
+		}
+	}
+	return root, merged, nil
+}
+
+// verkleHasher implements stateHasher in Verkle-Trie manner.
+type verkleHasher struct {
+	db     database.NodeDatabase
+	disk   ethdb.KeyValueStore
+	root   common.Hash
+	tr     *trie.VerkleTrie
+	sha256 *hasher
+}
+
+// newVerkleHasher constructs the verkle hasher with the given state root.
+func newVerkleHasher(db database.NodeDatabase, disk ethdb.KeyValueStore, root common.Hash, sha256 *hasher) (*verkleHasher, error) {
+	tr, err := trie.NewVerkleTrie(root, db, utils.NewPointCache(1024)) // TODO use the shared cache
+	if err != nil {
+		return nil, err
+	}
+	return &verkleHasher{
+		db:     db,
+		disk:   disk,
+		root:   root,
+		tr:     tr,
+		sha256: sha256,
+	}, nil
+}
+
+// UpdateAccount implements stateHasher, writing the provided account along with
+// the associated code length into the verkle trie.
+func (h *verkleHasher) UpdateAccount(address common.Address, blob []byte) error {
+	acct, err := types.FullAccount(blob)
+	if err != nil {
+		return err
+	}
+	var code []byte
+	if !bytes.Equal(acct.CodeHash, types.EmptyCodeHash.Bytes()) {
+		// Contract code is assumed to be available because:
+		//
+		// - There is no account deletion in Verkle, the scenario that account was
+		//   removed before and needs to be restored is unexpected;
+		//
+		// - The contract code in key-value store should be retained even if the
+		//   account is deleted;
+		code = rawdb.ReadCode(h.disk, common.BytesToHash(acct.CodeHash))
+		if len(code) == 0 {
+			return fmt.Errorf("account code is missing, address: %x, codeHash: %x", address, acct.CodeHash)
+		}
+	}
+	// TODO @gballet @rjl493456442
+	// (a) try to avoid unnecessary update if the code is not modified;
+	// (b) make sure the leftover code chunks can be deleted;
+	if err := h.tr.UpdateAccount(address, acct, len(code)); err != nil {
+		return err
+	}
+	return h.tr.UpdateContractCode(address, h.sha256.hash(code), code)
+}
+
+// DeleteAccount implements stateHasher, deleting the specified account from
+// the trie.
+func (h *verkleHasher) DeleteAccount(address common.Address) error {
+	return h.tr.RollBackAccount(address)
+}
+
+// UpdateStorage implements stateHasher, writing the provided storage slot into
+// the trie.
+func (h *verkleHasher) UpdateStorage(address common.Address, key storageKey, value []byte) error {
+	if !key.raw {
+		return errors.New("unexpected hashed storage key")
+	}
+	return h.tr.UpdateStorage(address, key.key.Bytes(), value)
+}
+
+// DeleteStorage implements stateHasher, deleting the specified storage slot from
+// the trie.
+func (h *verkleHasher) DeleteStorage(address common.Address, key storageKey) error {
+	if !key.raw {
+		return errors.New("unexpected hashed storage key")
+	}
+	// TODO(rjl493456442) rollback storage
+	return h.tr.DeleteStorage(address, key.key.Bytes())
+}
+
+// Commit implements stateHasher, gathering all modified trie nodes and returns
+// along with the new state root.
+func (h *verkleHasher) Commit() (common.Hash, *trienode.MergedNodeSet, error) {
+	merged := trienode.NewMergedNodeSet()
+	root, nodes := h.tr.Commit(false)
+	if nodes != nil {
+		merged.Merge(nodes)
+	}
+	return root, merged, nil
+}
+
 // context wraps all fields for executing state diffs.
 type context struct {
-	prevRoot      common.Hash
-	postRoot      common.Hash
 	accounts      map[common.Address][]byte
 	storages      map[common.Address]map[common.Hash][]byte
-	nodes         *trienode.MergedNodeSet
 	rawStorageKey bool
-
+	hasher        stateHasher
-	// TODO (rjl493456442) abstract out the state hasher
-	// for supporting verkle tree.
-	accountTrie *trie.Trie
 }
 
 // apply processes the given state diffs, updates the corresponding post-state
 // and returns the trie nodes that have been modified.
-func apply(db database.NodeDatabase, prevRoot common.Hash, postRoot common.Hash, rawStorageKey bool, accounts map[common.Address][]byte, storages map[common.Address]map[common.Hash][]byte) (map[common.Hash]map[string]*trienode.Node, error) {
+func apply(isVerkle bool, disk ethdb.KeyValueStore, nodeDb database.NodeDatabase, prevRoot common.Hash, postRoot common.Hash, rawStorageKey bool, accounts map[common.Address][]byte, storages map[common.Address]map[common.Hash][]byte) (map[common.Hash]map[string]*trienode.Node, error) {
-	tr, err := trie.New(trie.TrieID(postRoot), db)
+	var (
+		err    error
+		hr     stateHasher
+		sha256 = newHasher()
+	)
+	defer sha256.release()
+
+	if isVerkle {
+		hr, err = newVerkleHasher(nodeDb, disk, postRoot, sha256)
+	} else {
+		hr, err = newMerkleHasher(nodeDb, postRoot, sha256)
+	}
 	if err != nil {
 		return nil, err
 	}
 	ctx := &context{
-		prevRoot:      prevRoot,
-		postRoot:      postRoot,
 		accounts:      accounts,
 		storages:      storages,
-		accountTrie:   tr,
 		rawStorageKey: rawStorageKey,
-		nodes:         trienode.NewMergedNodeSet(),
+		hasher:        hr,
 	}
 	for addr, account := range accounts {
 		var err error
 		if len(account) == 0 {
-			err = deleteAccount(ctx, db, addr)
+			err = deleteAccount(ctx, addr)
 		} else {
-			err = updateAccount(ctx, db, addr)
+			err = updateAccount(ctx, addr)
 		}
 		if err != nil {
 			return nil, fmt.Errorf("failed to revert state, err: %w", err)
 		}
 	}
-	root, result := tr.Commit(false)
+	root, merged, err := ctx.hasher.Commit()
+	if err != nil {
+		return nil, err
+	}
 	if root != prevRoot {
 		return nil, fmt.Errorf("failed to revert state, want %#x, got %#x", prevRoot, root)
 	}
-	if err := ctx.nodes.Merge(result); err != nil {
+	return merged.Flatten(), nil
-		return nil, err
-	}
-	return ctx.nodes.Flatten(), nil
 }
 
 // updateAccount the account was present in prev-state, and may or may not
 // existent in post-state. Apply the reverse diff and verify if the storage
 // root matches the one in prev-state account.
-func updateAccount(ctx *context, db database.NodeDatabase, addr common.Address) error {
+func updateAccount(ctx *context, addr common.Address) error {
-	// The account was present in prev-state, decode it from the
-	// 'slim-rlp' format bytes.
-	h := newHasher()
-	defer h.release()
-
-	addrHash := h.hash(addr.Bytes())
-	prev, err := types.FullAccount(ctx.accounts[addr])
-	if err != nil {
-		return err
-	}
-	// The account may or may not existent in post-state, try to
-	// load it and decode if it's found.
-	blob, err := ctx.accountTrie.Get(addrHash.Bytes())
-	if err != nil {
-		return err
-	}
-	post := types.NewEmptyStateAccount()
-	if len(blob) != 0 {
-		if err := rlp.DecodeBytes(blob, &post); err != nil {
-			return err
-		}
-	}
-	// Apply all storage changes into the post-state storage trie.
-	st, err := trie.New(trie.StorageTrieID(ctx.postRoot, addrHash, post.Root), db)
-	if err != nil {
-		return err
-	}
 	for key, val := range ctx.storages[addr] {
-		tkey := key
-		if ctx.rawStorageKey {
-			tkey = h.hash(key.Bytes())
-		}
 		var err error
 		if len(val) == 0 {
-			err = st.Delete(tkey.Bytes())
+			err = ctx.hasher.DeleteStorage(addr, storageKey{key: key, raw: ctx.rawStorageKey})
 		} else {
-			err = st.Update(tkey.Bytes(), val)
+			err = ctx.hasher.UpdateStorage(addr, storageKey{key: key, raw: ctx.rawStorageKey}, val)
 		}
 		if err != nil {
 			return err
 		}
 	}
-	root, result := st.Commit(false)
+	return ctx.hasher.UpdateAccount(addr, ctx.accounts[addr])
-	if root != prev.Root {
-		return errors.New("failed to reset storage trie")
-	}
-	// The returned set can be nil if storage trie is not changed
-	// at all.
-	if result != nil {
-		if err := ctx.nodes.Merge(result); err != nil {
-			return err
-		}
-	}
-	// Write the prev-state account into the main trie
-	full, err := rlp.EncodeToBytes(prev)
-	if err != nil {
-		return err
-	}
-	return ctx.accountTrie.Update(addrHash.Bytes(), full)
 }
 
 // deleteAccount the account was not present in prev-state, and is expected
 // to be existent in post-state. Apply the reverse diff and verify if the
 // account and storage is wiped out correctly.
-func deleteAccount(ctx *context, db database.NodeDatabase, addr common.Address) error {
+func deleteAccount(ctx *context, addr common.Address) error {
-	// The account must be existent in post-state, load the account.
-	h := newHasher()
-	defer h.release()
-
-	addrHash := h.hash(addr.Bytes())
-	blob, err := ctx.accountTrie.Get(addrHash.Bytes())
-	if err != nil {
-		return err
-	}
-	if len(blob) == 0 {
-		return fmt.Errorf("account is non-existent %#x", addrHash)
-	}
-	var post types.StateAccount
-	if err := rlp.DecodeBytes(blob, &post); err != nil {
-		return err
-	}
-	st, err := trie.New(trie.StorageTrieID(ctx.postRoot, addrHash, post.Root), db)
-	if err != nil {
-		return err
-	}
 	for key, val := range ctx.storages[addr] {
 		if len(val) != 0 {
-			return errors.New("expect storage deletion")
+			return fmt.Errorf("unexpected storage update, addr: %x, key: %x, val: %v", addr, key, val)
 		}
-		tkey := key
+		if err := ctx.hasher.DeleteStorage(addr, storageKey{key: key, raw: ctx.rawStorageKey}); err != nil {
-		if ctx.rawStorageKey {
-			tkey = h.hash(key.Bytes())
-		}
-		if err := st.Delete(tkey.Bytes()); err != nil {
 			return err
 		}
 	}
-	root, result := st.Commit(false)
+	return ctx.hasher.DeleteAccount(addr)
-	if root != types.EmptyRootHash {
-		return errors.New("failed to clear storage trie")
-	}
-	// The returned set can be nil if storage trie is not changed
-	// at all.
-	if result != nil {
-		if err := ctx.nodes.Merge(result); err != nil {
-			return err
-		}
-	}
-	// Delete the post-state account from the main trie.
-	return ctx.accountTrie.Delete(addrHash.Bytes())
 }

triedb/pathdb/history_test.go

@@ -43,7 +43,7 @@ func randomStateSet(n int) (map[common.Address][]byte, map[common.Address]map[co
 			v, _ := rlp.EncodeToBytes(common.TrimLeftZeroes(testrand.Bytes(32)))
 			storages[addr][testrand.Hash()] = v
 		}
-		account := generateAccount(types.EmptyRootHash)
+		account := generateAccount(types.EmptyRootHash, types.EmptyCodeHash.Bytes())
 		accounts[addr] = types.SlimAccountRLP(account)
 	}
 	return accounts, storages