refactor StateManger/StateReader (#1886)

* refactor statemanger/statereader

* Update blockchain.go

Co-authored-by: CoderZhi <thecoderzhi@gmail.com>

action/protocol/account/protocol.go

@@ -103,14 +103,14 @@ func createAccount(sm protocol.StateManager, encodedAddr string, init *big.Int)
 		return errors.Wrap(err, "failed to get address public key hash from encoded address")
 	}
 	addrHash := hash.BytesToHash160(addr.Bytes())
-	err = sm.State(addrHash, &account)
+	_, err = sm.State(&account, protocol.LegacyKeyOption(addrHash))
 	switch errors.Cause(err) {
 	case nil:
 		return errors.Errorf("failed to create account %s", encodedAddr)
 	case state.ErrStateNotExist:
 		account.Balance = init
 		account.VotingWeight = big.NewInt(0)
-		if err := sm.PutState(addrHash, account); err != nil {
+		if _, err := sm.PutState(account, protocol.LegacyKeyOption(addrHash)); err != nil {
 			return errors.Wrapf(err, "failed to put state for account %x", addrHash)
 		}
 		return nil

action/protocol/account/protocol_test.go

@@ -33,21 +33,29 @@ func TestLoadOrCreateAccountState(t *testing.T) {
 	sm := mock_chainmanager.NewMockStateManager(ctrl)
 	cb := batch.NewCachedBatch()
 	sm.EXPECT().State(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
-			val, err := cb.Get("state", addrHash[:])
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
 			if err != nil {
-				return state.ErrStateNotExist
+				return 0, err
 			}
-			return state.Deserialize(account, val)
+			val, err := cb.Get("state", cfg.Key)
+			if err != nil {
+				return 0, state.ErrStateNotExist
+			}
+			return 0, state.Deserialize(account, val)
 		}).AnyTimes()
 	sm.EXPECT().PutState(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
+			if err != nil {
+				return 0, err
+			}
 			ss, err := state.Serialize(account)
 			if err != nil {
-				return err
+				return 0, err
 			}
-			cb.Put("state", addrHash[:], ss, "failed to put state")
-			return nil
+			cb.Put("state", cfg.Key, ss, "failed to put state")
+			return 0, nil
 		}).AnyTimes()
 
 	addrv1 := identityset.Address(27)
@@ -71,21 +79,29 @@ func TestProtocol_Initialize(t *testing.T) {
 	sm := mock_chainmanager.NewMockStateManager(ctrl)
 	cb := batch.NewCachedBatch()
 	sm.EXPECT().State(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
-			val, err := cb.Get("state", addrHash[:])
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
 			if err != nil {
-				return state.ErrStateNotExist
+				return 0, err
 			}
-			return state.Deserialize(account, val)
+			val, err := cb.Get("state", cfg.Key)
+			if err != nil {
+				return 0, state.ErrStateNotExist
+			}
+			return 0, state.Deserialize(account, val)
 		}).AnyTimes()
 	sm.EXPECT().PutState(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
+			if err != nil {
+				return 0, err
+			}
 			ss, err := state.Serialize(account)
 			if err != nil {
-				return err
+				return 0, err
 			}
-			cb.Put("state", addrHash[:], ss, "failed to put state")
-			return nil
+			cb.Put("state", cfg.Key, ss, "failed to put state")
+			return 0, nil
 		}).AnyTimes()
 
 	ge := config.Default.Genesis

action/protocol/account/transfer_test.go

@@ -41,21 +41,29 @@ func TestProtocol_HandleTransfer(t *testing.T) {
 	sm := mock_chainmanager.NewMockStateManager(ctrl)
 	cb := batch.NewCachedBatch()
 	sm.EXPECT().State(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
-			val, err := cb.Get("state", addrHash[:])
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
 			if err != nil {
-				return state.ErrStateNotExist
+				return 0, err
 			}
-			return state.Deserialize(account, val)
+			val, err := cb.Get("state", cfg.Key)
+			if err != nil {
+				return 0, state.ErrStateNotExist
+			}
+			return 0, state.Deserialize(account, val)
 		}).AnyTimes()
 	sm.EXPECT().PutState(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
+			if err != nil {
+				return 0, err
+			}
 			ss, err := state.Serialize(account)
 			if err != nil {
-				return err
+				return 0, err
 			}
-			cb.Put("state", addrHash[:], ss, "failed to put state")
-			return nil
+			cb.Put("state", cfg.Key, ss, "failed to put state")
+			return 0, nil
 		}).AnyTimes()
 
 	p := NewProtocol(rewarding.DepositGas)
