Use a local account object and refactor the store account interface (0…

…xPolygon#902)

* Use a local account object and refactor the store account interface

* Fix address/hash equal comparisions

jsonrpc/eth_endpoint.go

@@ -27,11 +27,16 @@ type ethTxPoolStore interface {
 	GetPendingTx(txHash types.Hash) (*types.Transaction, bool)
 }
 
+type Account struct {
+	Balance *big.Int
+	Nonce   uint64
+}
+
 type ethStateStore interface {
-	GetAccount(root types.Hash, addr types.Address) (*state.Account, error)
+	GetAccount(root types.Hash, addr types.Address) (*Account, error)
 	GetStorage(root types.Hash, addr types.Address, slot types.Hash) ([]byte, error)
 	GetForksInTime(blockNumber uint64) chain.ForksInTime
-	GetCode(hash types.Hash) ([]byte, error)
+	GetCode(root types.Hash, addr types.Address) ([]byte, error)
 }
 
 type ethBlockchainStore interface {
@@ -726,7 +731,7 @@ func (e *Eth) GetCode(address types.Address, filter BlockNumberOrHash) (interfac
 	}
 
 	emptySlice := []byte{}
-	acc, err := e.store.GetAccount(header.StateRoot, address)
+	code, err := e.store.GetCode(header.StateRoot, address)
 
 	if errors.Is(err, ErrStateNotFound) {
 		// If the account doesn't exist / is not initialized yet,
@@ -736,12 +741,6 @@ func (e *Eth) GetCode(address types.Address, filter BlockNumberOrHash) (interfac
 		return argBytesPtr(emptySlice), err
 	}
 
-	code, err := e.store.GetCode(types.BytesToHash(acc.CodeHash))
-	if err != nil {
-		// TODO This is just a workaround. Figure out why CodeHash is populated for regular accounts
-		return argBytesPtr(emptySlice), nil
-	}
-
 	return argBytesPtr(code), nil
 }
 

jsonrpc/eth_endpoint_test.go

@@ -5,7 +5,6 @@ import (
 	"math/big"
 	"testing"
 
-	"github.com/0xPolygon/polygon-edge/state"
 	"github.com/0xPolygon/polygon-edge/types"
 	"github.com/hashicorp/go-hclog"
 	"github.com/stretchr/testify/assert"
@@ -16,7 +15,7 @@ func TestEth_DecodeTxn(t *testing.T) {
 
 	tests := []struct {
 		name     string
-		accounts map[types.Address]*state.Account
+		accounts map[types.Address]*Account
 		arg      *txnArgs
 		res      *types.Transaction
 		err      error
@@ -77,7 +76,7 @@ func TestEth_DecodeTxn(t *testing.T) {
 		},
 		{
 			name: "should set latest nonce as default",
-			accounts: map[types.Address]*state.Account{
+			accounts: map[types.Address]*Account{
 				addr1: {
 					Nonce: 10,
 				},
@@ -171,11 +170,11 @@ func TestEth_GetNextNonce(t *testing.T) {
 	// Set up the mock accounts
 	accounts := []struct {
 		address types.Address
-		account *state.Account
+		account *Account
 	}{
 		{
 			types.StringToAddress("123"),
-			&state.Account{
+			&Account{
 				Nonce: 5,
 			},
 		},

jsonrpc/eth_state_test.go

@@ -1,9 +1,7 @@
 package jsonrpc
 
 import (
-	"bytes"
 	"errors"
-	"fmt"
 	"math/big"
 	"testing"
 
@@ -26,7 +24,7 @@ func TestEth_State_GetBalance(t *testing.T) {
 	store := &mockSpecialStore{
 		account: &mockAccount{
 			address: addr0,
-			account: &state.Account{
+			account: &Account{
 				Balance: big.NewInt(100),
 			},
 			storage: make(map[types.Hash][]byte),
@@ -158,7 +156,7 @@ func TestEth_State_GetTransactionCount(t *testing.T) {
 	store := &mockSpecialStore{
 		account: &mockAccount{
 			address: addr0,
-			account: &state.Account{
+			account: &Account{
 				Balance: big.NewInt(100),
 				Nonce:   100,
 			},
@@ -276,10 +274,9 @@ func TestEth_State_GetCode(t *testing.T) {
 	store := &mockSpecialStore{
 		account: &mockAccount{
 			address: addr0,
-			account: &state.Account{
-				Balance:  big.NewInt(100),
-				Nonce:    100,
-				CodeHash: types.BytesToHash(addr0.Bytes()).Bytes(),
+			account: &Account{
+				Balance: big.NewInt(100),
+				Nonce:   100,
 			},
 			code: code0,
 		},
@@ -401,7 +398,7 @@ func TestEth_State_GetStorageAt(t *testing.T) {
 	store := &mockSpecialStore{
 		account: &mockAccount{
 			address: addr0,
-			account: &state.Account{
+			account: &Account{
 				Balance: big.NewInt(100),
 				Nonce:   100,
 			},
@@ -550,7 +547,7 @@ func TestEth_State_GetStorageAt(t *testing.T) {
 			for addr, storage := range tt.initialStorage {
 				store.account = &mockAccount{
 					address: addr,
-					account: &state.Account{
+					account: &Account{
 						Balance: big.NewInt(100),
 						Nonce:   100,
 					},
@@ -598,7 +595,7 @@ func getExampleStore() *mockSpecialStore {
 	return &mockSpecialStore{
 		account: &mockAccount{
 			address: addr0,
-			account: &state.Account{
+			account: &Account{
 				Balance: big.NewInt(100),
 				Nonce:   0,
 			},
@@ -777,15 +774,15 @@ type mockSpecialStore struct {
 }
 
 func (m *mockSpecialStore) GetBlockByHash(hash types.Hash, full bool) (*types.Block, bool) {
-	if m.block.Header.Hash.String() != hash.String() {
+	if m.block.Header.Hash != hash {
 		return nil, false
 	}
 
 	return m.block, true
 }
 
-func (m *mockSpecialStore) GetAccount(root types.Hash, addr types.Address) (*state.Account, error) {
-	if m.account.address.String() != addr.String() {
+func (m *mockSpecialStore) GetAccount(root types.Hash, addr types.Address) (*Account, error) {
+	if m.account.address != addr {
 		return nil, ErrStateNotFound
 	}
 
@@ -809,7 +806,7 @@ func (m *mockSpecialStore) GetNonce(addr types.Address) uint64 {
 }
 
 func (m *mockSpecialStore) GetStorage(root types.Hash, addr types.Address, slot types.Hash) ([]byte, error) {
-	if m.account.address.String() != addr.String() {
+	if m.account.address != addr {
 		return nil, ErrStateNotFound
 	}
 
@@ -823,12 +820,12 @@ func (m *mockSpecialStore) GetStorage(root types.Hash, addr types.Address, slot
 	return val, nil
 }
 
-func (m *mockSpecialStore) GetCode(hash types.Hash) ([]byte, error) {
-	if bytes.Equal(m.account.account.CodeHash, hash.Bytes()) {
-		return m.account.code, nil
+func (m *mockSpecialStore) GetCode(root types.Hash, addr types.Address) ([]byte, error) {
+	if m.account.address != addr {
+		return nil, ErrStateNotFound
 	}
 
-	return nil, fmt.Errorf("code not found")
+	return m.account.code, nil
 }
 
 func (m *mockSpecialStore) GetForksInTime(blockNumber uint64) chain.ForksInTime {

jsonrpc/eth_txpool_test.go

@@ -4,7 +4,6 @@ import (
 	"math/big"
 	"testing"
 
-	"github.com/0xPolygon/polygon-edge/state"
 	"github.com/0xPolygon/polygon-edge/types"
 	"github.com/stretchr/testify/assert"
 )
@@ -69,7 +68,7 @@ func (m *mockStoreTxn) AddAccount(addr types.Address) *mockAccount {
 
 	acct := &mockAccount{
 		address: addr,
-		account: &state.Account{},
+		account: &Account{},
 		storage: make(map[types.Hash][]byte),
 	}
 	m.accounts[addr] = acct
@@ -81,7 +80,7 @@ func (m *mockStoreTxn) Header() *types.Header {
 	return &types.Header{}
 }
 
-func (m *mockStoreTxn) GetAccount(root types.Hash, addr types.Address) (*state.Account, error) {
+func (m *mockStoreTxn) GetAccount(root types.Hash, addr types.Address) (*Account, error) {
 	acct, ok := m.accounts[addr]
 	if !ok {
 		return nil, ErrStateNotFound

jsonrpc/mocks_test.go

@@ -5,14 +5,13 @@ import (
 	"sync"
 
 	"github.com/0xPolygon/polygon-edge/blockchain"
-	"github.com/0xPolygon/polygon-edge/state"
 	"github.com/0xPolygon/polygon-edge/types"
 )
 
 type mockAccount struct {
 	address types.Address
 	code    []byte
-	account *state.Account
+	account *Account
 	storage map[types.Hash][]byte
 }
 
@@ -21,9 +20,7 @@ func (m *mockAccount) Storage(k types.Hash, v []byte) {
 }
 
 func (m *mockAccount) Code(code []byte) {
-	codeHash := types.BytesToHash(m.address.Bytes())
 	m.code = code
-	m.account.CodeHash = codeHash.Bytes()
 }
 
 func (m *mockAccount) Nonce(n uint64) {
@@ -51,7 +48,7 @@ type mockStore struct {
 	subscription *blockchain.MockSubscription
 	receiptsLock sync.Mutex
 	receipts     map[types.Hash][]*types.Receipt
-	accounts     map[types.Address]*state.Account
+	accounts     map[types.Address]*Account
 
 	// headers is the list of historical headers
 	headers []*types.Header
@@ -61,7 +58,7 @@ func newMockStore() *mockStore {
 	m := &mockStore{
 		header:       &types.Header{Number: 0},
 		subscription: blockchain.NewMockSubscription(),
-		accounts:     map[types.Address]*state.Account{},
+		accounts:     map[types.Address]*Account{},
 	}
 	m.addHeader(m.header)
 
@@ -109,15 +106,15 @@ func (m *mockStore) emitEvent(evnt *mockEvent) {
 	m.subscription.Push(bEvnt)
 }
 
-func (m *mockStore) GetAccount(root types.Hash, addr types.Address) (*state.Account, error) {
+func (m *mockStore) GetAccount(root types.Hash, addr types.Address) (*Account, error) {
 	if acc, ok := m.accounts[addr]; ok {
 		return acc, nil
 	}
 
 	return nil, ErrStateNotFound
 }
 
-func (m *mockStore) SetAccount(addr types.Address, account *state.Account) {
+func (m *mockStore) SetAccount(addr types.Address, account *Account) {
 	m.accounts[addr] = account
 }
 

server/server.go

@@ -460,7 +460,7 @@ func (j *jsonRPCHub) getState(root types.Hash, slot []byte) ([]byte, error) {
 	return result, nil
 }
 
-func (j *jsonRPCHub) GetAccount(root types.Hash, addr types.Address) (*state.Account, error) {
+func (j *jsonRPCHub) getAccountImpl(root types.Hash, addr types.Address) (*state.Account, error) {
 	obj, err := j.getState(root, addr.Bytes())
 	if err != nil {
 		return nil, err
@@ -474,35 +474,51 @@ func (j *jsonRPCHub) GetAccount(root types.Hash, addr types.Address) (*state.Acc
 	return &account, nil
 }
 
+func (j *jsonRPCHub) GetAccount(root types.Hash, addr types.Address) (*jsonrpc.Account, error) {
+	acct, err := j.getAccountImpl(root, addr)
+	if err != nil {
+		return nil, err
+	}
+
+	account := &jsonrpc.Account{
+		Nonce:   acct.Nonce,
+		Balance: new(big.Int).Set(acct.Balance),
+	}
+
+	return account, nil
+}
+
 // GetForksInTime returns the active forks at the given block height
 func (j *jsonRPCHub) GetForksInTime(blockNumber uint64) chain.ForksInTime {
 	return j.Executor.GetForksInTime(blockNumber)
 }
 
 func (j *jsonRPCHub) GetStorage(root types.Hash, addr types.Address, slot types.Hash) ([]byte, error) {
-	account, err := j.GetAccount(root, addr)
-
+	account, err := j.getAccountImpl(root, addr)
 	if err != nil {
 		return nil, err
 	}
 
 	obj, err := j.getState(account.Root, slot.Bytes())
-
 	if err != nil {
 		return nil, err
 	}
 
 	return obj, nil
 }
 
-func (j *jsonRPCHub) GetCode(hash types.Hash) ([]byte, error) {
-	res, ok := j.state.GetCode(hash)
+func (j *jsonRPCHub) GetCode(root types.Hash, addr types.Address) ([]byte, error) {
+	account, err := j.getAccountImpl(root, addr)
+	if err != nil {
+		return nil, err
+	}
 
+	code, ok := j.state.GetCode(account.Root)
 	if !ok {
 		return nil, fmt.Errorf("unable to fetch code")
 	}
 
-	return res, nil
+	return code, nil
 }
 
 func (j *jsonRPCHub) ApplyTxn(