feat: tracing hooks added to stateDB

core/state/cached_reader.go

@@ -53,6 +53,11 @@ func (cr *CachedReader) ReadAccountData(address common.Address) (*accounts.Accou
 	return a, nil
 }
 
+// ReadAccountDataForDebug is called when an account needs to be fetched from the state
+func (cr *CachedReader) ReadAccountDataForDebug(address common.Address) (*accounts.Account, error) {
+	return cr.ReadAccountData(address)
+}
+
 // ReadAccountStorage is called when a storage item needs to be fetched from the state
 func (cr *CachedReader) ReadAccountStorage(address common.Address, incarnation uint64, key *common.Hash) ([]byte, error) {
 	addrBytes := address.Bytes()

core/state/cached_reader3.go

@@ -54,6 +54,23 @@ func (r *CachedReader3) ReadAccountData(address common.Address) (*accounts.Accou
 	return &a, nil
 }
 
+// ReadAccountDataForDebug - is like ReadAccountData, but without adding key to `readList`.
+// Used to get `prev` account balance
+func (r *CachedReader3) ReadAccountDataForDebug(address common.Address) (*accounts.Account, error) {
+	enc, err := r.cache.Get(address[:])
+	if err != nil {
+		return nil, err
+	}
+	if len(enc) == 0 {
+		return nil, nil
+	}
+	a := accounts.Account{}
+	if err = accounts.DeserialiseV3(&a, enc); err != nil {
+		return nil, err
+	}
+	return &a, nil
+}
+
 func (r *CachedReader3) ReadAccountStorage(address common.Address, incarnation uint64, key *common.Hash) ([]byte, error) {
 	compositeKey := append(address[:], key.Bytes()...)
 	enc, err := r.cache.Get(compositeKey)

core/state/database.go

@@ -37,6 +37,7 @@ const (
 
 type StateReader interface {
 	ReadAccountData(address common.Address) (*accounts.Account, error)
+	ReadAccountDataForDebug(address common.Address) (*accounts.Account, error)
 	ReadAccountStorage(address common.Address, incarnation uint64, key *common.Hash) ([]byte, error)
 	ReadAccountCode(address common.Address, incarnation uint64, codeHash common.Hash) ([]byte, error)
 	ReadAccountCodeSize(address common.Address, incarnation uint64, codeHash common.Hash) (int, error)

core/state/history_reader_v3.go

@@ -72,6 +72,12 @@ func (hr *HistoryReaderV3) ReadAccountData(address common.Address) (*accounts.Ac
 	return &a, nil
 }
 
+// ReadAccountDataForDebug - is like ReadAccountData, but without adding key to `readList`.
+// Used to get `prev` account balance
+func (hr *HistoryReaderV3) ReadAccountDataForDebug(address common.Address) (*accounts.Account, error) {
+	return hr.ReadAccountData(address)
+}
+
 func (hr *HistoryReaderV3) ReadAccountStorage(address common.Address, incarnation uint64, key *common.Hash) ([]byte, error) {
 	k := append(address[:], key.Bytes()...)
 	enc, _, err := hr.ttx.GetAsOf(kv.StorageDomain, k, nil, hr.txNum)

core/state/intra_block_state.go

@@ -96,6 +96,7 @@ type IntraBlockState struct {
 	validRevisions []revision
 	nextRevisionID int
 	trace          bool
+	tracingHooks   *tracing.Hooks
 	balanceInc     map[libcommon.Address]*BalanceIncrease // Map of balance increases (without first reading the account)
 }
 
@@ -116,6 +117,10 @@ func New(stateReader StateReader) *IntraBlockState {
 	}
 }
 
+func (sdb *IntraBlockState) SetHooks(hooks *tracing.Hooks) {
+	sdb.tracingHooks = hooks
+}
+
 func (sdb *IntraBlockState) SetTrace(trace bool) {
 	sdb.trace = trace
 }
@@ -167,6 +172,9 @@ func (sdb *IntraBlockState) AddLog(log2 *types.Log) {
 	sdb.journal.append(addLogChange{txIndex: sdb.txIndex})
 	log2.TxIndex = uint(sdb.txIndex)
 	log2.Index = sdb.logSize
+	if sdb.tracingHooks != nil && sdb.tracingHooks.OnLog != nil {
+		sdb.tracingHooks.OnLog(log2)
+	}
 	sdb.logSize++
 	for len(sdb.logs) <= sdb.txIndex {
 		sdb.logs = append(sdb.logs, nil)
@@ -383,6 +391,7 @@ func (sdb *IntraBlockState) AddBalance(addr libcommon.Address, amount *uint256.I
 	if sdb.trace {
 		fmt.Printf("AddBalance %x, %d\n", addr, amount)
 	}
+
 	// If this account has not been read, add to the balance increment map
 	_, needAccount := sdb.stateObjects[addr]
 	if !needAccount && addr == ripemd && amount.IsZero() {
@@ -393,11 +402,26 @@ func (sdb *IntraBlockState) AddBalance(addr libcommon.Address, amount *uint256.I
 			account:  &addr,
 			increase: *amount,
 		})
+
 		bi, ok := sdb.balanceInc[addr]
 		if !ok {
 			bi = &BalanceIncrease{}
 			sdb.balanceInc[addr] = bi
 		}
+
+		if sdb.tracingHooks != nil && sdb.tracingHooks.OnBalanceChange != nil {
+			// TODO: discuss if we should ignore error
+			prev := new(uint256.Int)
+			account, _ := sdb.stateReader.ReadAccountDataForDebug(addr)
+			if account != nil {
+				prev.Add(&account.Balance, &bi.increase)
+			} else {
+				prev.Add(prev, &bi.increase)
+			}
+
+			sdb.tracingHooks.OnBalanceChange(addr, prev, new(uint256.Int).Add(prev, amount), reason)
+		}
+
 		bi.increase.Add(&bi.increase, amount)
 		bi.count++
 		return
@@ -488,11 +512,18 @@ func (sdb *IntraBlockState) Selfdestruct(addr libcommon.Address) bool {
 	if stateObject == nil || stateObject.deleted {
 		return false
 	}
+
+	prevBalance := *stateObject.Balance()
 	sdb.journal.append(selfdestructChange{
 		account:     &addr,
 		prev:        stateObject.selfdestructed,
-		prevbalance: *stateObject.Balance(),
+		prevbalance: prevBalance,
 	})
+
+	if sdb.tracingHooks != nil && sdb.tracingHooks.OnBalanceChange != nil && !prevBalance.IsZero() {
+		sdb.tracingHooks.OnBalanceChange(addr, &prevBalance, uint256.NewInt(0), tracing.BalanceDecreaseSelfdestruct)
+	}
+
 	stateObject.markSelfdestructed()
 	stateObject.createdContract = false
 	stateObject.data.Balance.Clear()
@@ -684,9 +715,12 @@ func (sdb *IntraBlockState) GetRefund() uint64 {
 	return sdb.refund
 }
 
-func updateAccount(EIP161Enabled bool, isAura bool, stateWriter StateWriter, addr libcommon.Address, stateObject *stateObject, isDirty bool) error {
+func updateAccount(EIP161Enabled bool, isAura bool, stateWriter StateWriter, addr libcommon.Address, stateObject *stateObject, isDirty bool, tracingHooks *tracing.Hooks) error {
 	emptyRemoval := EIP161Enabled && stateObject.empty() && (!isAura || addr != SystemAddress)
 	if stateObject.selfdestructed || (isDirty && emptyRemoval) {
+		if tracingHooks != nil && tracingHooks.OnBalanceChange != nil && !stateObject.Balance().IsZero() && stateObject.selfdestructed {
+			tracingHooks.OnBalanceChange(stateObject.address, stateObject.Balance(), uint256.NewInt(0), tracing.BalanceDecreaseSelfdestructBurn)
+		}
 		if err := stateWriter.DeleteAccount(addr, &stateObject.original); err != nil {
 			return err
 		}
@@ -758,7 +792,7 @@ func (sdb *IntraBlockState) FinalizeTx(chainRules *chain.Rules, stateWriter Stat
 		}
 
 		//fmt.Printf("FinalizeTx: %x, balance=%d %T\n", addr, so.data.Balance.Uint64(), stateWriter)
-		if err := updateAccount(chainRules.IsSpuriousDragon, chainRules.IsAura, stateWriter, addr, so, true); err != nil {
+		if err := updateAccount(chainRules.IsSpuriousDragon, chainRules.IsAura, stateWriter, addr, so, true, sdb.tracingHooks); err != nil {
 			return err
 		}
 		so.newlyCreated = false
@@ -814,7 +848,7 @@ func (sdb *IntraBlockState) MakeWriteSet(chainRules *chain.Rules, stateWriter St
 	}
 	for addr, stateObject := range sdb.stateObjects {
 		_, isDirty := sdb.stateObjectsDirty[addr]
-		if err := updateAccount(chainRules.IsSpuriousDragon, chainRules.IsAura, stateWriter, addr, stateObject, isDirty); err != nil {
+		if err := updateAccount(chainRules.IsSpuriousDragon, chainRules.IsAura, stateWriter, addr, stateObject, isDirty, sdb.tracingHooks); err != nil {
 			return err
 		}
 	}

core/state/intra_block_state_logger_test.go

@@ -0,0 +1,139 @@
+// Copyright 2024 The Erigon Authors
+// This file is part of Erigon.
+//
+// Erigon is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Erigon is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with Erigon. If not, see <http://www.gnu.org/licenses/>.
+
+package state
+
+import (
+	"reflect"
+	"testing"
+
+	libcommon "github.com/erigontech/erigon-lib/common"
+	"github.com/erigontech/erigon-lib/kv/rawdbv3"
+	"github.com/erigontech/erigon-lib/log/v3"
+	stateLib "github.com/erigontech/erigon-lib/state"
+
+	"github.com/erigontech/erigon/core/tracing"
+	"github.com/erigontech/erigon/core/tracing/mocks"
+
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/require"
+	gomock "go.uber.org/mock/gomock"
+)
+
+func TestStateLogger(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name    string
+		prepare func(mockTracer *mocks.Mocktracer)
+		run     func(state *IntraBlockState)
+		checker func(t *testing.T, state *IntraBlockState)
+	}{
+		{
+			name: "multiple add balance",
+			prepare: func(mockTracer *mocks.Mocktracer) {
+				mockTracer.EXPECT().BalanceChangeHook(libcommon.Address{}, uint256.NewInt(0), uint256.NewInt(2), tracing.BalanceChangeUnspecified)
+				mockTracer.EXPECT().BalanceChangeHook(libcommon.Address{}, uint256.NewInt(2), uint256.NewInt(3), tracing.BalanceChangeUnspecified)
+			},
+			run: func(state *IntraBlockState) {
+				state.AddBalance(libcommon.Address{}, uint256.NewInt(2), tracing.BalanceChangeUnspecified)
+				state.AddBalance(libcommon.Address{}, uint256.NewInt(1), tracing.BalanceChangeUnspecified)
+			},
+			checker: func(t *testing.T, stateDB *IntraBlockState) {
+				bi, ok := stateDB.balanceInc[libcommon.Address{}]
+				if !ok {
+					t.Errorf("%s isn't present in balanceInc", libcommon.Address{})
+				}
+
+				if !reflect.DeepEqual(&bi.increase, uint256.NewInt(3)) {
+					t.Errorf("Incorrect BalanceInc for  %s expectedBalance: %s, got:%s", libcommon.Address{}, uint256.NewInt(3), &bi.increase)
+				}
+
+				if bi.count != 2 {
+					t.Errorf("Incorrect BalanceInc count for %s expected: %d, got:%d", libcommon.Address{}, 2, bi.count)
+				}
+
+				if len(stateDB.journal.entries) != 2 {
+					t.Errorf("Incorrect number of jounal entries expectedBalance: %d, got:%d", 2, len(stateDB.journal.entries))
+				}
+				for i := range stateDB.journal.entries {
+					switch balanceInc := stateDB.journal.entries[i].(type) {
+					case balanceIncrease:
+						var expectedInc *uint256.Int
+						if i == 0 {
+							expectedInc = uint256.NewInt(2)
+						} else {
+							expectedInc = uint256.NewInt(1)
+						}
+						if !reflect.DeepEqual(&balanceInc.increase, expectedInc) {
+							t.Errorf("Incorrect BalanceInc in jounal for  %s expectedBalance: %s, got:%s", libcommon.Address{}, expectedInc, &balanceInc.increase)
+						}
+					default:
+						t.Errorf("Invalid journal entry found:  %s", reflect.TypeOf(stateDB.journal.entries[i]))
+					}
+				}
+
+				so := stateDB.GetOrNewStateObject(libcommon.Address{})
+				if !reflect.DeepEqual(so.Balance(), uint256.NewInt(3)) {
+					t.Errorf("Incorrect Balance for  %s expectedBalance: %s, got:%s", libcommon.Address{}, uint256.NewInt(3), so.Balance())
+				}
+			},
+		},
+		{
+			name: "sub balance",
+			prepare: func(mockTracer *mocks.Mocktracer) {
+				mockTracer.EXPECT().BalanceChangeHook(libcommon.Address{}, uint256.NewInt(0), uint256.NewInt(2), tracing.BalanceChangeUnspecified)
+				mockTracer.EXPECT().BalanceChangeHook(libcommon.Address{}, uint256.NewInt(2), uint256.NewInt(1), tracing.BalanceChangeUnspecified)
+			},
+			run: func(state *IntraBlockState) {
+				state.AddBalance(libcommon.Address{}, uint256.NewInt(2), tracing.BalanceChangeUnspecified)
+				state.SubBalance(libcommon.Address{}, uint256.NewInt(1), tracing.BalanceChangeUnspecified)
+			},
+			checker: func(t *testing.T, stateDB *IntraBlockState) {
+				so := stateDB.GetOrNewStateObject(libcommon.Address{})
+				if !reflect.DeepEqual(so.Balance(), uint256.NewInt(1)) {
+					t.Errorf("Incorrect Balance for  %s expectedBalance: %s, got:%s", libcommon.Address{}, uint256.NewInt(1), so.Balance())
+				}
+			},
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			_, tx, _ := NewTestTemporalDb(t)
+
+			domains, err := stateLib.NewSharedDomains(tx, log.New())
+			require.NoError(t, err)
+			defer domains.Close()
+
+			domains.SetTxNum(1)
+			domains.SetBlockNum(1)
+			err = rawdbv3.TxNums.Append(tx, 1, 1)
+			require.NoError(t, err)
+
+			mockCtl := gomock.NewController(t)
+			defer mockCtl.Finish()
+			mockTracer := mocks.NewMocktracer(mockCtl)
+
+			state := New(NewReaderV3(domains))
+			state.SetHooks(mockTracer.Hooks())
+
+			tt.prepare(mockTracer)
+			tt.run(state)
+			tt.checker(t, state)
+		})
+	}
+}

core/state/rw_v3.go

@@ -611,6 +611,10 @@ func (r *ReaderV3) ReadAccountData(address common.Address) (*accounts.Account, e
 	return &acc, nil
 }
 
+func (r *ReaderV3) ReadAccountDataForDebug(address common.Address) (*accounts.Account, error) {
+	return r.ReadAccountData(address)
+}
+
 func (r *ReaderV3) ReadAccountStorage(address common.Address, incarnation uint64, key *common.Hash) ([]byte, error) {
 	r.composite = append(append(r.composite[:0], address[:]...), key.Bytes()...)
 	enc, _, err := r.tx.GetLatest(kv.StorageDomain, r.composite, nil)
@@ -709,6 +713,30 @@ func (r *ReaderParallelV3) ReadAccountData(address common.Address) (*accounts.Ac
 	return &acc, nil
 }
 
+// ReadAccountDataForDebug - is like ReadAccountData, but without adding key to `readList`.
+// Used to get `prev` account balance
+func (r *ReaderParallelV3) ReadAccountDataForDebug(address common.Address) (*accounts.Account, error) {
+	enc, _, err := r.sd.GetLatest(kv.AccountsDomain, address[:], nil)
+	if err != nil {
+		return nil, err
+	}
+	if len(enc) == 0 {
+		if r.trace {
+			fmt.Printf("ReadAccountData [%x] => [empty], txNum: %d\n", address, r.txNum)
+		}
+		return nil, nil
+	}
+
+	var acc accounts.Account
+	if err := accounts.DeserialiseV3(&acc, enc); err != nil {
+		return nil, err
+	}
+	if r.trace {
+		fmt.Printf("ReadAccountData [%x] => [nonce: %d, balance: %d, codeHash: %x], txNum: %d\n", address, acc.Nonce, &acc.Balance, acc.CodeHash, r.txNum)
+	}
+	return &acc, nil
+}
+
 func (r *ReaderParallelV3) ReadAccountStorage(address common.Address, incarnation uint64, key *common.Hash) ([]byte, error) {
 	r.composite = append(append(r.composite[:0], address[:]...), key.Bytes()...)
 	enc, _, err := r.sd.GetLatest(kv.StorageDomain, r.composite, nil)

core/state/state_object.go

@@ -222,6 +222,9 @@ func (so *stateObject) SetState(key *libcommon.Hash, value uint256.Int) {
 		key:      *key,
 		prevalue: prev,
 	})
+	if so.db.tracingHooks != nil && so.db.tracingHooks.OnStorageChange != nil {
+		so.db.tracingHooks.OnStorageChange(so.address, key, prev, value)
+	}
 	so.setState(key, value)
 }
 
@@ -295,6 +298,9 @@ func (so *stateObject) SetBalance(amount *uint256.Int, reason tracing.BalanceCha
 		account: &so.address,
 		prev:    so.data.Balance,
 	})
+	if so.db.tracingHooks != nil && so.db.tracingHooks.OnBalanceChange != nil {
+		so.db.tracingHooks.OnBalanceChange(so.address, so.Balance(), amount, reason)
+	}
 	so.setBalance(amount)
 }
 
@@ -342,6 +348,9 @@ func (so *stateObject) SetCode(codeHash libcommon.Hash, code []byte) {
 		prevhash: so.data.CodeHash,
 		prevcode: prevcode,
 	})
+	if so.db.tracingHooks != nil && so.db.tracingHooks.OnCodeChange != nil {
+		so.db.tracingHooks.OnCodeChange(so.address, so.data.CodeHash, prevcode, codeHash, code)
+	}
 	so.setCode(codeHash, code)
 }
 
@@ -356,6 +365,9 @@ func (so *stateObject) SetNonce(nonce uint64) {
 		account: &so.address,
 		prev:    so.data.Nonce,
 	})
+	if so.db.tracingHooks != nil && so.db.tracingHooks.OnNonceChange != nil {
+		so.db.tracingHooks.OnNonceChange(so.address, so.data.Nonce, nonce)
+	}
 	so.setNonce(nonce)
 }