add 2 overlay: post-bal and readerwithcache

core/blockchain.go

@@ -1902,7 +1902,15 @@ func (bc *BlockChain) processBlock(parentRoot common.Hash, block *types.Block, s
 
 	bc.cacheConfig.TrieCleanNoPrefetch = true
 	if bc.cacheConfig.TrieCleanNoPrefetch {
-		statedb, err = state.New(parentRoot, bc.statedb)
+		// statedb, err = state.New(parentRoot, bc.statedb)
+		// if err != nil {
+		// 	return nil, err
+		// }
+		reader, err := bc.statedb.ReaderWithCache(parentRoot)
+		if err != nil {
+			return nil, err
+		}
+		statedb, err = state.NewWithReader(parentRoot, bc.statedb, reader)
 		if err != nil {
 			return nil, err
 		}

core/parallel_state_processor.go

@@ -35,7 +35,7 @@ func NewParallelStateProcessor(config *params.ChainConfig, chain *HeaderChain) *
 }
 
 func (p *ParallelStateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg vm.Config) (*ProcessResult, error) {
-	fmt.Println("ParallelStateProcessor.Process called")
+	fmt.Println("ParallelStateProcessor.Process called:", block.NumberU64())
 	var (
 		header = block.Header()
 		gp     = new(GasPool).AddGas(block.GasLimit())
@@ -52,7 +52,8 @@ func (p *ParallelStateProcessor) Process(block *types.Block, statedb *state.Stat
 
 	if preStateType == BALPreState {
 		start := time.Now()
-		statedb.PrefetchStateBAL(block.NumberU64())
+		// statedb.PrefetchStateBAL(block.NumberU64())
+		statedb.PreComputePostState(block.NumberU64())
 		PrefetchBALTime += time.Since(start)
 	}
 
@@ -93,9 +94,8 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	case BALPreState:
 		{
 			start := time.Now()
-			statedb.MergePostBal()
+			// statedb.MergePostBal()
 			PrefetchMergeBALTime += time.Since(start)
-			preStateProvider = statedb
 		}
 
 	case SeqPreState:
@@ -116,24 +116,42 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	// Parallel executing the transaction
 	exeStart := time.Now()
 	postEntries := make([][]state.JournalEntry, len(block.Transactions()))
-	for i, tx := range block.Transactions() {
-		cleanStatedb, err := preStateProvider.PrestateAtIndex(i)
-		if err != nil {
-			return nil, err
-		}
 
+	initialdb := statedb.Copy()
+	for i, tx := range block.Transactions() {
 		i := i
-		gpcp := *gp
+
 		workers.Go(func() error {
+			var (
+				cleanStatedb *state.StateDB
+				err          error
+			)
+			switch preStateType {
+			case BALPreState:
+				{
+					cleanStatedb = initialdb.Copy()
+					cleanStatedb.SetTxContext(tx.Hash(), i)
+					err = cleanStatedb.SetTxBALReader()
+				}
+			case SeqPreState:
+				{
+					cleanStatedb, err = preStateProvider.PrestateAtIndex(i)
+					cleanStatedb.SetTxContext(tx.Hash(), i)
+				}
+			}
+			if err != nil {
+				return err
+			}
+
+			evm := vm.NewEVM(context, cleanStatedb, p.config, cfg)
+
 			usedGas := new(uint64)
 			msg, err := TransactionToMessage(tx, signer, header.BaseFee)
 			if err != nil {
 				return err
 			}
-			cleanStatedb.SetTxContext(tx.Hash(), i)
+			// todo: handle gp race
-
+			gpcp := *gp
-			evm := vm.NewEVM(context, cleanStatedb, p.config, cfg)
-
 			receipt, entries, err := ApplyTransactionWithParallelEVM(msg, &gpcp, cleanStatedb, blockNumber, blockHash, tx, usedGas, evm)
 			if err != nil {
 				return err

core/state/reader.go

@@ -18,6 +18,7 @@ package state
 
 import (
 	"errors"
+	"fmt"
 	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -476,3 +477,173 @@ func (r *readerWithCache) Storage(addr common.Address, slot common.Hash) (common
 
 	return value, nil
 }
+
+type readerWithBAL struct {
+	Reader  // Reader with cache
+	txIndex int
+	postBal map[int]types.TxPostValues // Shared accoss txs and mustn't be changed
+}
+
+func newReaderWithBAL(reader Reader, txIndex int, postBal map[int]types.TxPostValues) *readerWithBAL {
+	return &readerWithBAL{
+		reader, txIndex, postBal,
+	}
+}
+
+func (r *readerWithBAL) Account(addr common.Address) (*types.StateAccount, error) {
+	postVals, ok := r.postBal[r.txIndex-1]
+	if ok {
+		acctVal, ok := postVals[addr]
+		if ok {
+			// Cache codeHash
+			if acctVal.Code != nil {
+				if acctVal.CodeHash == nil {
+					acctVal.CodeHash = crypto.Keccak256(acctVal.Code)
+				}
+			}
+			if acctVal.Cached {
+				return &types.StateAccount{
+					Nonce:    acctVal.Nonce,
+					Balance:  acctVal.Balance,
+					Root:     acctVal.Root,
+					CodeHash: acctVal.CodeHash,
+				}, nil
+			}
+
+			acct, err := r.Reader.Account(addr)
+			acctVal.Cached = true
+			if acct == nil || err != nil {
+				// acct is nil, just return the postAcct
+				return &types.StateAccount{
+					Nonce:    acctVal.Nonce,
+					Balance:  acctVal.Balance,
+					CodeHash: acctVal.CodeHash,
+				}, nil
+			}
+
+			// acct mustn't be modified, or it'll change the cached pre-block values in readerwithcache
+			if acct.Nonce > acctVal.Nonce {
+				acctVal.Nonce = acct.Nonce
+			}
+			if acctVal.Balance == nil {
+				acctVal.Balance = acct.Balance
+			}
+			if acctVal.CodeHash == nil {
+				acctVal.CodeHash = acct.CodeHash
+			}
+			if acctVal.Root == (common.Hash{}) {
+				acctVal.Root = acct.Root
+			}
+			return &types.StateAccount{
+				Nonce:    acctVal.Nonce,
+				Balance:  acctVal.Balance,
+				Root:     acctVal.Root,
+				CodeHash: acctVal.CodeHash,
+			}, nil
+		}
+	}
+
+	acct, err := r.Reader.Account(addr)
+	if err != nil {
+		return nil, err
+	}
+	// cache account and codeHash except the first transaction
+	if postVals != nil && acct != nil {
+		postVals[addr] = &types.AcctPostValues{
+			Nonce:    acct.Nonce,
+			Balance:  acct.Balance,
+			Destruct: false,
+			CodeHash: acct.CodeHash,
+			Root:     acct.Root,
+			Cached:   true,
+		}
+	}
+	return acct, nil
+}
+
+func (r *readerWithBAL) Storage(addr common.Address, slot common.Hash) (common.Hash, error) {
+	postVals, ok := r.postBal[r.txIndex-1]
+	if ok {
+		acct, ok := postVals[addr]
+		if ok {
+			val, ok := acct.StorageKV[slot]
+			if ok {
+				return val, nil
+			}
+		}
+	}
+
+	val, err := r.Reader.Storage(addr, slot)
+	// Cache storage
+	if postVals != nil && postVals[addr] != nil {
+		if postVals[addr].StorageKV != nil {
+			postVals[addr].StorageKV[slot] = val
+		} else {
+			postVals[addr].StorageKV = map[common.Hash]common.Hash{
+				slot: val,
+			}
+		}
+	}
+	if err != nil {
+		return common.Hash{}, err
+	}
+
+	return val, nil
+}
+
+func (r *readerWithBAL) Code(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	postVals, ok := r.postBal[r.txIndex-1]
+	if ok {
+		acctVal, ok := postVals[addr]
+		if ok {
+			code := acctVal.Code
+			if code != nil {
+				return code, nil
+			}
+			if acctVal.CodeHash != nil {
+				code, err := r.Reader.Code(addr, common.Hash(acctVal.CodeHash))
+				// Cache code
+				acctVal.Code = code
+				if err != nil {
+					return nil, fmt.Errorf("cannot get code from code for addr:%x, err:%v", addr, err)
+				}
+				return code, nil
+			}
+		}
+	}
+	acct, err := r.Reader.Account(addr)
+	if err != nil || acct == nil {
+		return nil, fmt.Errorf("cannot get code from acct for addr:%x, err:%v", addr, err)
+	}
+	code, err := r.Reader.Code(addr, common.Hash(acct.CodeHash))
+	// Cache account, codeHash and code except the first transaction
+	// if postVals != nil {
+	// 	if postVals[addr] != nil {
+	// 		postVals[addr].Code = code
+	// 		postVals[addr].CodeHash = codeHash[:]
+	// 	} else {
+	// 		postVals[addr] = &types.AcctPostValues{
+	// 			Nonce:    acct.Nonce,
+	// 			Balance:  acct.Balance,
+	// 			Code:     code,
+	// 			Destruct: false,
+	// 			CodeHash: acct.CodeHash,
+	// 		}
+	// 	}
+	// }
+
+	if err != nil {
+		return nil, fmt.Errorf("cannot get code from codeHash for addr:%x, err:%v", addr, err)
+	}
+	return code, nil
+
+}
+
+func (r *readerWithBAL) CodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	// Here codeHash is not need, cause we'll fetch it from BAL reader
+	code, err := r.Code(addr, common.Hash{})
+	if err != nil {
+		return 0, err
+	}
+	return len(code), nil
+}

core/state/statedb.go

@@ -164,6 +164,7 @@ type StateDB struct {
 	blockNumber uint64
 	// postState after appling tx
 	postStates map[int]*StateDB
+	postBAL    map[int]types.TxPostValues
 }
 
 type BALType int
@@ -178,26 +179,15 @@ const (
 )
 
 type BALs struct {
-	Pre  map[uint64]types.AccessList     `json:"pre"`
+	Pre  map[uint64]types.AccessList           `json:"pre"`
-	Post map[uint64]map[int]TxPostValues `json:"post"`
+	Post map[uint64]map[int]types.TxPostValues `json:"post"`
 }
 
-type AcctPostValues struct {
-	Nonce     uint64                      `json:"nonce"`
-	Balance   *uint256.Int                `json:"balance"`
-	Code      []byte                      `json:"code"`
-	StorageKV map[common.Hash]common.Hash `json:"storageKV"`
-	Destruct  bool                        `json:"destruct"`
-}
-
-// For acccount destruct or storage clearing corresponding values would be 0
-type TxPostValues map[common.Address]*AcctPostValues
-
 var (
 	AllBlockBal         = BALs{}
 	AllBlockAccessLists = map[uint64]types.AccessList{}
 	// Blocknumber => TxIndex => TxPostValues
-	AllBlockTxPostValues = map[uint64]map[int]TxPostValues{}
+	AllBlockTxPostValues = map[uint64]map[int]types.TxPostValues{}
 )
 
 const (
@@ -282,6 +272,46 @@ func NewWithReader(root common.Hash, db Database, reader Reader) (*StateDB, erro
 	return sdb, nil
 }
 
+func (s *StateDB) PreComputePostState(blockNumber uint64) {
+	s.blockNumber = blockNumber
+	postBal := AllBlockTxPostValues[blockNumber]
+
+	for i := 0; i < len(postBal)-1; i++ {
+		prevVals := postBal[i]
+		postVals := postBal[i+1]
+		for addr, prevAcct := range prevVals {
+			if postAcct, ok := postVals[addr]; ok {
+				for k, v := range prevAcct.StorageKV {
+					if _, exists := postAcct.StorageKV[k]; !exists {
+						postAcct.StorageKV[k] = v
+					}
+				}
+				if postAcct.Balance == nil {
+					postAcct.Balance = prevAcct.Balance
+				}
+				if postAcct.Code == nil {
+					postAcct.Code = prevAcct.Code
+				}
+				// Storage changes but nonce didn't change
+				if postAcct.Nonce == 0 {
+					postAcct.Nonce = prevAcct.Nonce
+				}
+			} else {
+				postVals[addr] = prevAcct
+			}
+		}
+	}
+	s.postBAL = postBal
+}
+
+func (s *StateDB) SetTxBALReader() error {
+	if s.postBAL == nil {
+		return errors.New("cannot set bal reader without postBAL")
+	}
+	s.reader = newReaderWithBAL(s.reader, s.txIndex, s.postBAL)
+	return nil
+}
+
 func (s *StateDB) PrefetchStateBAL(blockNumber uint64) {
 	s.blockNumber = blockNumber
 	switch balType {
@@ -1022,7 +1052,9 @@ func (s *StateDB) Copy() *StateDB {
 		transientStorage: s.transientStorage.Copy(),
 		journal:          s.journal.copy(),
 		// The update journal is not copies to avoid duplicated updates in later transactions.
+		blockNumber:   s.blockNumber,
 		updateJournal: newJournal(),
+		postBAL:       s.postBAL,
 	}
 	if s.witness != nil {
 		state.witness = s.witness.Copy()

core/types/state_bal.go

@@ -0,0 +1,41 @@
+package types
+
+import (
+	"maps"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/holiman/uint256"
+)
+
+type AcctPostValues struct {
+	Nonce     uint64                      `json:"nonce"`
+	Balance   *uint256.Int                `json:"balance"`
+	Code      []byte                      `json:"code"`
+	StorageKV map[common.Hash]common.Hash `json:"storageKV"`
+	Destruct  bool                        `json:"destruct"`
+	CodeHash  []byte                      `json:"codeHash"`
+	Root      common.Hash                 `json:"root"`
+	Cached    bool
+}
+
+// For acccount destruct or storage clearing corresponding values would be 0
+type TxPostValues map[common.Address]*AcctPostValues
+
+// Balance and StorageKV must be cloned to avoid changing by setting postVals
+func (v *AcctPostValues) Clone() *AcctPostValues {
+	balance := v.Balance
+	if balance != nil {
+		balance = v.Balance.Clone()
+	}
+
+	return &AcctPostValues{
+		Nonce:     v.Nonce,
+		Balance:   balance,
+		Code:      v.Code,
+		StorageKV: maps.Clone(v.StorageKV),
+		Destruct:  false,
+		CodeHash:  v.CodeHash,
+		Root:      v.Root,
+		Cached: false,
+	}
+}