From ec1d0a4ec0d421611bba7d0f019439d3d77c6ba5 Mon Sep 17 00:00:00 2001
From: joeycli <joeycli0919@gmail.com>
Date: Wed, 4 Sep 2024 19:11:38 +0800
Subject: [PATCH] feat: support version and root as key to state

fix: genesis version to -1

fix: copy with version

fix: change version type from uint64 to int64
---
 cmd/utils/cmd.go                     |  2 +-
 core/block_validator.go              |  4 ++--
 core/blockchain.go                   | 29 ++++++++++++------------
 core/blockchain_reader.go            | 33 +++++++++++++++++++---------
 core/genesis.go                      |  2 +-
 core/state/caching_versa_db.go       |  7 +++---
 core/state/database.go               |  4 ++--
 core/txpool/blobpool/blobpool.go     |  6 ++---
 core/txpool/blobpool/interface.go    |  2 +-
 core/txpool/legacypool/legacypool.go |  8 +++----
 eth/api_backend.go                   |  4 ++--
 eth/api_debug.go                     |  6 ++---
 eth/handler.go                       |  4 ++--
 eth/state_accessor.go                |  4 ++--
 eth/sync.go                          |  7 +++---
 internal/ethapi/api.go               |  6 ++---
 miner/miner.go                       |  2 +-
 miner/worker.go                      |  2 +-
 18 files changed, 73 insertions(+), 59 deletions(-)

diff --git a/cmd/utils/cmd.go b/cmd/utils/cmd.go
index 4b57164665..6bdef6c960 100644
--- a/cmd/utils/cmd.go
+++ b/cmd/utils/cmd.go
@@ -343,7 +343,7 @@ func missingBlocks(chain *core.BlockChain, blocks []*types.Block) []*types.Block
 			continue
 		}
 		// If we're above the chain head, state availability is a must
-		if !chain.HasBlockAndState(block.Hash(), block.NumberU64()) {
+		if !chain.HasBlockAndState(block.Hash(), block.Number().Int64()) {
 			return blocks[i:]
 		}
 	}
