Computing the contract size and storing it in the Account struct

core/state/dump.go

@@ -53,6 +53,7 @@ type DumpAccount struct {
 	Root      hexutil.Bytes           `json:"root"`
 	CodeHash  hexutil.Bytes           `json:"codeHash"`
 	Code      hexutil.Bytes           `json:"code,omitempty"`
+	Size      string                  `json:"size"`
 	Storage   map[common.Hash]string  `json:"storage,omitempty"`
 	Address   *common.InternalAddress `json:"address,omitempty"` // Address only present in iterative (line-by-line) mode
 	SecureKey hexutil.Bytes           `json:"key,omitempty"`     // If we don't have address, we can output the key
@@ -149,6 +150,7 @@ func (s *StateDB) DumpToCollector(c DumpCollector, conf *DumpConfig) (nextKey []
 			Nonce:     data.Nonce,
 			Root:      data.Root[:],
 			CodeHash:  data.CodeHash,
+			Size:      data.Size.String(),
 			SecureKey: it.Key,
 		}
 		addrBytes := s.trie.GetKey(it.Key)

core/state/journal.go

@@ -116,6 +116,10 @@ type (
 		account            *common.InternalAddress
 		prevcode, prevhash []byte
 	}
+	sizeChange struct {
+		account *common.InternalAddress
+		prev    *big.Int
+	}
 
 	// Changes to other state values.
 	refundChange struct {
@@ -211,6 +215,14 @@ func (ch storageChange) dirtied() *common.InternalAddress {
 	return ch.account
 }
 
+func (ch sizeChange) revert(s *StateDB) {
+	s.getStateObject(*ch.account).SetSize(ch.prev)
+}
+
+func (ch sizeChange) dirtied() *common.InternalAddress {
+	return ch.account
+}
+
 func (ch refundChange) revert(s *StateDB) {
 	s.refund = ch.prev
 }

core/state/snapshot/account.go

@@ -33,13 +33,15 @@ type Account struct {
 	Balance  *big.Int
 	Root     []byte
 	CodeHash []byte
+	Size     *big.Int
 }
 
 // SlimAccount converts a state.Account content into a slim snapshot account
-func SlimAccount(nonce uint64, balance *big.Int, root common.Hash, codehash []byte) Account {
+func SlimAccount(nonce uint64, balance *big.Int, root common.Hash, codehash []byte, size *big.Int) Account {
 	slim := Account{
 		Nonce:   nonce,
 		Balance: balance,
+		Size:    size,
 	}
 	if root != emptyRoot {
 		slim.Root = root[:]
@@ -52,8 +54,8 @@ func SlimAccount(nonce uint64, balance *big.Int, root common.Hash, codehash []by
 
 // SlimAccountRLP converts a state.Account content into a slim snapshot
 // version RLP encoded.
-func SlimAccountRLP(nonce uint64, balance *big.Int, root common.Hash, codehash []byte) []byte {
-	data, err := rlp.EncodeToBytes(SlimAccount(nonce, balance, root, codehash))
+func SlimAccountRLP(nonce uint64, balance *big.Int, root common.Hash, codehash []byte, size *big.Int) []byte {
+	data, err := rlp.EncodeToBytes(SlimAccount(nonce, balance, root, codehash, size))
 	if err != nil {
 		panic(err)
 	}

core/state/snapshot/generate.go

@@ -566,6 +566,7 @@ func (dl *diskLayer) generate(stats *generatorStats) {
 			Balance  *big.Int
 			Root     common.Hash
 			CodeHash []byte
+			Size     *big.Int
 		}
 		if err := rlp.DecodeBytes(val, &acc); err != nil {
 			stats.logger.WithField("err", err).Fatal("Invalid account encountered during snapshot creation")
@@ -581,7 +582,7 @@ func (dl *diskLayer) generate(stats *generatorStats) {
 					dataLen -= 32
 				}
 			} else {
-				data := SlimAccountRLP(acc.Nonce, acc.Balance, acc.Root, acc.CodeHash)
+				data := SlimAccountRLP(acc.Nonce, acc.Balance, acc.Root, acc.CodeHash, acc.Size)
 				dataLen = len(data)
 				rawdb.WriteAccountSnapshot(batch, accountHash, data)
 			}

core/state/state_object.go

@@ -85,6 +85,8 @@ type stateObject struct {
 	dirtyStorage   Storage // Storage entries that have been modified in the current transaction execution
 	fakeStorage    Storage // Fake storage which constructed by caller for debugging purpose.
 
+	uniqueNewKeysStorage Storage //Storage cache for the new entries being considered to be added to the storage trie
+
 	// Cache flags.
 	// When an object is marked suicided it will be delete from the trie
 	// during the "update" phase of the state transition.
@@ -95,7 +97,7 @@ type stateObject struct {
 
 // empty returns whether the account is considered empty.
 func (s *stateObject) empty() bool {
-	return s.data.Nonce == 0 && s.data.Balance.Sign() == 0 && bytes.Equal(s.data.CodeHash, emptyCodeHash)
+	return s.data.Nonce == 0 && s.data.Balance.Sign() == 0 && bytes.Equal(s.data.CodeHash, emptyCodeHash) && s.data.Size.Sign() == 0
 }
 
 // Account is the Quai consensus representation of accounts.
@@ -105,6 +107,7 @@ type Account struct {
 	Balance  *big.Int
 	Root     common.Hash // merkle root of the storage trie
 	CodeHash []byte
+	Size     *big.Int // size of the storage trie for the contract
 }
 
 // newObject creates a state object.
@@ -118,14 +121,18 @@ func newObject(db *StateDB, address common.InternalAddress, data Account) *state
 	if data.Root == (common.Hash{}) {
 		data.Root = emptyRoot
 	}
+	if data.Size == nil {
+		data.Size = new(big.Int)
+	}
 	return &stateObject{
-		db:             db,
-		address:        address,
-		addrHash:       crypto.Keccak256Hash(address[:]),
-		data:           data,
-		originStorage:  make(Storage),
-		pendingStorage: make(Storage),
-		dirtyStorage:   make(Storage),
+		db:                   db,
+		address:              address,
+		addrHash:             crypto.Keccak256Hash(address[:]),
+		data:                 data,
+		originStorage:        make(Storage),
+		pendingStorage:       make(Storage),
+		dirtyStorage:         make(Storage),
+		uniqueNewKeysStorage: make(Storage),
 	}
 }
 
@@ -220,6 +227,11 @@ func (s *stateObject) GetCommittedState(db Database, key common.Hash) common.Has
 		}()
 	}
 
+	if enc, err = s.getTrie(db).TryGet(key.Bytes()); len(enc) == 0 && err == nil {
+		// This key didnt exist in  the database, we need to add it to the uniqueKeys
+		s.uniqueNewKeysStorage[key] = common.Hash{}
+	}
+
 	if s.db.snap != nil {
 		// If the object was destructed in *this* block (and potentially resurrected),
 		// the storage has been cleared out, and we should *not* consult the previous
@@ -347,10 +359,23 @@ func (s *stateObject) updateTrie(db Database) Trie {
 		var v []byte
 		if (value == common.Hash{}) {
 			s.setError(tr.TryDelete(key[:]))
+			// While writing the value is nil check in the uniqueNewKeysStorage
+			// and if the key doesnt exist that means we are deleting key that
+			// previously existed in the storage trie lower the size of the
+			// storage trie
+			_, exists := s.uniqueNewKeysStorage[key]
+			if !exists {
+				s.SubSize()
+			}
 		} else {
 			// Encoding []byte cannot fail, ok to ignore the error.
 			v, _ = rlp.EncodeToBytes(common.TrimLeftZeroes(value[:]))
 			s.setError(tr.TryUpdate(key[:], v))
+			// While writing the values to the trie, check if the key exists in the uniqueKeys
+			_, exists := s.uniqueNewKeysStorage[key]
+			if exists {
+				s.AddSize()
+			}
 		}
 		// If state snapshotting is active, cache the data til commit
 		if s.db.snap != nil {
@@ -371,6 +396,9 @@ func (s *stateObject) updateTrie(db Database) Trie {
 	if len(s.pendingStorage) > 0 {
 		s.pendingStorage = make(Storage)
 	}
+	if len(s.uniqueNewKeysStorage) > 0 {
+		s.uniqueNewKeysStorage = make(Storage)
+	}
 	return tr
 }
 
@@ -441,6 +469,26 @@ func (s *stateObject) setBalance(amount *big.Int) {
 	s.data.Balance = amount
 }
 
+func (s *stateObject) SetSize(size *big.Int) {
+	s.db.journal.append(sizeChange{
+		account: &s.address,
+		prev:    new(big.Int).Set(s.data.Size),
+	})
+	s.setSize(size)
+}
+
+func (s *stateObject) AddSize() {
+	s.SetSize(new(big.Int).Add(s.Size(), common.Big1))
+}
+
+func (s *stateObject) SubSize() {
+	s.SetSize(new(big.Int).Sub(s.Size(), common.Big1))
+}
+
+func (s *stateObject) setSize(size *big.Int) {
+	s.data.Size = size
+}
+
 func (s *stateObject) deepCopy(db *StateDB) *stateObject {
 	stateObject := newObject(db, s.address, s.data)
 	if s.trie != nil {
@@ -450,6 +498,7 @@ func (s *stateObject) deepCopy(db *StateDB) *stateObject {
 	stateObject.dirtyStorage = s.dirtyStorage.Copy()
 	stateObject.originStorage = s.originStorage.Copy()
 	stateObject.pendingStorage = s.pendingStorage.Copy()
+	stateObject.uniqueNewKeysStorage = s.uniqueNewKeysStorage.Copy()
 	stateObject.suicided = s.suicided
 	stateObject.dirtyCode = s.dirtyCode
 	stateObject.deleted = s.deleted
@@ -538,6 +587,10 @@ func (s *stateObject) Nonce() uint64 {
 	return s.data.Nonce
 }
 
+func (s *stateObject) Size() *big.Int {
+	return s.data.Size
+}
+
 // Never called, but must be present to allow stateObject to be used
 // as a vm.Account interface that also satisfies the vm.ContractRef
 // interface. Interfaces are awesome.

core/state/statedb.go

@@ -502,6 +502,7 @@ func (s *StateDB) Suicide(addr common.InternalAddress) bool {
 	})
 	stateObject.markSuicided()
 	stateObject.data.Balance = new(big.Int)
+	stateObject.data.Size = new(big.Int)
 
 	return true
 }
@@ -532,7 +533,7 @@ func (s *StateDB) updateStateObject(obj *stateObject) {
 	// enough to track account updates at commit time, deletions need tracking
 	// at transaction boundary level to ensure we capture state clearing.
 	if s.snap != nil {
-		s.snapAccounts[obj.addrHash] = snapshot.SlimAccountRLP(obj.data.Nonce, obj.data.Balance, obj.data.Root, obj.data.CodeHash)
+		s.snapAccounts[obj.addrHash] = snapshot.SlimAccountRLP(obj.data.Nonce, obj.data.Balance, obj.data.Root, obj.data.CodeHash, obj.data.Size)
 	}
 }
 
@@ -801,6 +802,7 @@ func (s *StateDB) getDeletedStateObject(addr common.InternalAddress) *stateObjec
 				Balance:  acc.Balance,
 				CodeHash: acc.CodeHash,
 				Root:     common.BytesToHash(acc.Root),
+				Size:     acc.Size,
 			}
 			if len(data.CodeHash) == 0 {
 				data.CodeHash = emptyCodeHash
@@ -894,6 +896,7 @@ func (s *StateDB) CreateAccount(addr common.InternalAddress) {
 	newObj, prev := s.createObject(addr)
 	if prev != nil {
 		newObj.setBalance(prev.data.Balance)
+		newObj.setSize(prev.data.Size)
 	}
 }
 

core/vm/interface.go

@@ -30,6 +30,7 @@ type StateDB interface {
 	SubBalance(common.InternalAddress, *big.Int)
 	AddBalance(common.InternalAddress, *big.Int)
 	GetBalance(common.InternalAddress) *big.Int
+	GetSize(common.InternalAddress) *big.Int
 
 	GetNonce(common.InternalAddress) uint64
 	SetNonce(common.InternalAddress, uint64)