[Flow EVM] update StateDB to store codes by code hash

fvm/evm/emulator/state/base.go

@@ -131,11 +131,12 @@ func (v *BaseView) GetCode(addr gethCommon.Address) ([]byte, error) {
 
 // GetCodeHash returns the code hash of an address
 //
-// for non-existent accounts or accounts without a code (e.g. EOAs) it returns default empty
+// for non-existent accounts it returns gethCommon.Hash{}
+// and for accounts without a code (e.g. EOAs) it returns default empty
 // hash value (gethTypes.EmptyCodeHash)
 func (v *BaseView) GetCodeHash(addr gethCommon.Address) (gethCommon.Hash, error) {
 	acc, err := v.getAccount(addr)
-	codeHash := gethTypes.EmptyCodeHash
+	codeHash := gethCommon.Hash{}
 	if acc != nil {
 		codeHash = acc.CodeHash
 	}
@@ -204,7 +205,7 @@ func (v *BaseView) CreateAccount(
 	var colID []byte
 	// if is an smart contract account
 	if len(code) > 0 {
-		err := v.storeCode(addr, code)
+		err := v.updateAccountCode(addr, code, codeHash)
 		if err != nil {
 			return err
 		}
@@ -235,17 +236,17 @@ func (v *BaseView) UpdateAccount(
 	if acc == nil {
 		return v.CreateAccount(addr, balance, nonce, code, codeHash)
 	}
-	// if it has a code change
-	if codeHash != acc.CodeHash {
-		err := v.storeCode(addr, code)
-		if err != nil {
-			return err
-		}
-		// TODO: maybe purge the state in the future as well
-		// currently the behaviour of stateDB doesn't purge the data
-		// We don't need to check if the code is empty and we purge the state
-		// this is not possible right now.
+
+	// update account code
+	err = v.updateAccountCode(addr, code, codeHash)
+	if err != nil {
+		return err
 	}
+	// TODO: maybe purge the state in the future as well
+	// currently the behaviour of stateDB doesn't purge the data
+	// We don't need to check if the code is empty and we purge the state
+	// this is not possible right now.
+
 	newAcc := NewAccount(addr, balance, nonce, codeHash, acc.CollectionID)
 	// no need to update the cache , storeAccount would update the cache
 	return v.storeAccount(newAcc)
@@ -263,23 +264,23 @@ func (v *BaseView) DeleteAccount(addr gethCommon.Address) error {
 		return fmt.Errorf("account doesn't exist to be deleted")
 	}
 
-	// 2. update the cache
+	// 2. remove the code
+	if acc.HasCode() {
+		err = v.updateAccountCode(addr, nil, gethTypes.EmptyCodeHash)
+		if err != nil {
+			return err
+		}
+	}
+
+	// 3. update the cache
 	delete(v.cachedAccounts, addr)
 
-	// 3. collections
+	// 4. collections
 	err = v.accounts.Remove(addr.Bytes())
 	if err != nil {
 		return err
 	}
 
-	// 4. remove the code
-	if acc.HasCode() {
-		err = v.storeCode(addr, nil)
-		if err != nil {
-			return err
-		}
-	}
-
 	// 5. remove storage slots
 	if len(acc.CollectionID) > 0 {
 		col, found := v.slots[addr]
@@ -338,6 +339,16 @@ func (v *BaseView) Commit() error {
 	return nil
 }
 
+// NumberOfContracts returns the number of unique contracts
+func (v *BaseView) NumberOfContracts() uint64 {
+	return v.codes.Size()
+}
+
+// NumberOfContracts returns the number of accounts
+func (v *BaseView) NumberOfAccounts() uint64 {
+	return v.accounts.Size()
+}
+
 func (v *BaseView) fetchOrCreateCollection(path string) (collection *Collection, created bool, error error) {
 	collectionID, err := v.ledger.GetValue(v.rootAddress[:], []byte(path))
 	if err != nil {
@@ -391,29 +402,110 @@ func (v *BaseView) getCode(addr gethCommon.Address) ([]byte, error) {
 	if found {
 		return code, nil
 	}
-	// check if account exist in cache and has codeHash
-	acc, found := v.cachedAccounts[addr]
-	if found && !acc.HasCode() {
+
+	// get account
+	acc, err := v.getAccount(addr)
+	if err != nil {
+		return nil, err
+	}
+
+	if acc == nil || !acc.HasCode() {
 		return nil, nil
 	}
-	// then collect it from the code collection
-	code, err := v.codes.Get(addr.Bytes())
+
+	// collect the container from the code collection by codeHash
+	encoded, err := v.codes.Get(acc.CodeHash.Bytes())
 	if err != nil {
 		return nil, err
 	}
-	if code != nil {
+	if len(encoded) == 0 {
+		return nil, nil
+	}
+
+	codeCont, err := CodeContainerFromEncoded(encoded)
+	if err != nil {
+		return nil, err
+	}
+	code = codeCont.Code()
+	if len(code) > 0 {
 		v.cachedCodes[addr] = code
 	}
 	return code, nil
 }
 
-func (v *BaseView) storeCode(addr gethCommon.Address, code []byte) error {
-	if len(code) == 0 {
+func (v *BaseView) updateAccountCode(addr gethCommon.Address, code []byte, codeHash gethCommon.Hash) error {
+	// get account
+	acc, err := v.getAccount(addr)
+	if err != nil {
+		return err
+	}
+	// if is a new account
+	if acc == nil {
+		if len(code) == 0 {
+			return nil
+		}
+		v.cachedCodes[addr] = code
+		return v.addCode(code, codeHash)
+	}
+
+	// skip if is the same code
+	if acc.CodeHash == codeHash {
+		return nil
+	}
+
+	// clean old code first if exist
+	if acc.HasCode() {
 		delete(v.cachedCodes, addr)
-		return v.codes.Remove(addr.Bytes())
+		err = v.removeCode(acc.CodeHash)
+		if err != nil {
+			return err
+		}
+	}
+
+	// add new code
+	if len(code) == 0 {
+		return nil
 	}
 	v.cachedCodes[addr] = code
-	return v.codes.Set(addr.Bytes(), code)
+	return v.addCode(code, codeHash)
+}
+
+func (v *BaseView) removeCode(codeHash gethCommon.Hash) error {
+	encoded, err := v.codes.Get(codeHash.Bytes())
+	if err != nil {
+		return err
+	}
+	if len(encoded) == 0 {
+		return nil
+	}
+
+	cc, err := CodeContainerFromEncoded(encoded)
+	if err != nil {
+		return err
+	}
+	if cc.DecRefCount() {
+		return v.codes.Remove(codeHash.Bytes())
+	}
+	return v.codes.Set(codeHash.Bytes(), cc.Encode())
+}
+
+func (v *BaseView) addCode(code []byte, codeHash gethCommon.Hash) error {
+	encoded, err := v.codes.Get(codeHash.Bytes())
+	if err != nil {
+		return err
+	}
+	// if is the first time the code is getting deployed
+	if len(encoded) == 0 {
+		return v.codes.Set(codeHash.Bytes(), NewCodeContainer(code).Encode())
+	}
+
+	// otherwise update the cc
+	cc, err := CodeContainerFromEncoded(encoded)
+	if err != nil {
+		return err
+	}
+	cc.IncRefCount()
+	return v.codes.Set(codeHash.Bytes(), cc.Encode())
 }
 
 func (v *BaseView) getSlot(sk types.SlotAddress) (gethCommon.Hash, error) {

fvm/evm/emulator/state/base_test.go

@@ -6,6 +6,7 @@ import (
 
 	gethCommon "github.com/ethereum/go-ethereum/common"
 	gethTypes "github.com/ethereum/go-ethereum/core/types"
+	gethCrypto "github.com/ethereum/go-ethereum/crypto"
 	"github.com/stretchr/testify/require"
 
 	"github.com/onflow/flow-go/fvm/evm/emulator/state"
@@ -210,6 +211,97 @@ func TestBaseView(t *testing.T) {
 		require.Equal(t, false, addrFound)
 		require.Equal(t, false, slotFound)
 	})
+
+	t.Run("test code storage", func(t *testing.T) {
+		ledger := testutils.GetSimpleValueStore()
+		rootAddr := flow.Address{1, 2, 3, 4, 5, 6, 7, 8}
+		view, err := state.NewBaseView(ledger, rootAddr)
+		require.NoError(t, err)
+
+		bal := new(big.Int)
+		nonce := uint64(0)
+
+		addr1 := testutils.RandomCommonAddress(t)
+		var code1 []byte
+		codeHash1 := gethTypes.EmptyCodeHash
+		err = view.CreateAccount(addr1, bal, nonce, code1, codeHash1)
+		require.NoError(t, err)
+
+		ret, err := view.GetCode(addr1)
+		require.NoError(t, err)
+		require.Equal(t, code1, ret)
+
+		addr2 := testutils.RandomCommonAddress(t)
+		code2 := []byte("code2")
+		codeHash2 := gethCrypto.Keccak256Hash(code2)
+		err = view.CreateAccount(addr2, bal, nonce, code2, codeHash2)
+		require.NoError(t, err)
+
+		ret, err = view.GetCode(addr2)
+		require.NoError(t, err)
+		require.Equal(t, code2, ret)
+
+		err = view.Commit()
+		require.NoError(t, err)
+		orgSize := ledger.TotalStorageSize()
+		require.Equal(t, uint64(1), view.NumberOfContracts())
+
+		err = view.UpdateAccount(addr1, bal, nonce, code2, codeHash2)
+		require.NoError(t, err)
+
+		err = view.Commit()
+		require.NoError(t, err)
+		require.Equal(t, orgSize, ledger.TotalStorageSize())
+		require.Equal(t, uint64(1), view.NumberOfContracts())
+
+		ret, err = view.GetCode(addr1)
+		require.NoError(t, err)
+		require.Equal(t, code2, ret)
+
+		// now remove the code from account 1
+		err = view.UpdateAccount(addr1, bal, nonce, code1, codeHash1)
+		require.NoError(t, err)
+
+		// there should not be any side effect on the code return for account 2
+		// and no impact on storage size
+		ret, err = view.GetCode(addr2)
+		require.NoError(t, err)
+		require.Equal(t, code2, ret)
+
+		ret, err = view.GetCode(addr1)
+		require.NoError(t, err)
+		require.Equal(t, code1, ret)
+
+		err = view.Commit()
+		require.NoError(t, err)
+		require.Equal(t, orgSize, ledger.TotalStorageSize())
+		require.Equal(t, uint64(1), view.NumberOfContracts())
+
+		// now update account 2 and there should a reduction in storage
+		err = view.UpdateAccount(addr2, bal, nonce, code1, codeHash1)
+		require.NoError(t, err)
+
+		ret, err = view.GetCode(addr2)
+		require.NoError(t, err)
+		require.Equal(t, code1, ret)
+
+		err = view.Commit()
+		require.NoError(t, err)
+		require.Greater(t, orgSize, ledger.TotalStorageSize())
+		require.Equal(t, uint64(0), view.NumberOfContracts())
+
+		// delete account 2
+		err = view.DeleteAccount(addr2)
+		require.NoError(t, err)
+
+		ret, err = view.GetCode(addr2)
+		require.NoError(t, err)
+		require.Len(t, ret, 0)
+
+		require.Greater(t, orgSize, ledger.TotalStorageSize())
+		require.Equal(t, uint64(1), view.NumberOfAccounts())
+	})
+
 }
 
 func checkAccount(t *testing.T,