fix: ensure state compatibility with v1 (#9)

keeper/keeper.go

@@ -46,8 +46,8 @@ type Keeper struct {
 	EntitlementsOwner  collections.Item[string]
 	Paused             collections.Item[bool]
 	PublicCapabilities collections.Map[string, bool]
-	RoleCapabilities   collections.Map[collections.Pair[string, uint64], bool]
-	UserRoles          collections.Map[collections.Pair[[]byte, uint64], bool]
+	RoleCapabilities   collections.Map[[]byte, bool]
+	UserRoles          collections.Map[[]byte, bool]
 
 	addressCodec      address.Codec
 	accountKeeper     types.AccountKeeper
@@ -89,8 +89,8 @@ func NewKeeper(
 		EntitlementsOwner:  collections.NewItem(builder, entitlements.OwnerKey, "entitlements_owner", collections.StringValue),
 		Paused:             collections.NewItem(builder, entitlements.PausedKey, "entitlements_paused", collections.BoolValue),
 		PublicCapabilities: collections.NewMap(builder, entitlements.PublicPrefix, "entitlements_public_capabilities", collections.StringKey, collections.BoolValue),
-		RoleCapabilities:   collections.NewMap(builder, entitlements.CapabilityPrefix, "entitlements_role_capabilities", collections.PairKeyCodec(collections.StringKey, collections.Uint64Key), collections.BoolValue),
-		UserRoles:          collections.NewMap(builder, entitlements.UserPrefix, "entitlements_user_roles", collections.PairKeyCodec(collections.BytesKey, collections.Uint64Key), collections.BoolValue),
+		RoleCapabilities:   collections.NewMap(builder, entitlements.CapabilityPrefix, "entitlements_role_capabilities", collections.BytesKey, collections.BoolValue),
+		UserRoles:          collections.NewMap(builder, entitlements.UserPrefix, "entitlements_user_roles", collections.BytesKey, collections.BoolValue),
 
 		accountKeeper:     accountKeeper,
 		bankKeeper:        bankKeeper,

keeper/msg_server_entitlements_test.go

@@ -319,7 +319,7 @@ func TestEntitlementsUserRoles(t *testing.T) {
 	tmp := k.UserRoles
 	k.UserRoles = collections.NewMap(
 		collections.NewSchemaBuilder(mocks.FailingStore(mocks.Set, utils.GetKVStore(ctx, types.ModuleName))),
-		entitlements.UserPrefix, "entitlements_user_roles", collections.PairKeyCodec(collections.BytesKey, collections.Uint64Key), collections.BoolValue,
+		entitlements.UserPrefix, "entitlements_user_roles", collections.BytesKey, collections.BoolValue,
 	)
 
 	// ACT: Attempt set user role with failing UserRoles collection store.
@@ -464,7 +464,7 @@ func TestEntitlementsUserCapability(t *testing.T) {
 	tmpRole := k.RoleCapabilities
 	k.RoleCapabilities = collections.NewMap(
 		collections.NewSchemaBuilder(mocks.FailingStore(mocks.Set, utils.GetKVStore(ctx, types.ModuleName))),
-		entitlements.CapabilityPrefix, "entitlements_role_capabilities", collections.PairKeyCodec(collections.StringKey, collections.Uint64Key), collections.BoolValue,
+		entitlements.CapabilityPrefix, "entitlements_role_capabilities", collections.BytesKey, collections.BoolValue,
 	)
 
 	// ACT: Attempt set role capability with failing RoleCapabilities collection store.

keeper/state_entitlements.go

@@ -8,6 +8,7 @@ package keeper
 
 import (
 	"context"
+	"encoding/binary"
 	"slices"
 
 	"cosmossdk.io/collections"
@@ -80,24 +81,30 @@ func (k *Keeper) SetPublicCapability(ctx context.Context, method string, enabled
 func (k *Keeper) GetCapabilityRoles(ctx context.Context, method string) []entitlements.Role {
 	var roles []entitlements.Role
 
-	_ = k.RoleCapabilities.Walk(ctx, collections.NewPrefixedPairRange[string, uint64](method), func(key collections.Pair[string, uint64], enabled bool) (stop bool, err error) {
+	itr, _ := k.RoleCapabilities.Iterate(ctx, new(collections.Range[[]byte]).Prefix([]byte(method)))
+
+	defer itr.Close()
+
+	for ; itr.Valid(); itr.Next() {
+		key, _ := itr.Key()
+		enabled, _ := itr.Value()
+
 		if enabled {
-			roles = append(roles, entitlements.Role(key.K2()))
+			role := binary.BigEndian.Uint64(key[len(key)-8:])
+			roles = append(roles, entitlements.Role(role))
 		}
-
-		return false, nil
-	})
+	}
 
 	return roles
 }
 
 func (k *Keeper) GetAllCapabilityRoles(ctx context.Context) []entitlements.RoleCapability {
 	var capabilityRoles []entitlements.RoleCapability
 
-	_ = k.RoleCapabilities.Walk(ctx, nil, func(key collections.Pair[string, uint64], enabled bool) (stop bool, err error) {
+	_ = k.RoleCapabilities.Walk(ctx, nil, func(key []byte, enabled bool) (stop bool, err error) {
 		capabilityRoles = append(capabilityRoles, entitlements.RoleCapability{
-			Method:  key.K1(),
-			Role:    entitlements.Role(key.K2()),
+			Method:  string(key[:len(key)-8]),
+			Role:    entitlements.Role(binary.BigEndian.Uint64(key[len(key)-8:])),
 			Enabled: enabled,
 		})
 
@@ -108,33 +115,39 @@ func (k *Keeper) GetAllCapabilityRoles(ctx context.Context) []entitlements.RoleC
 }
 
 func (k *Keeper) SetRoleCapability(ctx context.Context, method string, role entitlements.Role, enabled bool) error {
-	return k.RoleCapabilities.Set(ctx, collections.Join(method, uint64(role)), enabled)
+	return k.RoleCapabilities.Set(ctx, entitlements.CapabilityRoleKey(method, role), enabled)
 }
 
 //
 
-func (k *Keeper) GetUserRoles(ctx context.Context, user []byte) []entitlements.Role {
+func (k *Keeper) GetUserRoles(ctx context.Context, address []byte) []entitlements.Role {
 	var roles []entitlements.Role
 
-	_ = k.UserRoles.Walk(ctx, collections.NewPrefixedPairRange[[]byte, uint64](user), func(key collections.Pair[[]byte, uint64], enabled bool) (stop bool, err error) {
+	itr, _ := k.UserRoles.Iterate(ctx, new(collections.Range[[]byte]).Prefix(address))
+
+	defer itr.Close()
+
+	for ; itr.Valid(); itr.Next() {
+		key, _ := itr.Key()
+		enabled, _ := itr.Value()
+
 		if enabled {
-			roles = append(roles, entitlements.Role(key.K2()))
+			role := binary.BigEndian.Uint64(key[len(key)-8:])
+			roles = append(roles, entitlements.Role(role))
 		}
-
-		return false, nil
-	})
+	}
 
 	return roles
 }
 
 func (k *Keeper) GetAllUserRoles(ctx context.Context) []entitlements.UserRole {
 	var userRoles []entitlements.UserRole
 
-	_ = k.UserRoles.Walk(ctx, nil, func(key collections.Pair[[]byte, uint64], enabled bool) (stop bool, err error) {
-		address, _ := k.addressCodec.BytesToString(key.K1())
+	_ = k.UserRoles.Walk(ctx, nil, func(key []byte, enabled bool) (stop bool, err error) {
+		address, _ := k.addressCodec.BytesToString(key[:len(key)-8])
 		userRoles = append(userRoles, entitlements.UserRole{
 			User:    address,
-			Role:    entitlements.Role(key.K2()),
+			Role:    entitlements.Role(binary.BigEndian.Uint64(key[len(key)-8:])),
 			Enabled: enabled,
 		})
 
@@ -145,10 +158,10 @@ func (k *Keeper) GetAllUserRoles(ctx context.Context) []entitlements.UserRole {
 }
 
 func (k *Keeper) HasRole(ctx context.Context, address []byte, role entitlements.Role) bool {
-	enabled, _ := k.UserRoles.Get(ctx, collections.Join(address, uint64(role)))
+	enabled, _ := k.UserRoles.Get(ctx, entitlements.UserRoleKey(address, role))
 	return enabled
 }
 
 func (k *Keeper) SetUserRole(ctx context.Context, address []byte, role entitlements.Role, enabled bool) error {
-	return k.UserRoles.Set(ctx, collections.Join(address, uint64(role)), enabled)
+	return k.UserRoles.Set(ctx, entitlements.UserRoleKey(address, role), enabled)
 }

types/aggregator/keys.go

@@ -6,19 +6,9 @@
 
 package aggregator
 
-import "encoding/binary"
-
-const SubmoduleName = "halo-aggregator"
-
 var (
 	OwnerKey       = []byte("aggregator/owner")
 	LastRoundIDKey = []byte("aggregator/last_round_id")
 	NextPriceKey   = []byte("aggregator/next_price")
 	RoundPrefix    = []byte("aggregator/round/")
 )
-
-func RoundKey(id uint64) []byte {
-	bz := make([]byte, 8)
-	binary.BigEndian.PutUint64(bz, id)
-	return append(RoundPrefix, bz...)
-}

types/entitlements/keys.go

@@ -8,8 +8,6 @@ package entitlements
 
 import "encoding/binary"
 
-const SubmoduleName = "halo-entitlements"
-
 var (
 	OwnerKey         = []byte("entitlements/owner")
 	PausedKey        = []byte("entitlements/paused")
@@ -18,26 +16,14 @@ var (
 	UserPrefix       = []byte("entitlements/user/")
 )
 
-func PublicKey(method string) []byte {
-	return append(PublicPrefix, []byte(method)...)
-}
-
-func CapabilityKey(method string) []byte {
-	return append(CapabilityPrefix, []byte(method)...)
-}
-
 func CapabilityRoleKey(method string, role Role) []byte {
 	bz := make([]byte, 8)
 	binary.BigEndian.PutUint64(bz, uint64(role))
-	return append(CapabilityKey(method), bz...)
-}
-
-func UserKey(address []byte) []byte {
-	return append(UserPrefix, address...)
+	return append([]byte(method), bz...)
 }
 
 func UserRoleKey(address []byte, role Role) []byte {
 	bz := make([]byte, 8)
 	binary.BigEndian.PutUint64(bz, uint64(role))
-	return append(UserKey(address), bz...)
+	return append(address, bz...)
 }

types/keys.go

@@ -16,7 +16,3 @@ var (
 	OwnerKey    = []byte("owner")
 	NoncePrefix = []byte("nonce/")
 )
-
-func NonceKey(address []byte) []byte {
-	return append(NoncePrefix, address...)
-}