core/state: fix balance hooks

ethereum · Oct 21, 2024 · 350d4bb · 350d4bb
1 parent 8c7526c
commit 350d4bb
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 14 deletions.
diff --git a/core/state/state_object.go b/core/state/state_object.go
@@ -459,7 +459,7 @@ func (s *stateObject) AddBalance(amount *uint256.Int) uint256.Int {
 	return s.SetBalance(new(uint256.Int).Add(s.Balance(), amount))
 }
 
-// SetBalance sets the balance for the object, and returns the prevous balance
+// SetBalance sets the balance for the object, and returns the previous balance
 func (s *stateObject) SetBalance(amount *uint256.Int) uint256.Int {
 	prev := *s.data.Balance
 	s.db.journal.balanceChange(s.address, s.data.Balance)

diff --git a/core/state/statedb.go b/core/state/statedb.go
@@ -402,31 +402,30 @@ func (s *StateDB) HasSelfDestructed(addr common.Address) bool {
 // AddBalance adds amount to the account associated with addr.
 func (s *StateDB) AddBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) uint256.Int {
 	stateObject := s.getOrNewStateObject(addr)
-	if stateObject != nil {
-		return stateObject.AddBalance(amount)
+	if stateObject == nil {
+		return uint256.Int{}
 	}
-	return uint256.Int{}
+	return stateObject.AddBalance(amount)
 }
 
 // SubBalance subtracts amount from the account associated with addr.
 func (s *StateDB) SubBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) uint256.Int {
 	stateObject := s.getOrNewStateObject(addr)
-	var prev uint256.Int
-	if amount.IsZero() {
-		return prev
+	if stateObject == nil {
+		return uint256.Int{}
 	}
-	if stateObject != nil {
-		prev = *(stateObject.Balance())
-		stateObject.SetBalance(new(uint256.Int).Sub(stateObject.Balance(), amount))
+	if amount.IsZero() {
+		return *(stateObject.Balance())
 	}
-	return prev
+	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)
-	if stateObject != nil {
-		stateObject.SetBalance(amount)
+	if stateObject == nil {
+		return uint256.Int{}
 	}
+	return stateObject.SetBalance(amount)
 }
 
 func (s *StateDB) SetNonce(addr common.Address, nonce uint64) {

diff --git a/core/state/statedb_hooked.go b/core/state/statedb_hooked.go
@@ -60,6 +60,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.StateDB.SetBalance(addr, amount, reason)
+	if s.hooks.OnBalanceChange != nil && amount.Cmp(&prev) != 0 {
+		s.hooks.OnBalanceChange(addr, prev.ToBig(), amount.ToBig(), reason)
+	}
+	return prev
+}
+
 func (s *hookedStateDB) SetNonce(address common.Address, nonce uint64) {
 	s.StateDB.SetNonce(address, nonce)
 	if s.hooks.OnNonceChange != nil {