Merge pull request #6555 from onflow/ramtin/evm-add-iterators-to-base…

…-storage

[EVM] Adding account/slot/code iterators to the base storage

fvm/evm/emulator/state/account.go

@@ -42,10 +42,16 @@ func NewAccount(
 	}
 }
 
+// HasCode returns true if account has code
 func (a *Account) HasCode() bool {
 	return a.CodeHash != gethTypes.EmptyCodeHash
 }
 
+// HasStoredValues returns true if account has stored values
+func (a *Account) HasStoredValues() bool {
+	return len(a.CollectionID) != 0
+}
+
 // Encode encodes the account
 func (a *Account) Encode() ([]byte, error) {
 	return rlp.EncodeToBytes(a)

fvm/evm/emulator/state/base.go

@@ -21,6 +21,8 @@ const (
 	CodesStorageIDKey = "CodesStorageIDKey"
 )
 
+var EmptyHash = gethCommon.Hash{}
+
 // BaseView implements a types.BaseView
 // it acts as the base layer of state queries for the stateDB
 // it stores accounts, codes and storage slots.
@@ -389,6 +391,65 @@ func (v *BaseView) NumberOfAccounts() uint64 {
 	return v.accounts.Size()
 }
 
+// AccountIterator returns an account iterator
+//
+// Warning! this is an expensive operation and should only be used
+// for testing and exporting state operations, while no changes
+// are applied to accounts. Note that the iteration order is not guaranteed.
+func (v *BaseView) AccountIterator() (*AccountIterator, error) {
+	itr, err := v.accounts.ReadOnlyIterator()
+	if err != nil {
+		return nil, err
+	}
+	return &AccountIterator{colIterator: itr}, nil
+}
+
+// CodeIterator returns a code iterator
+//
+// Warning! this is an expensive operation and should only be used
+// for testing and exporting state operations, while no changes
+// are applied to codes. Note that the iteration order is not guaranteed.
+func (v *BaseView) CodeIterator() (*CodeIterator, error) {
+	itr, err := v.codes.ReadOnlyIterator()
+	if err != nil {
+		return nil, err
+	}
+	return &CodeIterator{colIterator: itr}, nil
+}
+
+// AccountStorageIterator returns an account storage iterator
+// for the given address
+//
+// Warning! this is an expensive operation and should only be used
+// for testing and exporting state operations, while no changes
+// are applied to accounts. Note that the iteration order is not guaranteed.
+func (v *BaseView) AccountStorageIterator(
+	addr gethCommon.Address,
+) (*AccountStorageIterator, error) {
+	acc, err := v.getAccount(addr)
+	if err != nil {
+		return nil, err
+	}
+	if acc == nil || !acc.HasStoredValues() {
+		return nil, fmt.Errorf("account %s has no stored value", addr.String())
+	}
+	col, found := v.slots[addr]
+	if !found {
+		col, err = v.collectionProvider.CollectionByID(acc.CollectionID)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load storage collection for account %s: %w", addr.String(), err)
+		}
+	}
+	itr, err := col.ReadOnlyIterator()
+	if err != nil {
+		return nil, err
+	}
+	return &AccountStorageIterator{
+		address:     addr,
+		colIterator: itr,
+	}, nil
+}
+
 func (v *BaseView) fetchOrCreateCollection(path string) (collection *Collection, created bool, error error) {
 	collectionID, err := v.ledger.GetValue(v.rootAddress[:], []byte(path))
 	if err != nil {
@@ -592,8 +653,7 @@ func (v *BaseView) storeSlot(sk types.SlotAddress, data gethCommon.Hash) error {
 		return err
 	}
 
-	emptyValue := gethCommon.Hash{}
-	if data == emptyValue {
+	if data == EmptyHash {
 		delete(v.cachedSlots, sk)
 		return col.Remove(sk.Key.Bytes())
 	}
@@ -631,3 +691,83 @@ func (v *BaseView) getSlotCollection(acc *Account) (*Collection, error) {
 	}
 	return col, nil
 }
+
+// AccountIterator iterates over accounts
+type AccountIterator struct {
+	colIterator *CollectionIterator
+}
+
+// Next returns the next account
+// if no more accounts next would return nil (no error)
+func (ai *AccountIterator) Next() (*Account, error) {
+	_, value, err := ai.colIterator.Next()
+	if err != nil {
+		return nil, fmt.Errorf("account iteration failed: %w", err)
+	}
+	return DecodeAccount(value)
+}
+
+// CodeIterator iterates over codes stored in EVM
+// code storage only stores unique codes
+type CodeIterator struct {
+	colIterator *CollectionIterator
+}
+
+// Next returns the next code
+// if no more codes, it return nil (no error)
+func (ci *CodeIterator) Next() (
+	*CodeInContext,
+	error,
+) {
+	ch, encodedCC, err := ci.colIterator.Next()
+	if err != nil {
+		return nil, fmt.Errorf("code iteration failed: %w", err)
+	}
+	// no more keys
+	if ch == nil {
+		return nil, nil
+	}
+	if len(encodedCC) == 0 {
+		return nil,
+			fmt.Errorf("encoded code container is empty (code hash: %x)", ch)
+	}
+
+	codeCont, err := CodeContainerFromEncoded(encodedCC)
+	if err != nil {
+		return nil, fmt.Errorf("code container decoding failed (code hash: %x)", ch)
+
+	}
+	return &CodeInContext{
+		Hash:      gethCommon.BytesToHash(ch),
+		Code:      codeCont.Code(),
+		RefCounts: codeCont.RefCount(),
+	}, nil
+}
+
+// AccountStorageIterator iterates over slots of an account
+type AccountStorageIterator struct {
+	address     gethCommon.Address
+	colIterator *CollectionIterator
+}
+
+// Next returns the next slot in the storage
+// if no more keys, it returns nil (no error)
+func (asi *AccountStorageIterator) Next() (
+	*types.SlotEntry,
+	error,
+) {
+	k, v, err := asi.colIterator.Next()
+	if err != nil {
+		return nil, fmt.Errorf("account storage iteration failed: %w", err)
+	}
+	// no more keys
+	if k == nil {
+		return nil, nil
+	}
+	return &types.SlotEntry{
+		Address: asi.address,
+		Key:     gethCommon.BytesToHash(k),
+		Value:   gethCommon.BytesToHash(v),
+	}, nil
+
+}

fvm/evm/emulator/state/base_test.go

@@ -302,6 +302,170 @@ func TestBaseView(t *testing.T) {
 		require.Equal(t, uint64(1), view.NumberOfAccounts())
 	})
 
+	t.Run("test account iterator", 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)
+
+		accountCounts := 10
+		nonces := make(map[gethCommon.Address]uint64)
+		balances := make(map[gethCommon.Address]*uint256.Int)
+		codeHashes := make(map[gethCommon.Address]gethCommon.Hash)
+		for i := 0; i < accountCounts; i++ {
+			addr := testutils.RandomCommonAddress(t)
+			balance := testutils.RandomUint256Int(1000)
+			nonce := testutils.RandomBigInt(1000).Uint64()
+			code := testutils.RandomData(t)
+			codeHash := testutils.RandomCommonHash(t)
+
+			err = view.CreateAccount(addr, balance, nonce, code, codeHash)
+			require.NoError(t, err)
+
+			nonces[addr] = nonce
+			balances[addr] = balance
+			codeHashes[addr] = codeHash
+		}
+		err = view.Commit()
+		require.NoError(t, err)
+
+		ai, err := view.AccountIterator()
+		require.NoError(t, err)
+
+		counter := 0
+		for {
+			acc, err := ai.Next()
+			require.NoError(t, err)
+			if acc == nil {
+				break
+			}
+			require.Equal(t, nonces[acc.Address], acc.Nonce)
+			delete(nonces, acc.Address)
+			require.Equal(t, balances[acc.Address].Uint64(), acc.Balance.Uint64())
+			delete(balances, acc.Address)
+			require.Equal(t, codeHashes[acc.Address], acc.CodeHash)
+			delete(codeHashes, acc.Address)
+			counter += 1
+		}
+
+		require.Equal(t, accountCounts, counter)
+	})
+
+	t.Run("test code iterator", 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)
+
+		codeCounts := 10
+		codeByCodeHash := make(map[gethCommon.Hash][]byte)
+		refCountByCodeHash := make(map[gethCommon.Hash]uint64)
+		for i := 0; i < codeCounts; i++ {
+
+			code := testutils.RandomData(t)
+			codeHash := testutils.RandomCommonHash(t)
+			refCount := 0
+			// we add each code couple of times through different accounts
+			for j := 1; j <= i+1; j++ {
+				addr := testutils.RandomCommonAddress(t)
+				balance := testutils.RandomUint256Int(1000)
+				nonce := testutils.RandomBigInt(1000).Uint64()
+				err = view.CreateAccount(addr, balance, nonce, code, codeHash)
+				require.NoError(t, err)
+				refCount += 1
+			}
+			codeByCodeHash[codeHash] = code
+			refCountByCodeHash[codeHash] = uint64(refCount)
+		}
+		err = view.Commit()
+		require.NoError(t, err)
+
+		ci, err := view.CodeIterator()
+		require.NoError(t, err)
+
+		counter := 0
+		for {
+			cic, err := ci.Next()
+			require.NoError(t, err)
+			if cic == nil {
+				break
+			}
+			require.Equal(t, codeByCodeHash[cic.Hash], cic.Code)
+			delete(codeByCodeHash, cic.Hash)
+			require.Equal(t, refCountByCodeHash[cic.Hash], cic.RefCounts)
+			delete(refCountByCodeHash, cic.Hash)
+			counter += 1
+		}
+
+		require.Equal(t, codeCounts, counter)
+	})
+
+	t.Run("test account storage iterator", 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)
+
+		addr := testutils.RandomCommonAddress(t)
+		code := []byte("code")
+		balance := testutils.RandomUint256Int(1000)
+		nonce := testutils.RandomBigInt(1000).Uint64()
+		codeHash := gethCrypto.Keccak256Hash(code)
+		err = view.CreateAccount(addr, balance, nonce, code, codeHash)
+		require.NoError(t, err)
+
+		slotCounts := 10
+		values := make(map[gethCommon.Hash]gethCommon.Hash)
+
+		for i := 0; i < slotCounts; i++ {
+			key := testutils.RandomCommonHash(t)
+			value := testutils.RandomCommonHash(t)
+
+			err = view.UpdateSlot(
+				types.SlotAddress{
+					Address: addr,
+					Key:     key,
+				}, value)
+			require.NoError(t, err)
+			values[key] = value
+		}
+		err = view.Commit()
+		require.NoError(t, err)
+
+		asi, err := view.AccountStorageIterator(addr)
+		require.NoError(t, err)
+
+		counter := 0
+		for {
+			slot, err := asi.Next()
+			require.NoError(t, err)
+			if slot == nil {
+				break
+			}
+			require.Equal(t, addr, slot.Address)
+			require.Equal(t, values[slot.Key], slot.Value)
+			delete(values, slot.Key)
+			counter += 1
+		}
+
+		require.Equal(t, slotCounts, counter)
+
+		// test non existing address
+		addr2 := testutils.RandomCommonAddress(t)
+		_, err = view.AccountStorageIterator(addr2)
+		require.Error(t, err)
+
+		// test address without storage
+		err = view.CreateAccount(addr2, balance, nonce, code, codeHash)
+		require.NoError(t, err)
+
+		err = view.Commit()
+		require.NoError(t, err)
+
+		_, err = view.AccountStorageIterator(addr2)
+		require.Error(t, err)
+	})
+
 }
 
 func checkAccount(t *testing.T,

fvm/evm/emulator/state/code.go

@@ -3,6 +3,9 @@ package state
 import (
 	"encoding/binary"
 	"fmt"
+
+	gethCommon "github.com/onflow/go-ethereum/common"
+	"github.com/onflow/go-ethereum/rlp"
 )
 
 // CodeContainer contains codes and keeps
@@ -77,3 +80,24 @@ func (cc *CodeContainer) Encode() []byte {
 	copy(encoded[8:], cc.code)
 	return encoded
 }
+
+// CodeInContext captures a code in its context
+type CodeInContext struct {
+	Hash      gethCommon.Hash
+	Code      []byte
+	RefCounts uint64
+}
+
+// Encoded returns the encoded content of the code in context
+func (cic *CodeInContext) Encode() ([]byte, error) {
+	return rlp.EncodeToBytes(cic)
+}
+
+// CodeInContextFromEncoded constructs a code in context from the encoded data
+func CodeInContextFromEncoded(encoded []byte) (*CodeInContext, error) {
+	if len(encoded) == 0 {
+		return nil, nil
+	}
+	cic := &CodeInContext{}
+	return cic, rlp.DecodeBytes(encoded, cic)
+}