More StateDB hook

consensus/bor/bor.go

@@ -881,7 +881,7 @@ func (c *Bor) Finalize(chain consensus.ChainHeaderReader, header *types.Header,
 
 	// Extract the underlying state to access methods like `IntermediateRoot` and `Copy`
 	// required for bor consensus operations
-	state := wrappedState.(*state.StateDB)
+	// state := wrappedState.(*state.StateDB)
 
 	var (
 		stateSyncData []*types.StateSyncData
@@ -892,24 +892,24 @@ func (c *Bor) Finalize(chain consensus.ChainHeaderReader, header *types.Header,
 		start := time.Now()
 		cx := statefull.ChainContext{Chain: chain, Bor: c}
 		// check and commit span
-		if err := c.checkAndCommitSpan(state, header, cx, tracer); err != nil {
+		if err := c.checkAndCommitSpan(wrappedState, header, cx, tracer); err != nil {
 			log.Error("Error while committing span", "error", err)
 			return
 		}
 
 		if c.HeimdallClient != nil {
 			// commit states
-			stateSyncData, err = c.CommitStates(state, header, cx, tracer)
+			stateSyncData, err = c.CommitStates(wrappedState, header, cx, tracer)
 			if err != nil {
 				log.Error("Error while committing states", "error", err)
 				return
 			}
 		}
 
-		state.BorConsensusTime = time.Since(start)
+		wrappedState.SetBorConsensusTime(time.Since(start))
 	}
 
-	if err = c.changeContractCodeIfNeeded(headerNumber, state); err != nil {
+	if err = c.changeContractCodeIfNeeded(headerNumber, wrappedState); err != nil {
 		log.Error("Error changing contract code", "error", err)
 		return
 	}
@@ -934,7 +934,7 @@ func decodeGenesisAlloc(i interface{}) (types.GenesisAlloc, error) {
 	return alloc, nil
 }
 
-func (c *Bor) changeContractCodeIfNeeded(headerNumber uint64, state *state.StateDB) error {
+func (c *Bor) changeContractCodeIfNeeded(headerNumber uint64, state vm.StateDB) error {
 	for blockNumber, genesisAlloc := range c.config.BlockAlloc {
 		if blockNumber == strconv.FormatUint(headerNumber, 10) {
 			allocs, err := decodeGenesisAlloc(genesisAlloc)
@@ -977,7 +977,7 @@ func (c *Bor) FinalizeAndAssemble(chain consensus.ChainHeaderReader, header *typ
 		cx := statefull.ChainContext{Chain: chain, Bor: c}
 
 		// check and commit span
-		if err = c.checkAndCommitSpan(state, header, cx, tracer); err != nil {
+		if err = c.checkAndCommitSpan(vm.StateDB(state), header, cx, tracer); err != nil {
 			log.Error("Error while committing span", "error", err)
 			return nil, err
 		}
@@ -1168,7 +1168,7 @@ func (c *Bor) Close() error {
 }
 
 func (c *Bor) checkAndCommitSpan(
-	state *state.StateDB,
+	state vm.StateDB,
 	header *types.Header,
 	chain core.ChainContext,
 	tracer *tracing.Hooks,
@@ -1217,7 +1217,7 @@ func (c *Bor) needToCommitSpan(currentSpan *span.Span, headerNumber uint64) bool
 func (c *Bor) FetchAndCommitSpan(
 	ctx context.Context,
 	newSpanID uint64,
-	state *state.StateDB,
+	state vm.StateDB,
 	header *types.Header,
 	chain core.ChainContext,
 	tracer *tracing.Hooks,
@@ -1255,7 +1255,7 @@ func (c *Bor) FetchAndCommitSpan(
 
 // CommitStates commit states
 func (c *Bor) CommitStates(
-	state *state.StateDB,
+	state vm.StateDB,
 	header *types.Header,
 	chain statefull.ChainContext,
 	tracer *tracing.Hooks,
@@ -1272,7 +1272,7 @@ func (c *Bor) CommitStates(
 
 	if c.config.IsIndore(header.Number) {
 		// Fetch the LastStateId from contract via current state instance
-		lastStateIDBig, err = c.GenesisContractsClient.LastStateId(state.Copy(), number-1, header.ParentHash)
+		lastStateIDBig, err = c.GenesisContractsClient.LastStateId(state.Clone().(vm.StateDB), number-1, header.ParentHash)
 		if err != nil {
 			return nil, err
 		}

consensus/bor/contract/client.go

@@ -2,11 +2,12 @@ package contract
 
 import (
 	"context"
-	"github.com/ethereum/go-ethereum/core/tracing"
 	"math"
 	"math/big"
 	"strings"
 
+	"github.com/ethereum/go-ethereum/core/tracing"
+
 	"github.com/ethereum/go-ethereum/accounts/abi"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
@@ -15,6 +16,7 @@ import (
 	"github.com/ethereum/go-ethereum/consensus/bor/statefull"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/internal/ethapi"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
@@ -67,7 +69,7 @@ func NewGenesisContractsClient(
 
 func (gc *GenesisContractsClient) CommitState(
 	event *clerk.EventRecordWithTime,
-	state *state.StateDB,
+	state vm.StateDB,
 	header *types.Header,
 	chCtx statefull.ChainContext,
 	tracer *tracing.Hooks,
@@ -105,7 +107,7 @@ func (gc *GenesisContractsClient) CommitState(
 	return gasUsed, nil
 }
 
-func (gc *GenesisContractsClient) LastStateId(state *state.StateDB, number uint64, hash common.Hash) (*big.Int, error) {
+func (gc *GenesisContractsClient) LastStateId(stateDB vm.StateDB, number uint64, hash common.Hash) (*big.Int, error) {
 	blockNr := rpc.BlockNumber(number)
 
 	const method = "lastStateId"
@@ -121,13 +123,16 @@ func (gc *GenesisContractsClient) LastStateId(state *state.StateDB, number uint6
 	toAddress := common.HexToAddress(gc.StateReceiverContract)
 	gas := (hexutil.Uint64)(uint64(math.MaxUint64 / 2))
 
+	// The unhooked version always return the *state.StateDB inner object
+	original := stateDB.Unhooked().(*state.StateDB)
+
 	// BOR: Do a 'CallWithState' so that we can fetch the last state ID from a given (incoming)
 	// state instead of local(canonical) chain's state.
 	result, err := gc.ethAPI.CallWithState(context.Background(), ethapi.TransactionArgs{
 		Gas:  &gas,
 		To:   &toAddress,
 		Data: &msgData,
-	}, &rpc.BlockNumberOrHash{BlockNumber: &blockNr, BlockHash: &hash}, state, nil, nil)
+	}, &rpc.BlockNumberOrHash{BlockNumber: &blockNr, BlockHash: &hash}, original, nil, nil)
 	if err != nil {
 		return nil, err
 	}

consensus/bor/genesis.go

@@ -1,18 +1,19 @@
 package bor
 
 import (
-	"github.com/ethereum/go-ethereum/core/tracing"
 	"math/big"
 
+	"github.com/ethereum/go-ethereum/core/tracing"
+	"github.com/ethereum/go-ethereum/core/vm"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/consensus/bor/clerk"
 	"github.com/ethereum/go-ethereum/consensus/bor/statefull"
-	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 )
 
 //go:generate mockgen -destination=./genesis_contract_mock.go -package=bor . GenesisContract
 type GenesisContract interface {
-	CommitState(event *clerk.EventRecordWithTime, state *state.StateDB, header *types.Header, chCtx statefull.ChainContext, tracer *tracing.Hooks) (uint64, error)
+	CommitState(event *clerk.EventRecordWithTime, state vm.StateDB, header *types.Header, chCtx statefull.ChainContext, tracer *tracing.Hooks) (uint64, error)
-	LastStateId(state *state.StateDB, number uint64, hash common.Hash) (*big.Int, error)
+	LastStateId(state vm.StateDB, number uint64, hash common.Hash) (*big.Int, error)
 }

consensus/bor/genesis_contract_mock.go

@@ -11,9 +11,9 @@ import (
 	common "github.com/ethereum/go-ethereum/common"
 	clerk "github.com/ethereum/go-ethereum/consensus/bor/clerk"
 	statefull "github.com/ethereum/go-ethereum/consensus/bor/statefull"
-	state "github.com/ethereum/go-ethereum/core/state"
-	types "github.com/ethereum/go-ethereum/core/types"
 	tracing "github.com/ethereum/go-ethereum/core/tracing"
+	types "github.com/ethereum/go-ethereum/core/types"
+	vm "github.com/ethereum/go-ethereum/core/vm"
 	gomock "github.com/golang/mock/gomock"
 )
 
@@ -41,7 +41,7 @@ func (m *MockGenesisContract) EXPECT() *MockGenesisContractMockRecorder {
 }
 
 // CommitState mocks base method.
-func (m *MockGenesisContract) CommitState(arg0 *clerk.EventRecordWithTime, arg1 *state.StateDB, arg2 *types.Header, arg3 statefull.ChainContext, arg4 *tracing.Hooks) (uint64, error) {
+func (m *MockGenesisContract) CommitState(arg0 *clerk.EventRecordWithTime, arg1 vm.StateDB, arg2 *types.Header, arg3 statefull.ChainContext, arg4 *tracing.Hooks) (uint64, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "CommitState", arg0, arg1, arg2, arg3, arg4)
 	ret0, _ := ret[0].(uint64)
@@ -56,7 +56,7 @@ func (mr *MockGenesisContractMockRecorder) CommitState(arg0, arg1, arg2, arg3, a
 }
 
 // LastStateId mocks base method.
-func (m *MockGenesisContract) LastStateId(arg0 *state.StateDB, arg1 uint64, arg2 common.Hash) (*big.Int, error) {
+func (m *MockGenesisContract) LastStateId(arg0 vm.StateDB, arg1 uint64, arg2 common.Hash) (*big.Int, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "LastStateId", arg0, arg1, arg2)
 	ret0, _ := ret[0].(*big.Int)

consensus/bor/heimdall/span/spanner.go

@@ -13,9 +13,9 @@ import (
 	"github.com/ethereum/go-ethereum/consensus/bor/statefull"
 	"github.com/ethereum/go-ethereum/consensus/bor/valset"
 	"github.com/ethereum/go-ethereum/core"
-	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/tracing"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/internal/ethapi"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
@@ -288,7 +288,7 @@ func (c *ChainSpanner) GetCurrentValidatorsByHash(ctx context.Context, headerHas
 
 const method = "commitSpan"
 
-func (c *ChainSpanner) CommitSpan(ctx context.Context, heimdallSpan HeimdallSpan, state *state.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error {
+func (c *ChainSpanner) CommitSpan(ctx context.Context, heimdallSpan HeimdallSpan, state vm.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error {
 	// get validators bytes
 	validators := make([]valset.MinimalVal, 0, len(heimdallSpan.ValidatorSet.Validators))
 	for _, val := range heimdallSpan.ValidatorSet.Validators {

consensus/bor/span.go

@@ -2,13 +2,14 @@ package bor
 
 import (
 	"context"
+
 	"github.com/ethereum/go-ethereum/core/tracing"
+	"github.com/ethereum/go-ethereum/core/vm"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/consensus/bor/heimdall/span"
 	"github.com/ethereum/go-ethereum/consensus/bor/valset"
 	"github.com/ethereum/go-ethereum/core"
-	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/rpc"
 )
@@ -18,5 +19,5 @@ type Spanner interface {
 	GetCurrentSpan(ctx context.Context, headerHash common.Hash) (*span.Span, error)
 	GetCurrentValidatorsByHash(ctx context.Context, headerHash common.Hash, blockNumber uint64) ([]*valset.Validator, error)
 	GetCurrentValidatorsByBlockNrOrHash(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash, blockNumber uint64) ([]*valset.Validator, error)
-	CommitSpan(ctx context.Context, heimdallSpan span.HeimdallSpan, state *state.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error
+	CommitSpan(ctx context.Context, heimdallSpan span.HeimdallSpan, state vm.StateDB, header *types.Header, chainContext core.ChainContext, tracer *tracing.Hooks) error
 }

consensus/bor/span_mock.go

@@ -12,9 +12,9 @@ import (
 	span "github.com/ethereum/go-ethereum/consensus/bor/heimdall/span"
 	valset "github.com/ethereum/go-ethereum/consensus/bor/valset"
 	core "github.com/ethereum/go-ethereum/core"
-	state "github.com/ethereum/go-ethereum/core/state"
 	tracing "github.com/ethereum/go-ethereum/core/tracing"
 	types "github.com/ethereum/go-ethereum/core/types"
+	vm "github.com/ethereum/go-ethereum/core/vm"
 	rpc "github.com/ethereum/go-ethereum/rpc"
 	gomock "github.com/golang/mock/gomock"
 )
@@ -43,7 +43,7 @@ func (m *MockSpanner) EXPECT() *MockSpannerMockRecorder {
 }
 
 // CommitSpan mocks base method.
-func (m *MockSpanner) CommitSpan(arg0 context.Context, arg1 span.HeimdallSpan, arg2 *state.StateDB, arg3 *types.Header, arg4 core.ChainContext, arg5 *tracing.Hooks) error {
+func (m *MockSpanner) CommitSpan(arg0 context.Context, arg1 span.HeimdallSpan, arg2 vm.StateDB, arg3 *types.Header, arg4 core.ChainContext, arg5 *tracing.Hooks) error {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "CommitSpan", arg0, arg1, arg2, arg3, arg4, arg5)
 	ret0, _ := ret[0].(error)

consensus/bor/statefull/processor.go

@@ -3,16 +3,16 @@ package statefull
 import (
 	"bytes"
 	"context"
-	"github.com/ethereum/go-ethereum/core/tracing"
-	"github.com/ethereum/go-ethereum/crypto"
 	"math"
 	"math/big"
 
+	"github.com/ethereum/go-ethereum/core/tracing"
+	"github.com/ethereum/go-ethereum/crypto"
+
 	ethereum "github.com/ethereum/go-ethereum"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/consensus"
 	"github.com/ethereum/go-ethereum/core"
-	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/log"
@@ -71,7 +71,7 @@ func GetSystemMessage(toAddress common.Address, data []byte) Callmsg {
 func ApplyMessage(
 	_ context.Context,
 	msg Callmsg,
-	state *state.StateDB,
+	state vm.StateDB,
 	header *types.Header,
 	chainConfig *params.ChainConfig,
 	chainContext core.ChainContext,

core/state/statedb.go

@@ -808,13 +808,16 @@ func (s *StateDB) SubBalance(addr common.Address, amount *uint256.Int, reason tr
 	return stateObject.SetBalance(new(uint256.Int).Sub(stateObject.Balance(), amount))
 }
 
-func (s *StateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) {
+func (s *StateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) uint256.Int {
 	stateObject := s.getOrNewStateObject(addr)
+	var prevBalance uint256.Int
+
 	if stateObject != nil {
 		stateObject = s.mvRecordWritten(stateObject)
-		stateObject.SetBalance(amount)
+		prevBalance = stateObject.SetBalance(amount)
 		MVWrite(s, blockstm.NewSubpathKey(addr, BalancePath))
 	}
+	return prevBalance
 }
 
 func (s *StateDB) SetNonce(addr common.Address, nonce uint64, reason tracing.NonceChangeReason) {
@@ -1942,3 +1945,18 @@ func (s *StateDB) Witness() *stateless.Witness {
 func (s *StateDB) AccessEvents() *AccessEvents {
 	return s.accessEvents
 }
+
+// Polygon specific
+
+func (s *StateDB) Clone() any {
+	return s.Copy()
+}
+
+func (s *StateDB) Unhooked() any {
+	// Already unhooked, just return self
+	return s
+}
+
+func (s *StateDB) SetBorConsensusTime(borConsensusTime time.Duration) {
+	s.BorConsensusTime = borConsensusTime
+}

core/state/statedb_hooked.go

@@ -18,6 +18,7 @@ package state
 
 import (
 	"math/big"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/blockstm"
@@ -184,6 +185,14 @@ func (s *hookedStateDB) AddBalance(addr common.Address, amount *uint256.Int, rea
 	return prev
 }
 
+func (s *hookedStateDB) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) uint256.Int {
+	prev := s.inner.SetBalance(addr, amount, reason)
+	if s.hooks.OnBalanceChange != nil {
+		s.hooks.OnBalanceChange(addr, prev.ToBig(), amount.ToBig(), reason)
+	}
+	return prev
+}
+
 func (s *hookedStateDB) SetNonce(address common.Address, nonce uint64, reason tracing.NonceChangeReason) {
 	prev := s.inner.GetNonce(address)
 	s.inner.SetNonce(address, nonce, reason)
@@ -304,3 +313,12 @@ func (s *hookedStateDB) TxIndex() int {
 func (s *hookedStateDB) SetTxContext(txHash common.Hash, txIndex int) {
 	s.inner.SetTxContext(txHash, txIndex)
 }
+func (s *hookedStateDB) Clone() any {
+	return NewHookedState(s.inner.Copy(), s.hooks)
+}
+func (s *hookedStateDB) Unhooked() any {
+	return s.inner
+}
+func (s *hookedStateDB) SetBorConsensusTime(borConsensusTime time.Duration) {
+	s.inner.SetBorConsensusTime(borConsensusTime)
+}

core/vm/interface.go

@@ -17,6 +17,8 @@
 package vm
 
 import (
+	"time"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/blockstm"
 	"github.com/ethereum/go-ethereum/core/state"
@@ -111,4 +113,20 @@ type StateDB interface {
 	GetLogs(txHash common.Hash, blockNumber uint64, blockHash common.Hash) []*types.Log
 	TxIndex() int
 	SetTxContext(txHash common.Hash, txIndex int)
+	SetBalance(common.Address, *uint256.Int, tracing.BalanceChangeReason) uint256.Int
+	// Clone is used to create a copy of the StateDB, same as `Copy` on *state.StateDB but rename so interface has its own name
+	//
+	//   state.Clone().(vm.StateDB)
+	//
+	// The `any` return type is required to avoid import cycles.
+	Clone() any
+	// Unhooked is used to return the underlying state without any hooks applied, in Polygon, some potential
+	// state modifying operations can be called on a vm.StateDB interface, which might be hooked but we want those
+	// operation to always be non-recorded, this method ensures this.
+	//
+	//   state.Unhooked().(vm.StateDB)
+	//
+	// The `any` return type is required to avoid import cycles.
+	Unhooked() any
+	SetBorConsensusTime(borConsensusTime time.Duration)
 }