@@ -104,9 +112,12 @@ func TestProtocol_HandleTransfer(t *testing.T) {
 	pubKeyBravo := hash.BytesToHash160(identityset.Address(29).Bytes())
 	pubKeyCharlie := hash.BytesToHash160(identityset.Address(30).Bytes())
 
-	require.NoError(sm.PutState(pubKeyAlfa, &accountAlfa))
-	require.NoError(sm.PutState(pubKeyBravo, &accountBravo))
-	require.NoError(sm.PutState(pubKeyCharlie, &accountCharlie))
+	_, err := sm.PutState(&accountAlfa, protocol.LegacyKeyOption(pubKeyAlfa))
+	require.NoError(err)
+	_, err = sm.PutState(&accountBravo, protocol.LegacyKeyOption(pubKeyBravo))
+	require.NoError(err)
+	_, err = sm.PutState(&accountCharlie, protocol.LegacyKeyOption(pubKeyCharlie))
+	require.NoError(err)
 
 	transfer, err := action.NewTransfer(
 		uint64(1),
@@ -138,17 +149,20 @@ func TestProtocol_HandleTransfer(t *testing.T) {
 	require.Equal(uint64(iotextypes.ReceiptStatus_Success), receipt.Status)
 
 	var acct state.Account
-	require.NoError(sm.State(pubKeyAlfa, &acct))
+	_, err = sm.State(&acct, protocol.LegacyKeyOption(pubKeyAlfa))
+	require.NoError(err)
 	require.Equal("40003", acct.Balance.String())
 	require.Equal(uint64(1), acct.Nonce)
-	require.NoError(sm.State(pubKeyBravo, &acct))
+	_, err = sm.State(&acct, protocol.LegacyKeyOption(pubKeyBravo))
+	require.NoError(err)
 	require.Equal("2", acct.Balance.String())
 
 	contractAcct := state.Account{
 		CodeHash: []byte("codeHash"),
 	}
 	contractAddr := hash.BytesToHash160(identityset.Address(32).Bytes())
-	require.NoError(sm.PutState(contractAddr, &contractAcct))
+	_, err = sm.PutState(&contractAcct, protocol.LegacyKeyOption(contractAddr))
+	require.NoError(err)
 	transfer, err = action.NewTransfer(
 		uint64(2),
 		big.NewInt(3),
@@ -162,7 +176,8 @@ func TestProtocol_HandleTransfer(t *testing.T) {
 	receipt, err = p.Handle(ctx, transfer, sm)
 	require.NoError(err)
 	require.Equal(uint64(iotextypes.ReceiptStatus_Failure), receipt.Status)
-	require.NoError(sm.State(pubKeyAlfa, &acct))
+	_, err = sm.State(&acct, protocol.LegacyKeyOption(pubKeyAlfa))
+	require.NoError(err)
 	require.Equal(uint64(2), acct.Nonce)
 	require.Equal("20003", acct.Balance.String())
 }

action/protocol/account/util/util.go

@@ -38,14 +38,14 @@ func LoadOrCreateAccount(sm protocol.StateManager, encodedAddr string) (*state.A
 		return &account, errors.Wrap(err, "failed to get address public key hash from encoded address")
 	}
 	addrHash := hash.BytesToHash160(addr.Bytes())
-	err = sm.State(addrHash, &account)
+	_, err = sm.State(&account, protocol.LegacyKeyOption(addrHash))
 	if err == nil {
 		return &account, nil
 	}
 	if errors.Cause(err) == state.ErrStateNotExist {
 		account.Balance = big.NewInt(0)
 		account.VotingWeight = big.NewInt(0)
-		if err := sm.PutState(addrHash, account); err != nil {
+		if _, err := sm.PutState(account, protocol.LegacyKeyOption(addrHash)); err != nil {
 			return nil, errors.Wrapf(err, "failed to put state for account %x", addrHash)
 		}
 		return &account, nil
@@ -56,7 +56,7 @@ func LoadOrCreateAccount(sm protocol.StateManager, encodedAddr string) (*state.A
 // LoadAccount loads an account state
 func LoadAccount(sm protocol.StateReader, addrHash hash.Hash160) (*state.Account, error) {
 	var account state.Account
-	if err := sm.State(addrHash, &account); err != nil {
+	if _, err := sm.State(&account, protocol.LegacyKeyOption(addrHash)); err != nil {
 		if errors.Cause(err) == state.ErrStateNotExist {
 			account = state.EmptyAccount()
 			return &account, nil
@@ -73,13 +73,14 @@ func StoreAccount(sm protocol.StateManager, encodedAddr string, account *state.A
 		return errors.Wrap(err, "failed to get address public key hash from encoded address")
 	}
 	addrHash := hash.BytesToHash160(addr.Bytes())
-	return sm.PutState(addrHash, account)
+	_, err = sm.PutState(account, protocol.LegacyKeyOption(addrHash))
+	return err
 }
 
 // Recorded tests if an account has been actually stored
 func Recorded(sm protocol.StateReader, addr address.Address) (bool, error) {
 	var account state.Account
-	err := sm.State(hash.BytesToHash160(addr.Bytes()), &account)
+	_, err := sm.State(&account, protocol.LegacyKeyOption(hash.BytesToHash160(addr.Bytes())))
 	if err == nil {
 		return true, nil
 	}
@@ -97,7 +98,7 @@ func AccountState(sr protocol.StateReader, encodedAddr string) (*state.Account,
 	}
 	pkHash := hash.BytesToHash160(addr.Bytes())
 	var account state.Account
-	if err := sr.State(pkHash, &account); err != nil {
+	if _, err := sr.State(&account, protocol.LegacyKeyOption(pkHash)); err != nil {
 		if errors.Cause(err) == state.ErrStateNotExist {
 			account = state.EmptyAccount()
 			return &account, nil
@@ -115,7 +116,7 @@ func AccountStateAtHeight(sr protocol.StateReader, encodedAddr string, height ui
 	}
 	pkHash := hash.BytesToHash160(addr.Bytes())
 	var account state.Account
-	if err := sr.State(pkHash, &account, protocol.BlockHeightOption(height)); err != nil {
+	if _, err := sr.State(&account, protocol.BlockHeightOption(height), protocol.LegacyKeyOption(pkHash)); err != nil {
 		if errors.Cause(err) == state.ErrStateNotExist {
 			account = state.EmptyAccount()
 			return &account, nil

action/protocol/execution/evm/contract_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/require"
 
+	"github.com/iotexproject/iotex-core/action/protocol"
 	accountutil "github.com/iotexproject/iotex-core/action/protocol/account/util"
 	"github.com/iotexproject/iotex-core/config"
 	"github.com/iotexproject/iotex-core/db"
@@ -40,21 +41,29 @@ func TestCreateContract(t *testing.T) {
 	sm := mock_chainmanager.NewMockStateManager(ctrl)
 	cb := batch.NewCachedBatch()
 	sm.EXPECT().State(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
-			val, err := cb.Get("state", addrHash[:])
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
 			if err != nil {
-				return state.ErrStateNotExist
+				return 0, err
 			}
-			return state.Deserialize(account, val)
+			val, err := cb.Get("state", cfg.Key)
+			if err != nil {
+				return 0, state.ErrStateNotExist
+			}
+			return 0, state.Deserialize(account, val)
 		}).AnyTimes()
 	sm.EXPECT().PutState(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
+			if err != nil {
+				return 0, err
+			}
 			ss, err := state.Serialize(account)
 			if err != nil {
-				return err
+				return 0, err
 			}
-			cb.Put("state", addrHash[:], ss, "failed to put state")
-			return nil
+			cb.Put("state", cfg.Key, ss, "failed to put state")
+			return 0, nil
 		}).AnyTimes()
 
 	flusher, err := db.NewKVStoreFlusher(db.NewMemKVStore(), cb)

action/protocol/execution/evm/evm_test.go

@@ -36,8 +36,8 @@ func TestExecuteContractFailure(t *testing.T) {
 	flusher, err := db.NewKVStoreFlusher(db.NewMemKVStore(), batch.NewCachedBatch())
 	require.NoError(t, err)
 	sm.EXPECT().GetDB().Return(flusher.KVStoreWithBuffer()).AnyTimes()
-	sm.EXPECT().State(gomock.Any(), gomock.Any()).Return(state.ErrStateNotExist).AnyTimes()
-	sm.EXPECT().PutState(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	sm.EXPECT().State(gomock.Any(), gomock.Any()).Return(uint64(0), state.ErrStateNotExist).AnyTimes()
+	sm.EXPECT().PutState(gomock.Any(), gomock.Any()).Return(uint64(0), nil).AnyTimes()
 	sm.EXPECT().Snapshot().Return(1).AnyTimes()
 
 	e, err := action.NewExecution(

action/protocol/execution/evm/evmstatedbadapter.go

@@ -289,7 +289,7 @@ func (stateDB *StateDBAdapter) Suicide(evmAddr common.Address) bool {
 	s.Balance = nil
 	s.Balance = big.NewInt(0)
 	addrHash := hash.BytesToHash160(evmAddr.Bytes())
-	if err := stateDB.sm.PutState(addrHash, s); err != nil {
+	if _, err := stateDB.sm.PutState(s, protocol.LegacyKeyOption(addrHash)); err != nil {
 		log.L().Error("Failed to kill contract.", zap.Error(err))
 		stateDB.logError(err)
 		return false
@@ -639,7 +639,7 @@ func (stateDB *StateDBAdapter) CommitContracts() error {
 		}
 		state := contract.SelfState()
 		// store the account (with new storage trie root) into account trie
-		if err := stateDB.sm.PutState(addr, state); err != nil {
+		if _, err := stateDB.sm.PutState(state, protocol.LegacyKeyOption(addr)); err != nil {
 			stateDB.logError(err)
 			return errors.Wrap(err, "failed to update pending account changes to trie")
 		}
@@ -658,7 +658,7 @@ func (stateDB *StateDBAdapter) CommitContracts() error {
 			return errors.Wrap(err, "failed to decode address hash")
 		}
 		copy(addr[:], addrBytes)
-		if err := stateDB.sm.DelState(addr); err != nil {
+		if _, err := stateDB.sm.DelState(protocol.LegacyKeyOption(addr)); err != nil {
 			stateDB.logError(err)
 			return errors.Wrapf(err, "failed to delete suicide account/contract %x", addr[:])
 		}

action/protocol/execution/evm/evmstatedbadapter_test.go

@@ -30,21 +30,29 @@ func initMockStateManager(ctrl *gomock.Controller) (protocol.StateManager, error
 	sm := mock_chainmanager.NewMockStateManager(ctrl)
 	cb := batch.NewCachedBatch()
 	sm.EXPECT().State(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
-			val, err := cb.Get("state", addrHash[:])
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
 			if err != nil {
-				return state.ErrStateNotExist
+				return 0, err
 			}
-			return state.Deserialize(account, val)
+			val, err := cb.Get("state", cfg.Key)
+			if err != nil {
+				return 0, state.ErrStateNotExist
+			}
+			return 0, state.Deserialize(account, val)
 		}).AnyTimes()
 	sm.EXPECT().PutState(gomock.Any(), gomock.Any()).DoAndReturn(
-		func(addrHash hash.Hash160, account interface{}) error {
+		func(account interface{}, opts ...protocol.StateOption) (uint64, error) {
+			cfg, err := protocol.CreateStateConfig(opts...)
+			if err != nil {
+				return 0, err
+			}
 			ss, err := state.Serialize(account)
 			if err != nil {
-				return err
+				return 0, err
 			}
-			cb.Put("state", addrHash[:], ss, "failed to put state")
-			return nil
+			cb.Put("state", cfg.Key, ss, "failed to put state")
+			return 0, nil
 		}).AnyTimes()
 	flusher, err := db.NewKVStoreFlusher(db.NewMemKVStore(), cb)
 	if err != nil {
@@ -408,7 +416,7 @@ func TestGetBalanceOnError(t *testing.T) {
 		errors.New("other error"),
 	}
 	for _, err := range errs {
-		sm.EXPECT().State(gomock.Any(), gomock.Any()).Return(err).Times(1)
+		sm.EXPECT().State(gomock.Any(), gomock.Any()).Return(uint64(0), err).Times(1)
 		addr := common.HexToAddress("test address")
 		stateDB := NewStateDBAdapter(sm, 1, true, hash.ZeroHash256)
 		amount := stateDB.GetBalance(addr)