diff --git a/core/block_validator.go b/core/block_validator.go
index 19f64b2a7f..bf929cf0d4 100644
--- a/core/block_validator.go
+++ b/core/block_validator.go
@@ -103,7 +103,7 @@ func ValidateListsInBody(block *types.Block) error {
 // validated at this point.
 func (v *BlockValidator) ValidateBody(block *types.Block) error {
 	// Check whether the block is already imported.
-	if v.bc.HasBlockAndState(block.Hash(), block.NumberU64()) {
+	if v.bc.HasBlockAndState(block.Hash(), block.Number().Int64()) {
 		return ErrKnownBlock
 	}
 	if v.bc.isCachedBadBlock(block) {
@@ -149,7 +149,7 @@ func (v *BlockValidator) ValidateBody(block *types.Block) error {
 			return nil
 		},
 		func() error {
-			if !v.bc.HasBlockAndState(block.ParentHash(), block.NumberU64()-1) {
+			if !v.bc.HasBlockAndState(block.ParentHash(), block.Number().Int64()-1) {
 				if !v.bc.HasBlock(block.ParentHash(), block.NumberU64()-1) {
 					return consensus.ErrUnknownAncestor
 				}
diff --git a/core/blockchain.go b/core/blockchain.go
index c8d7f30008..8e9beb6229 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -429,7 +429,7 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, genesis *Genesis
 	// if there is no available state, waiting for state sync.
 	head := bc.CurrentBlock()
 	if bc.triedb.Scheme() != rawdb.VersionScheme {
-		if !bc.HasState(head.Root) {
+		if !bc.HasState(head.Number.Int64(), head.Root) {
 			if head.Number.Uint64() == 0 {
 				// The genesis state is missing, which is only possible in the path-based
 				// scheme. This situation occurs when the initial state sync is not finished
@@ -447,7 +447,7 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, genesis *Genesis
 				}
 				if bc.triedb.Scheme() == rawdb.PathScheme && !bc.NoTries() {
 					recoverable, _ := bc.triedb.Recoverable(diskRoot)
-					if !bc.HasState(diskRoot) && !recoverable {
+					if !bc.HasState(0, diskRoot) && !recoverable {
 						diskRoot = bc.triedb.Head()
 					}
 				}
@@ -776,8 +776,7 @@ func (bc *BlockChain) loadLastState() error {
 		archiveVersion, archiveRoot := versa.LatestStoreDiskVersionInfo()
 		// first start
 		if archiveVersion == -1 {
-			archiveVersion = 0
-			archiveRoot = bc.genesisBlock.Root()
+			log.Crit("please clear data, resync")
 		}
 
 		if int64(headBlock.NumberU64()) < archiveVersion {
@@ -989,7 +988,7 @@ func (bc *BlockChain) rewindHashHead(head *types.Header, root common.Hash) (*typ
 		}
 		// If the associated state is not reachable, continue searching
 		// backwards until an available state is found.
-		if !bc.HasState(head.Root) {
+		if !bc.HasState(head.Number.Int64(), head.Root) {
 			// If the chain is gapped in the middle, return the genesis
 			// block as the new chain head.
 			parent := bc.GetHeader(head.ParentHash, head.Number.Uint64()-1)
@@ -1029,7 +1028,7 @@ func (bc *BlockChain) rewindPathHead(head *types.Header, root common.Hash) (*typ
 
 		// noState represents if the target state requested for search
 		// is unavailable and impossible to be recovered.
-		noState = !bc.HasState(root) && !bc.stateRecoverable(root)
+		noState = !bc.HasState(head.Number.Int64(), root) && !bc.stateRecoverable(root)
 
 		start  = time.Now() // Timestamp the rewinding is restarted
 		logged = time.Now() // Timestamp last progress log was printed
@@ -1050,13 +1049,13 @@ func (bc *BlockChain) rewindPathHead(head *types.Header, root common.Hash) (*typ
 		// If the root threshold hasn't been crossed but the available
 		// state is reached, quickly determine if the target state is
 		// possible to be reached or not.
-		if !beyondRoot && noState && bc.HasState(head.Root) {
+		if !beyondRoot && noState && bc.HasState(head.Number.Int64(), head.Root) {
 			beyondRoot = true
 			log.Info("Disable the search for unattainable state", "root", root)
 		}
 		// Check if the associated state is available or recoverable if
 		// the requested root has already been crossed.
-		if beyondRoot && (bc.HasState(head.Root) || bc.stateRecoverable(head.Root)) {
+		if beyondRoot && (bc.HasState(head.Number.Int64(), head.Root) || bc.stateRecoverable(head.Root)) {
 			break
 		}
 		// If pivot block is reached, return the genesis block as the
@@ -1083,7 +1082,7 @@ func (bc *BlockChain) rewindPathHead(head *types.Header, root common.Hash) (*typ
 		}
 	}
 	// Recover if the target state if it's not available yet.
-	if !bc.HasState(head.Root) {
+	if !bc.HasState(head.Number.Int64(), head.Root) {
 		if err := bc.triedb.Recover(head.Root); err != nil {
 			log.Crit("Failed to rollback state", "err", err)
 		}
@@ -1178,7 +1177,7 @@ func (bc *BlockChain) setHeadBeyondRoot(head uint64, time uint64, root common.Ha
 			// the pivot point. In this scenario, there is no possible recovery
 			// approach except for rerunning a snap sync. Do nothing here until the
 			// state syncer picks it up.
-			if !bc.HasState(newHeadBlock.Root) {
+			if !bc.HasState(newHeadBlock.Number.Int64(), newHeadBlock.Root) {
 				if newHeadBlock.Number.Uint64() != 0 {
 					log.Crit("Chain is stateless at a non-genesis block")
 				}
@@ -1290,7 +1289,7 @@ func (bc *BlockChain) SnapSyncCommitHead(hash common.Hash) error {
 			return err
 		}
 	}
-	if !bc.NoTries() && !bc.HasState(root) {
+	if !bc.NoTries() && !bc.HasState(0, root) {
 		return fmt.Errorf("non existent state [%x..]", root[:4])
 	}
 	// If all checks out, manually set the head block.
@@ -2347,7 +2346,7 @@ func (bc *BlockChain) insertChain(chain types.Blocks, setHead bool) (int, error)
 			parent = bc.GetHeader(block.ParentHash(), block.NumberU64()-1)
 		}
 
-		bc.stateCache.SetVersion(int64(block.NumberU64()))
+		bc.stateCache.SetVersion(int64(block.NumberU64()) - 1)
 		statedb, err := state.NewWithSharedPool(parent.Root, bc.stateCache, bc.snaps)
 		if err != nil {
 			bc.stateCache.Release()
@@ -2615,7 +2614,7 @@ func (bc *BlockChain) insertSideChain(block *types.Block, it *insertIterator) (i
 		numbers []uint64
 	)
 	parent := it.previous()
-	for parent != nil && !bc.HasState(parent.Root) {
+	for parent != nil && !bc.HasState(parent.Number.Int64(), parent.Root) {
 		if bc.stateRecoverable(parent.Root) {
 			if err := bc.triedb.Recover(parent.Root); err != nil {
 				return 0, err
@@ -2682,7 +2681,7 @@ func (bc *BlockChain) recoverAncestors(block *types.Block) (common.Hash, error)
 		numbers []uint64
 		parent  = block
 	)
-	for parent != nil && !bc.HasState(parent.Root()) {
+	for parent != nil && !bc.HasState(parent.Number().Int64(), parent.Root()) {
 		if bc.stateRecoverable(parent.Root()) {
 			if err := bc.triedb.Recover(parent.Root()); err != nil {
 				return common.Hash{}, err
@@ -2953,7 +2952,7 @@ func (bc *BlockChain) SetCanonical(head *types.Block) (common.Hash, error) {
 	defer bc.chainmu.Unlock()
 
 	// Re-execute the reorged chain in case the head state is missing.
-	if !bc.HasState(head.Root()) {
+	if !bc.HasState(head.Number().Int64(), head.Root()) {
 		if latestValidHash, err := bc.recoverAncestors(head); err != nil {
 			return latestValidHash, err
 		}
diff --git a/core/blockchain_reader.go b/core/blockchain_reader.go
index a7c6144665..e4398b3afb 100644
--- a/core/blockchain_reader.go
+++ b/core/blockchain_reader.go
@@ -18,6 +18,7 @@ package core
 
 import (
 	"errors"
+	"fmt"
 	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -338,7 +339,7 @@ func (bc *BlockChain) GetTd(hash common.Hash, number uint64) *big.Int {
 }
 
 // HasState checks if state trie is fully present in the database or not.
-func (bc *BlockChain) HasState(hash common.Hash) bool {
+func (bc *BlockChain) HasState(number int64, hash common.Hash) bool {
 	if bc.NoTries() {
 		return bc.snaps != nil && bc.snaps.Snapshot(hash) != nil
 	}
@@ -348,18 +349,24 @@ func (bc *BlockChain) HasState(hash common.Hash) bool {
 			return true
 		}
 	}
-	return bc.stateCache.HasState(hash)
+	return bc.stateCache.HasState(number, hash)
 }
 
 // HasBlockAndState checks if a block and associated state trie is fully present
 // in the database or not, caching it if present.
-func (bc *BlockChain) HasBlockAndState(hash common.Hash, number uint64) bool {
+func (bc *BlockChain) HasBlockAndState(hash common.Hash, number int64) bool {
 	// Check first that the block itself is known
-	block := bc.GetBlock(hash, number)
-	if block == nil {
-		return false
+	var root common.Hash
+	if number < 0 {
+		root = types.EmptyRootHash
+	} else {
+		block := bc.GetBlock(hash, uint64(number))
+		if block == nil {
+			return false
+		}
+		root = block.Root()
 	}
-	return bc.HasState(block.Root())
+	return bc.HasState(number, root)
 }
 
 // stateRecoverable checks if the specified state is recoverable.
@@ -390,13 +397,19 @@ func (bc *BlockChain) ContractCodeWithPrefix(hash common.Hash) ([]byte, error) {
 
 // State returns a new mutable state based on the current HEAD block.
 func (bc *BlockChain) State() (*state.StateDB, error) {
-	return bc.StateAt(bc.CurrentBlock().Root)
+	return bc.StateAt(bc.CurrentBlock().Number.Int64(), bc.CurrentBlock().Root)
 }
 
 // StateAt returns a new mutable state based on a particular point in time.
-func (bc *BlockChain) StateAt(root common.Hash) (*state.StateDB, error) {
+func (bc *BlockChain) StateAt(number int64, root common.Hash) (*state.StateDB, error) {
 	// new state db with no need commit mode
-	stateDb, err := state.New(root, state.NewDatabaseWithNodeDB(bc.db, bc.triedb, false), bc.snaps)
+	has := bc.HasState(number, root)
+	if !has {
+		return nil, fmt.Errorf(fmt.Sprintf("do not has state, verison: %d, root: %s", number, root.String()))
+	}
+	sdb := state.NewDatabaseWithNodeDB(bc.db, bc.triedb, false)
+	sdb.SetVersion(number)
+	stateDb, err := state.New(root, sdb, bc.snaps)
 	if err != nil {
 		return nil, err
 	}
diff --git a/core/genesis.go b/core/genesis.go
index d72aba7959..6a58fc2015 100644
--- a/core/genesis.go
+++ b/core/genesis.go
@@ -158,7 +158,7 @@ func flushAlloc(ga *types.GenesisAlloc, db ethdb.Database, triedb *triedb.Databa
 		triedbConfig.NoTries = false
 	}
 	cachingdb := state.NewDatabaseWithNodeDB(db, triedb, true)
-	cachingdb.SetVersion(0)
+	cachingdb.SetVersion(-1)
 	defer cachingdb.Release()
 	statedb, err := state.New(types.EmptyRootHash, cachingdb, nil)
 	if err != nil {
diff --git a/core/state/caching_versa_db.go b/core/state/caching_versa_db.go
index e2699086d9..9661a7b1bb 100644
--- a/core/state/caching_versa_db.go
+++ b/core/state/caching_versa_db.go
@@ -60,6 +60,7 @@ func (cv *cachingVersaDB) Copy() Database {
 	cp.triedb = cv.triedb
 	cp.versionDB = cv.versionDB
 	cp.codeDB = cv.codeDB
+	cp.version = cv.version
 	cp.mode = versa.S_RW // it is important
 
 	// TODO:: maybe add lock for cv.root
@@ -109,8 +110,8 @@ func (cv *cachingVersaDB) CopyTrie(tr Trie) Trie {
 	return nil
 }
 
-func (cv *cachingVersaDB) HasState(root common.Hash) bool {
-	return cv.versionDB.HasState(root)
+func (cv *cachingVersaDB) HasState(version int64, root common.Hash) bool {
+	return cv.versionDB.HasState(version, root)
 }
 
 func (cv *cachingVersaDB) OpenTrie(root common.Hash) (Trie, error) {
@@ -214,7 +215,7 @@ func (cv *cachingVersaDB) Flush() error {
 }
 
 func (cv *cachingVersaDB) SetVersion(version int64) {
-	cv.version = version - 1
+	cv.version = version
 	//cv.debug = NewDebugVersionState(cv.codeDB, cv.versionDB)
 	//cv.debug.Version = version
 }
diff --git a/core/state/database.go b/core/state/database.go
index 26501023d4..1422d8ab53 100644
--- a/core/state/database.go
+++ b/core/state/database.go
@@ -87,7 +87,7 @@ type Database interface {
 	Copy() Database
 
 	// HasState returns the state data whether in the triedb.
-	HasState(root common.Hash) bool
+	HasState(version int64, root common.Hash) bool
 
 	// HasTreeExpired used for caching versa db, whether the state where the opened tree resides has been closed
 	HasTreeExpired(tr Trie) bool
@@ -402,7 +402,7 @@ func (db *cachingDB) Copy() Database {
 	return db
 }
 
-func (db *cachingDB) HasState(root common.Hash) bool {
+func (db *cachingDB) HasState(_ int64, root common.Hash) bool {
 	_, err := db.OpenTrie(root)
 	return err == nil
 }
diff --git a/core/txpool/blobpool/blobpool.go b/core/txpool/blobpool/blobpool.go
index a30f0e1417..c0c3e3850c 100644
--- a/core/txpool/blobpool/blobpool.go
+++ b/core/txpool/blobpool/blobpool.go
@@ -366,9 +366,9 @@ func (p *BlobPool) Init(gasTip uint64, head *types.Header, reserve txpool.Addres
 	// Initialize the state with head block, or fallback to empty one in
 	// case the head state is not available (might occur when node is not
 	// fully synced).
-	state, err := p.chain.StateAt(head.Root)
+	state, err := p.chain.StateAt(head.Number.Int64(), head.Root)
 	if err != nil {
-		state, err = p.chain.StateAt(types.EmptyRootHash)
+		state, err = p.chain.StateAt(-1, types.EmptyRootHash)
 	}
 	if err != nil {
 		return err
@@ -793,7 +793,7 @@ func (p *BlobPool) Reset(oldHead, newHead *types.Header) {
 		resettimeHist.Update(time.Since(start).Nanoseconds())
 	}(time.Now())
 
-	statedb, err := p.chain.StateAt(newHead.Root)
+	statedb, err := p.chain.StateAt(newHead.Number.Int64(), newHead.Root)
 	if err != nil {
 		log.Error("Failed to reset blobpool state", "err", err)
 		return
diff --git a/core/txpool/blobpool/interface.go b/core/txpool/blobpool/interface.go
index 6f296a54bd..e471617ddb 100644
--- a/core/txpool/blobpool/interface.go
+++ b/core/txpool/blobpool/interface.go
@@ -40,5 +40,5 @@ type BlockChain interface {
 	GetBlock(hash common.Hash, number uint64) *types.Block
 
 	// StateAt returns a state database for a given root hash (generally the head).
-	StateAt(root common.Hash) (*state.StateDB, error)
+	StateAt(number int64, root common.Hash) (*state.StateDB, error)
 }
diff --git a/core/txpool/legacypool/legacypool.go b/core/txpool/legacypool/legacypool.go
index 91cd01e7b4..295183c9e0 100644
--- a/core/txpool/legacypool/legacypool.go
+++ b/core/txpool/legacypool/legacypool.go
@@ -120,7 +120,7 @@ type BlockChain interface {
 	GetBlock(hash common.Hash, number uint64) *types.Block
 
 	// StateAt returns a state database for a given root hash (generally the head).
-	StateAt(root common.Hash) (*state.StateDB, error)
+	StateAt(number int64, root common.Hash) (*state.StateDB, error)
 }
 
 // Config are the configuration parameters of the transaction pool.
@@ -311,9 +311,9 @@ func (pool *LegacyPool) Init(gasTip uint64, head *types.Header, reserve txpool.A
 	// Initialize the state with head block, or fallback to empty one in
 	// case the head state is not available (might occur when node is not
 	// fully synced).
-	statedb, err := pool.chain.StateAt(head.Root)
+	statedb, err := pool.chain.StateAt(head.Number.Int64(), head.Root)
 	if err != nil {
-		statedb, err = pool.chain.StateAt(types.EmptyRootHash)
+		statedb, err = pool.chain.StateAt(-1, types.EmptyRootHash)
 	}
 	if err != nil {
 		return err
@@ -1492,7 +1492,7 @@ func (pool *LegacyPool) reset(oldHead, newHead *types.Header) {
 	if newHead == nil {
 		newHead = pool.chain.CurrentBlock() // Special case during testing
 	}
-	statedb, err := pool.chain.StateAt(newHead.Root)
+	statedb, err := pool.chain.StateAt(newHead.Number.Int64(), newHead.Root)
 	if err != nil {
 		log.Error("Failed to reset txpool state", "err", err)
 		return
diff --git a/eth/api_backend.go b/eth/api_backend.go
index ae698ade36..c7096e36b5 100644
--- a/eth/api_backend.go
+++ b/eth/api_backend.go
@@ -204,7 +204,7 @@ func (b *EthAPIBackend) StateAndHeaderByNumber(ctx context.Context, number rpc.B
 	if header == nil {
 		return nil, nil, errors.New("header not found")
 	}
-	stateDb, err := b.eth.BlockChain().StateAt(header.Root)
+	stateDb, err := b.eth.BlockChain().StateAt(header.Number.Int64(), header.Root)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -226,7 +226,7 @@ func (b *EthAPIBackend) StateAndHeaderByNumberOrHash(ctx context.Context, blockN
 		if blockNrOrHash.RequireCanonical && b.eth.blockchain.GetCanonicalHash(header.Number.Uint64()) != hash {
 			return nil, nil, errors.New("hash is not currently canonical")
 		}
-		stateDb, err := b.eth.BlockChain().StateAt(header.Root)
+		stateDb, err := b.eth.BlockChain().StateAt(header.Number.Int64(), header.Root)
 		if err != nil {
 			return nil, nil, err
 		}
diff --git a/eth/api_debug.go b/eth/api_debug.go
index baedd2632f..6185b093ac 100644
--- a/eth/api_debug.go
+++ b/eth/api_debug.go
@@ -81,7 +81,7 @@ func (api *DebugAPI) DumpBlock(blockNr rpc.BlockNumber) (state.Dump, error) {
 	if header == nil {
 		return state.Dump{}, fmt.Errorf("block #%d not found", blockNr)
 	}
-	stateDb, err := api.eth.BlockChain().StateAt(header.Root)
+	stateDb, err := api.eth.BlockChain().StateAt(header.Number.Int64(), header.Root)
 	if err != nil {
 		return state.Dump{}, err
 	}
@@ -167,7 +167,7 @@ func (api *DebugAPI) AccountRange(blockNrOrHash rpc.BlockNumberOrHash, start hex
 			if header == nil {
 				return state.Dump{}, fmt.Errorf("block #%d not found", number)
 			}
-			stateDb, err = api.eth.BlockChain().StateAt(header.Root)
+			stateDb, err = api.eth.BlockChain().StateAt(header.Number.Int64(), header.Root)
 			if err != nil {
 				return state.Dump{}, err
 			}
@@ -177,7 +177,7 @@ func (api *DebugAPI) AccountRange(blockNrOrHash rpc.BlockNumberOrHash, start hex
 		if block == nil {
 			return state.Dump{}, fmt.Errorf("block %s not found", hash.Hex())
 		}
-		stateDb, err = api.eth.BlockChain().StateAt(block.Root())
+		stateDb, err = api.eth.BlockChain().StateAt(block.Number().Int64(), block.Root())
 		if err != nil {
 			return state.Dump{}, err
 		}
diff --git a/eth/handler.go b/eth/handler.go
index 43ec78b256..7eed59ccbb 100644
--- a/eth/handler.go
+++ b/eth/handler.go
@@ -220,13 +220,13 @@ func newHandler(config *handlerConfig) (*handler, error) {
 				}
 				h.snapSync.Store(true)
 				log.Warn("Switch sync mode from full sync to snap sync", "reason", "snap sync incomplete")
-			} else if !h.chain.NoTries() && !h.chain.HasState(fullBlock.Root) {
+			} else if !h.chain.NoTries() && !h.chain.HasState(fullBlock.Number.Int64(), fullBlock.Root) {
 				h.snapSync.Store(true)
 				log.Warn("Switch sync mode from full sync to snap sync", "reason", "head state missing")
 			}
 		} else {
 			head := h.chain.CurrentBlock()
-			if head.Number.Uint64() > 0 && h.chain.HasState(head.Root) {
+			if head.Number.Uint64() > 0 && h.chain.HasState(head.Number.Int64(), head.Root) {
 				// Print warning log if database is not empty to run snap sync.
 				log.Warn("Switch sync mode from snap sync to full sync", "reason", "snap sync complete")
 			} else {
diff --git a/eth/state_accessor.go b/eth/state_accessor.go
index 6b045f2809..669cbb87a7 100644
--- a/eth/state_accessor.go
+++ b/eth/state_accessor.go
@@ -55,7 +55,7 @@ func (eth *Ethereum) hashState(ctx context.Context, block *types.Block, reexec u
 		// The state is available in live database, create a reference
 		// on top to prevent garbage collection and return a release
 		// function to deref it.
-		if statedb, err = eth.blockchain.StateAt(block.Root()); err == nil {
+		if statedb, err = eth.blockchain.StateAt(block.Number().Int64(), block.Root()); err == nil {
 			eth.blockchain.TrieDB().Reference(block.Root(), common.Hash{})
 			return statedb, func() {
 				eth.blockchain.TrieDB().Dereference(block.Root())
@@ -185,7 +185,7 @@ func (eth *Ethereum) hashState(ctx context.Context, block *types.Block, reexec u
 
 func (eth *Ethereum) pathState(block *types.Block) (*state.StateDB, func(), error) {
 	// Check if the requested state is available in the live chain.
-	statedb, err := eth.blockchain.StateAt(block.Root())
+	statedb, err := eth.blockchain.StateAt(block.Number().Int64(), block.Root())
 	if err == nil {
 		return statedb, noopReleaser, nil
 	}
diff --git a/eth/sync.go b/eth/sync.go
index 861eceb6c6..450d9b4c7b 100644
--- a/eth/sync.go
+++ b/eth/sync.go
@@ -17,6 +17,7 @@
 package eth
 
 import (
+	"fmt"
 	"math/big"
 	"time"
 
@@ -215,15 +216,15 @@ func (cs *chainSyncer) modeAndLocalHead() (downloader.SyncMode, *big.Int) {
 		// We are in a full sync, but the associated head state is missing. To complete
 		// the head state, forcefully rerun the snap sync. Note it doesn't mean the
 		// persistent state is corrupted, just mismatch with the head block.
-		if !cs.handler.chain.NoTries() && !cs.handler.chain.HasState(head.Root) {
+		if !cs.handler.chain.NoTries() && !cs.handler.chain.HasState(head.Number.Int64(), head.Root) {
 			block := cs.handler.chain.CurrentSnapBlock()
 			td := cs.handler.chain.GetTd(block.Hash(), block.Number.Uint64())
 			log.Info("Reenabled snap sync as chain is stateless")
 			return downloader.SnapSync, td
 		}
 	} else {
-		if !cs.handler.chain.HasState(head.Root) {
-			panic("version db not support snap sync")
+		if !cs.handler.chain.HasState(head.Number.Int64(), head.Root) {
+			panic(fmt.Sprintf("version db not support snap sync, version: %d, root: %s", head.Number.Int64(), head.Root))
 		}
 	}
 	// Nope, we're really full syncing
diff --git a/internal/ethapi/api.go b/internal/ethapi/api.go
index e31cee15ec..5598944b01 100644
--- a/internal/ethapi/api.go
+++ b/internal/ethapi/api.go
@@ -1367,11 +1367,11 @@ func (s *BlockChainAPI) needToReplay(ctx context.Context, block *types.Block, ac
 	if err != nil {
 		return false, fmt.Errorf("block not found for block number (%d): %v", block.NumberU64()-1, err)
 	}
-	parentState, err := s.b.Chain().StateAt(parent.Root())
+	parentState, err := s.b.Chain().StateAt(parent.Number().Int64(), parent.Root())
 	if err != nil {
 		return false, fmt.Errorf("statedb not found for block number (%d): %v", block.NumberU64()-1, err)
 	}
-	currentState, err := s.b.Chain().StateAt(block.Root())
+	currentState, err := s.b.Chain().StateAt(parent.Number().Int64(), block.Root())
 	if err != nil {
 		return false, fmt.Errorf("statedb not found for block number (%d): %v", block.NumberU64(), err)
 	}
@@ -1397,7 +1397,7 @@ func (s *BlockChainAPI) replay(ctx context.Context, block *types.Block, accounts
 	if err != nil {
 		return nil, nil, fmt.Errorf("block not found for block number (%d): %v", block.NumberU64()-1, err)
 	}
-	statedb, err := s.b.Chain().StateAt(parent.Root())
+	statedb, err := s.b.Chain().StateAt(parent.Number().Int64(), parent.Root())
 	if err != nil {
 		return nil, nil, fmt.Errorf("state not found for block number (%d): %v", block.NumberU64()-1, err)
 	}
diff --git a/miner/miner.go b/miner/miner.go
index 41f93ea388..deb5f3957a 100644
--- a/miner/miner.go
+++ b/miner/miner.go
@@ -245,7 +245,7 @@ func (miner *Miner) Pending() (*types.Block, *state.StateDB) {
 	if block == nil {
 		return nil, nil
 	}
-	stateDb, err := miner.worker.chain.StateAt(block.Root)
+	stateDb, err := miner.worker.chain.StateAt(block.Number.Int64(), block.Root)
 	if err != nil {
 		return nil, nil
 	}
diff --git a/miner/worker.go b/miner/worker.go
index 1dc05554f8..4976e503bb 100644
--- a/miner/worker.go
+++ b/miner/worker.go
@@ -690,7 +690,7 @@ func (w *worker) makeEnv(parent *types.Header, header *types.Header, coinbase co
 	prevEnv *environment) (*environment, error) {
 	// Retrieve the parent state to execute on top and start a prefetcher for
 	// the miner to speed block sealing up a bit
-	state, err := w.chain.StateAt(parent.Root)
+	state, err := w.chain.StateAt(parent.Number.Int64(), parent.Root)
 	if err != nil {
 		return nil, err
 	}