Merge pull request ethereum#44 from OffchainLabs/track-balance-delta

Track total balance delta in StateDB

core/state/statedb.go

@@ -38,6 +38,8 @@ import (
 type revision struct {
  id int
  journalIndex int
+ // Arbitrum: track the total balance change across all accounts
+ totalBalanceDelta *big.Int
 }
 
 var (
@@ -62,6 +64,9 @@ func (n *proofList) Delete(key []byte) error {
 // * Contracts
 // * Accounts
 type StateDB struct {
+ // Arbitrum: track the total balance change across all accounts
+ totalBalanceDelta *big.Int
+
  db Database
  prefetcher *triePrefetcher
  originalRoot common.Hash // The pre-state root, before any changes were made
@@ -131,6 +136,7 @@ func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error)
  return nil, err
  }
  sdb := &StateDB{
+ totalBalanceDelta: new(big.Int),
  db: db,
  trie: tr,
  originalRoot: root,
@@ -378,6 +384,7 @@ func (s *StateDB) HasSuicided(addr common.Address) bool {
 func (s *StateDB) AddBalance(addr common.Address, amount *big.Int) {
  stateObject := s.GetOrNewStateObject(addr)
  if stateObject != nil {
+ s.totalBalanceDelta.Add(s.totalBalanceDelta, amount)
  stateObject.AddBalance(amount)
  }
 }
@@ -386,13 +393,18 @@ func (s *StateDB) AddBalance(addr common.Address, amount *big.Int) {
 func (s *StateDB) SubBalance(addr common.Address, amount *big.Int) {
  stateObject := s.GetOrNewStateObject(addr)
  if stateObject != nil {
+ s.totalBalanceDelta.Sub(s.totalBalanceDelta, amount)
  stateObject.SubBalance(amount)
  }
 }
 
 func (s *StateDB) SetBalance(addr common.Address, amount *big.Int) {
  stateObject := s.GetOrNewStateObject(addr)
  if stateObject != nil {
+ prevBalance := stateObject.Balance()
+ s.totalBalanceDelta.Add(s.totalBalanceDelta, amount)
+ s.totalBalanceDelta.Sub(s.totalBalanceDelta, prevBalance)
+
  stateObject.SetBalance(amount)
  }
 }
@@ -443,6 +455,7 @@ func (s *StateDB) Suicide(addr common.Address) bool {
  prevbalance: new(big.Int).Set(stateObject.Balance()),
  })
  stateObject.markSuicided()
+ s.totalBalanceDelta.Sub(s.totalBalanceDelta, stateObject.data.Balance)
  stateObject.data.Balance = new(big.Int)
 
  return true
@@ -647,6 +660,7 @@ func (db *StateDB) ForEachStorage(addr common.Address, cb func(key, value common
 func (s *StateDB) Copy() *StateDB {
  // Copy all the basic fields, initialize the memory ones
  state := &StateDB{
+ totalBalanceDelta: new(big.Int).Set(s.totalBalanceDelta),
  db: s.db,
  trie: s.db.CopyTrie(s.trie),
  stateObjects: make(map[common.Address]*stateObject, len(s.journal.dirties)),
@@ -746,7 +760,7 @@ func (s *StateDB) Copy() *StateDB {
 func (s *StateDB) Snapshot() int {
  id := s.nextRevisionId
  s.nextRevisionId++
- s.validRevisions = append(s.validRevisions, revision{id, s.journal.length()})
+ s.validRevisions = append(s.validRevisions, revision{id, s.journal.length(), new(big.Int).Set(s.totalBalanceDelta)})
  return id
 }
 
@@ -759,7 +773,9 @@ func (s *StateDB) RevertToSnapshot(revid int) {
  if idx == len(s.validRevisions) || s.validRevisions[idx].id != revid {
  panic(fmt.Errorf("revision id %v cannot be reverted", revid))
  }
- snapshot := s.validRevisions[idx].journalIndex
+ revision := s.validRevisions[idx]
+ snapshot := revision.journalIndex
+ s.totalBalanceDelta = new(big.Int).Set(revision.totalBalanceDelta)
 
  // Replay the journal to undo changes and remove invalidated snapshots
  s.journal.revert(s, snapshot)
@@ -771,6 +787,11 @@ func (s *StateDB) GetRefund() uint64 {
  return s.refund
 }
 
+// GetTotalBalanceDelta returns the total change in balances since the last commit to the database
+func (s *StateDB) GetTotalBalanceDelta() *big.Int {
+ return new(big.Int).Set(s.totalBalanceDelta)
+}
+
 // Finalise finalises the state by removing the s destructed objects and clears
 // the journal as well as the refunds. Finalise, however, will not push any updates
 // into the tries just yet. Only IntermediateRoot or Commit will do that.
@@ -982,6 +1003,9 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
  }
  s.snap, s.snapDestructs, s.snapAccounts, s.snapStorage = nil, nil, nil, nil
  }
+ if err == nil {
+ s.totalBalanceDelta.Set(new(big.Int))
+ }
  return root, err
 }
 

core/state_transition.go

@@ -374,7 +374,9 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
  }
 
  res, err := st.transitionDbImpl()
+ transitionSuccess := true
  if err != nil && !errors.Is(err, ErrNonceTooLow) && !errors.Is(err, ErrNonceTooHigh) && st.msg.UnderlyingTransaction() != nil {
+ transitionSuccess = false
  res = &ExecutionResult{
  UsedGas: st.gasUsed(),
  Err: err,
@@ -384,7 +386,7 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
  }
 
  if err == nil {
- st.evm.ProcessingHook.EndTxHook(st.gas, res.Err == nil)
+ st.evm.ProcessingHook.EndTxHook(st.gas, transitionSuccess, res.Err == nil)
  }
  return res, err
 }

core/vm/evm_arbitrum.go

@@ -28,7 +28,7 @@ func (evm *EVM) Depth() int {
 type TxProcessingHook interface {
  StartTxHook() (bool, uint64, error, []byte) // return 4-tuple rather than *struct to avoid an import cycle
  GasChargingHook(gasRemaining *uint64) error
- EndTxHook(totalGasUsed uint64, success bool)
+ EndTxHook(totalGasUsed uint64, transitionSuccess bool, evmSuccess bool)
  NonrefundableGas() uint64
  PushCaller(addr common.Address)
  PopCaller()
@@ -46,7 +46,7 @@ func (p DefaultTxProcessor) GasChargingHook(gasRemaining *uint64) error {
  return nil
 }
 
-func (p DefaultTxProcessor) EndTxHook(totalGasUsed uint64, success bool) {
+func (p DefaultTxProcessor) EndTxHook(totalGasUsed uint64, transitionSuccess bool, evmSuccess bool) {
  return
 }