refactor: Migrate to Collections

.github/workflows/lint.yml

@@ -28,7 +28,7 @@ jobs:
       - uses: golangci/golangci-lint-action@v3
         with:
           version: latest
-          args: --timeout 10m --tests=false --skip-dirs="e2e"
+          args: --timeout 10m --tests=false --exclude-dirs="e2e"
           skip-pkg-cache: true
 
   protobuf:

app/app.go

@@ -610,7 +610,7 @@ func NewApp(
 
 	app.WasmStorageKeeper = *wasmstoragekeeper.NewKeeper(
 		appCodec,
-		keys[wasmstoragetypes.StoreKey],
+		runtime.NewKVStoreService(keys[wasmstoragetypes.StoreKey]),
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 		contractKeeper,
 	)

x/randomness/genesis.go

@@ -9,12 +9,20 @@ import (
 
 // InitGenesis puts data from genesis state into store.
 func InitGenesis(ctx sdk.Context, k keeper.Keeper, data types.GenesisState) {
-	k.SetSeed(ctx, data.Seed)
+	err := k.Seed.Set(ctx, data.Seed)
+	if err != nil {
+		panic(err)
+	}
 }
 
 // ExportGenesis extracts data from store to genesis state.
 func ExportGenesis(ctx sdk.Context, k keeper.Keeper) types.GenesisState {
+	seed, err := k.Seed.Get(ctx)
+	if err != nil {
+		return types.GenesisState{}
+	}
+
 	return types.GenesisState{
-		Seed: k.GetSeed(ctx),
+		Seed: seed,
 	}
 }

x/randomness/keeper/abci.go

@@ -56,7 +56,10 @@ func (h *ProposalHandler) PrepareProposalHandler(
 
 		// Seed transaction
 		// alpha = (seed_{i-1} || timestamp)
-		prevSeed := keeper.GetSeed(ctx)
+		prevSeed, err := keeper.GetSeed(ctx)
+		if err != nil {
+			return nil, err
+		}
 		if prevSeed == "" {
 			return nil, fmt.Errorf("previous seed is empty - this should never happen")
 		}
@@ -161,7 +164,10 @@ func (h *ProposalHandler) ProcessProposalHandler(
 			return nil, err
 		}
 
-		prevSeed := keeper.GetSeed(ctx)
+		prevSeed, err := keeper.GetSeed(ctx)
+		if err != nil {
+			return nil, err
+		}
 		if prevSeed == "" {
 			panic("seed should never be empty")
 		}

x/randomness/keeper/keeper.go

@@ -4,39 +4,50 @@ import (
 	"context"
 	"fmt"
 
+	"cosmossdk.io/collections"
+	storetypes "cosmossdk.io/core/store"
 	"cosmossdk.io/log"
-	storetypes "cosmossdk.io/store/types"
 
 	"github.com/cosmos/cosmos-sdk/codec"
 	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
 
 	"github.com/sedaprotocol/seda-chain/x/randomness/types"
 )
 
+var (
+	SeedPrefix         = collections.NewPrefix(0)
+	ValidatorVRFPrefix = collections.NewPrefix(1)
+)
+
+// GetValidatorVRFKeyPrefixFull gets the key for the validator VRF object.
+func GetValidatorVRFKeyPrefixFull(consensusAddr sdk.ConsAddress) []byte {
+	return append(ValidatorVRFPrefix, address.MustLengthPrefix(consensusAddr)...)
+}
+
 type Keeper struct {
-	cdc      codec.BinaryCodec
-	storeKey storetypes.StoreKey
+	Schema              collections.Schema
+	Seed                collections.Item[string]
+	ValidatorVRFPubKeys collections.Map[string, cryptotypes.PubKey]
 }
 
-func NewKeeper(cdc codec.BinaryCodec, storeKey storetypes.StoreKey) *Keeper {
+func NewKeeper(cdc codec.BinaryCodec, storeService storetypes.KVStoreService) *Keeper {
+	sb := collections.NewSchemaBuilder(storeService)
+
 	return &Keeper{
-		cdc:      cdc,
-		storeKey: storeKey,
+		Seed:                collections.NewItem(sb, SeedPrefix, "seed", collections.StringValue),
+		ValidatorVRFPubKeys: collections.NewMap(sb, ValidatorVRFPrefix, "validator-vrf-pubkeys", collections.StringKey, codec.CollInterfaceValue[cryptotypes.PubKey](cdc)),
 	}
 }
 
 // GetSeed returns the seed.
-func (k Keeper) GetSeed(ctx sdk.Context) string {
-	store := ctx.KVStore(k.storeKey)
-	bz := store.Get(types.KeyPrefixSeed)
-	return string(bz)
-}
-
-// SetSeed stores the seed.
-func (k Keeper) SetSeed(ctx sdk.Context, seed string) {
-	store := ctx.KVStore(k.storeKey)
-	store.Set(types.KeyPrefixSeed, []byte(seed))
+func (k Keeper) GetSeed(ctx sdk.Context) (string, error) {
+	seed, err := k.Seed.Get(ctx)
+	if err != nil {
+		return "", err
+	}
+	return seed, nil
 }
 
 // GetValidatorVRFPubKey retrieves from the store the VRF public key
@@ -46,18 +57,12 @@ func (k Keeper) GetValidatorVRFPubKey(ctx sdk.Context, consensusAddr string) (cr
 	if err != nil {
 		return nil, err
 	}
-
-	store := ctx.KVStore(k.storeKey)
-	bz := store.Get(types.GetValidatorVRFKey(addr))
-	if bz == nil {
-		return nil, fmt.Errorf("vrf pubkey not found for %s", consensusAddr)
-	}
-
-	var vrfPubKey cryptotypes.PubKey
-	err = k.cdc.UnmarshalInterface(bz, &vrfPubKey)
+	validatorVRFKeyPrefixFull := GetValidatorVRFKeyPrefixFull(addr)
+	vrfPubKey, err := k.ValidatorVRFPubKeys.Get(ctx, string(validatorVRFKeyPrefixFull))
 	if err != nil {
 		return nil, err
 	}
+
 	return vrfPubKey, nil
 }
 
@@ -66,14 +71,8 @@ func (k Keeper) SetValidatorVRFPubKey(goCtx context.Context, consensusAddr strin
 	if err != nil {
 		return err
 	}
-
-	store := sdk.UnwrapSDKContext(goCtx).KVStore(k.storeKey)
-	bz, err := k.cdc.MarshalInterface(vrfPubKey)
-	if err != nil {
-		return err
-	}
-	store.Set(types.GetValidatorVRFKey(addr), bz)
-	return nil
+	validatorVRFKeyPrefixFull := GetValidatorVRFKeyPrefixFull(addr)
+	return k.ValidatorVRFPubKeys.Set(goCtx, string(validatorVRFKeyPrefixFull), vrfPubKey)
 }
 
 func (k Keeper) Logger(ctx sdk.Context) log.Logger {

x/randomness/keeper/msg_server.go

@@ -22,6 +22,12 @@ var _ types.MsgServer = msgServer{}
 
 func (k msgServer) NewSeed(goCtx context.Context, msg *types.MsgNewSeed) (*types.MsgNewSeedResponse, error) {
 	sdkCtx := sdk.UnwrapSDKContext(goCtx)
-	k.Keeper.SetSeed(sdkCtx, msg.Beta)
+	err := k.Keeper.Seed.Set(sdkCtx, msg.Beta)
+	if err != nil {
+		return nil, err
+	}
+
+	// TO-DO event?
+
 	return &types.MsgNewSeedResponse{}, nil
 }

x/randomness/keeper/querier.go

@@ -20,8 +20,13 @@ func NewQuerierImpl(keeper Keeper) *Querier {
 
 func (q Querier) Seed(c context.Context, _ *types.QuerySeedRequest) (*types.QuerySeedResponse, error) {
 	ctx := sdk.UnwrapSDKContext(c)
+	seed, err := q.Keeper.Seed.Get(ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	return &types.QuerySeedResponse{
-		Seed:        q.GetSeed(ctx),
+		Seed:        seed,
 		BlockHeight: ctx.BlockHeight(),
 	}, nil
 }

x/randomness/types/keys.go

@@ -1,10 +1,5 @@
 package types
 
-import (
-	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/address"
-)
-
 const (
 	// ModuleName defines the module name
 	ModuleName = "randomness"
@@ -15,14 +10,3 @@ const (
 	// RouterKey defines the module's message routing key
 	RouterKey = ModuleName
 )
-
-// KeyPrefixSeed defines prefix to store the current block's seed.
-var KeyPrefixSeed = []byte{0x00}
-
-// KeyPrefixValidatorVRF defines prefix to store the validator VRF object.
-var KeyPrefixValidatorVRF = []byte{0x01}
-
-// GetValidatorVRFKey gets the key for the validator VRF object.
-func GetValidatorVRFKey(consensusAddr sdk.ConsAddress) []byte {
-	return append(KeyPrefixValidatorVRF, address.MustLengthPrefix(consensusAddr)...)
-}

x/wasm-storage/genesis.go

@@ -13,25 +13,38 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, data types.GenesisState) {
 		wasm := data.Wasms[i]
 		if wasm.WasmType == types.WasmTypeDataRequest ||
 			wasm.WasmType == types.WasmTypeTally {
-			k.SetDataRequestWasm(ctx, &wasm)
+			err := k.SetDataRequestWasm(ctx, &wasm)
+			if err != nil {
+				panic(err)
+			}
 		}
 		if wasm.WasmType == types.WasmTypeDataRequestExecutor ||
 			wasm.WasmType == types.WasmTypeRelayer {
-			k.SetOverlayWasm(ctx, &wasm)
+			err := k.SetOverlayWasm(ctx, &wasm)
+			if err != nil {
+				panic(err)
+			}
 		}
 	}
 	if data.ProxyContractRegistry != "" {
 		proxyAddr, err := sdk.AccAddressFromBech32(data.ProxyContractRegistry)
 		if err != nil {
 			panic(err)
 		}
-		k.SetProxyContractRegistry(ctx, proxyAddr)
+		err = k.ProxyContractRegistry.Set(ctx, proxyAddr)
+		if err != nil {
+			panic(err)
+		}
 	}
 }
 
 // ExportGenesis extracts all data from store to genesis state.
 func ExportGenesis(ctx sdk.Context, k keeper.Keeper) types.GenesisState {
 	wasms := k.GetAllWasms(ctx)
-	proxy := k.GetProxyContractRegistry(ctx)
-	return types.NewGenesisState(wasms, proxy.String())
+	proxy, err := k.ProxyContractRegistry.Get(ctx)
+	if err != nil {
+		return types.NewGenesisState(wasms, "")
+	}
+	accAddress := sdk.AccAddress(proxy).String()
+	return types.NewGenesisState(wasms, accAddress)
 }

x/wasm-storage/keeper/common_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/runtime"
 	"github.com/cosmos/cosmos-sdk/testutil"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
@@ -63,7 +64,7 @@ func setupKeeper(t *testing.T, authority string) (*keeper.Keeper, moduletestutil
 	encCfg := moduletestutil.MakeTestEncodingConfig(wasmstorage.AppModuleBasic{})
 	wasmstoragetypes.RegisterInterfaces(encCfg.InterfaceRegistry)
 
-	wasmStorageKeeper := keeper.NewKeeper(encCfg.Codec, key, authority, nil)
+	wasmStorageKeeper := keeper.NewKeeper(encCfg.Codec, runtime.NewKVStoreService(key), authority, nil)
 
 	return wasmStorageKeeper, encCfg, ctx
 }