github.com/lmittmann/w3@v0.20.0/w3types/state.go (about)

     1  package w3types
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"maps"
     7  	"math/big"
     8  	"sync/atomic"
     9  
    10  	"github.com/ethereum/go-ethereum/common"
    11  	"github.com/ethereum/go-ethereum/common/hexutil"
    12  	"github.com/ethereum/go-ethereum/core/types"
    13  	"github.com/lmittmann/w3/internal/crypto"
    14  )
    15  
    16  type State map[common.Address]*Account
    17  
    18  // SetGenesisAlloc copies the given [types.GenesisAlloc] to the state and
    19  // returns it.
    20  func (s State) SetGenesisAlloc(alloc types.GenesisAlloc) State {
    21  	s = make(State, len(alloc))
    22  	for addr, acc := range alloc {
    23  		s[addr] = &Account{
    24  			Nonce:   acc.Nonce,
    25  			Balance: acc.Balance,
    26  			Code:    acc.Code,
    27  			Storage: acc.Storage,
    28  		}
    29  	}
    30  	return s
    31  }
    32  
    33  // Merge returns a new state that is the result of merging the called state with the given state.
    34  // All state in other state will overwrite the state in the called state.
    35  func (s State) Merge(other State) (merged State) {
    36  	merged = make(State, len(s))
    37  
    38  	// copy all accounts from s
    39  	for addr, acc := range s {
    40  		merged[addr] = acc.deepCopy()
    41  	}
    42  
    43  	// merge all accounts from other
    44  	for addr, acc := range other {
    45  		if mergedAcc, ok := merged[addr]; ok {
    46  			mergedAcc.merge(acc)
    47  		} else {
    48  			merged[addr] = acc.deepCopy()
    49  		}
    50  	}
    51  	return merged
    52  }
    53  
    54  type Account struct {
    55  	Nonce   uint64
    56  	Balance *big.Int
    57  	Code    []byte
    58  	Storage Storage
    59  
    60  	codeHash atomic.Pointer[common.Hash] // caches the code hash
    61  }
    62  
    63  // deepCopy returns a deep copy of the account.
    64  func (acc *Account) deepCopy() *Account {
    65  	newAcc := &Account{Nonce: acc.Nonce}
    66  	if acc.Balance != nil {
    67  		newAcc.Balance = new(big.Int).Set(acc.Balance)
    68  	}
    69  	if acc.Code != nil {
    70  		newAcc.Code = bytes.Clone(acc.Code)
    71  	}
    72  	if len(acc.Storage) > 0 {
    73  		newAcc.Storage = maps.Clone(acc.Storage)
    74  	}
    75  	return newAcc
    76  }
    77  
    78  // merge merges the given account into the called account.
    79  func (dst *Account) merge(src *Account) {
    80  	// merge account fields
    81  	srcIsZero := src.Nonce == 0 && src.Balance == nil && len(src.Code) == 0
    82  	if !srcIsZero {
    83  		dst.Nonce = src.Nonce
    84  		if src.Balance != nil {
    85  			dst.Balance = new(big.Int).Set(src.Balance)
    86  		} else {
    87  			dst.Balance = nil
    88  		}
    89  		if len(src.Code) > 0 {
    90  			dst.Code = bytes.Clone(src.Code)
    91  		} else {
    92  			dst.Code = nil
    93  		}
    94  	}
    95  
    96  	// merge storage
    97  	if dst.Storage == nil && len(src.Storage) > 0 {
    98  		dst.Storage = maps.Clone(src.Storage)
    99  	} else if len(src.Storage) > 0 {
   100  		maps.Copy(dst.Storage, src.Storage)
   101  	}
   102  }
   103  
   104  // CodeHash returns the hash of the account's code.
   105  func (acc *Account) CodeHash() common.Hash {
   106  	if codeHash := acc.codeHash.Load(); codeHash != nil {
   107  		return *codeHash
   108  	}
   109  
   110  	if len(acc.Code) == 0 {
   111  		acc.codeHash.Store(&types.EmptyCodeHash)
   112  		return types.EmptyCodeHash
   113  	}
   114  
   115  	codeHash := crypto.Keccak256Hash(acc.Code)
   116  	acc.codeHash.Store(&codeHash)
   117  	return codeHash
   118  }
   119  
   120  // MarshalJSON implements the [json.Marshaler].
   121  func (acc *Account) MarshalJSON() ([]byte, error) {
   122  	type account struct {
   123  		Nonce   hexutil.Uint64 `json:"nonce,omitempty"`
   124  		Balance *hexutil.Big   `json:"balance,omitempty"`
   125  		Code    hexutil.Bytes  `json:"code,omitempty"`
   126  		Storage Storage        `json:"stateDiff,omitempty"`
   127  	}
   128  
   129  	var enc account
   130  	if acc.Nonce > 0 {
   131  		enc.Nonce = hexutil.Uint64(acc.Nonce)
   132  	}
   133  	if acc.Balance != nil {
   134  		enc.Balance = (*hexutil.Big)(acc.Balance)
   135  	}
   136  	if len(acc.Code) > 0 {
   137  		enc.Code = acc.Code
   138  	}
   139  	if len(acc.Storage) > 0 {
   140  		enc.Storage = acc.Storage
   141  	}
   142  	return json.Marshal(&enc)
   143  }
   144  
   145  type Storage map[common.Hash]common.Hash