refactor: optimize delegations by target queries (#33)

## Description

This PR optimizes `PoolDelegations`, `OperatorDelegations` and
`ServiceDelegations` gRPC queries by utilizing store indexes.

<!-- Add a description of the changes that this PR introduces and the
files that
are the most critical to review. -->

---

### Author Checklist

*All items are required. Please add a note to the item if the item is
not applicable and
please add links to any relevant follow up issues.*

I have...

- [ ] included the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] added `!` to the type prefix if API or client breaking change
- [ ] targeted the correct branch (see [PR
Targeting](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#pr-targeting))
- [ ] provided a link to the relevant issue or specification
- [ ] followed the guidelines for [building
modules](https://docs.cosmos.network/v0.44/building-modules/intro.html)
- [ ] included the necessary unit and integration
[tests](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#testing)
- [ ] added a changelog entry to `CHANGELOG.md`
- [ ] included comments for [documenting Go
code](https://blog.golang.org/godoc)
- [ ] updated the relevant documentation or specification
- [ ] reviewed "Files changed" and left comments if necessary
- [ ] confirmed all CI checks have passed

### Reviewers Checklist

*All items are required. Please add a note if the item is not applicable
and please add
your handle next to the items reviewed if you only reviewed selected
items.*

I have...

- [ ] confirmed the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] confirmed `!` in the type prefix if API or client breaking change
- [ ] confirmed all author checklist items have been addressed
- [ ] reviewed state machine logic
- [ ] reviewed API design and naming
- [ ] reviewed documentation is accurate
- [ ] reviewed tests and test coverage
- [ ] manually tested (if applicable)

---------

Co-authored-by: Riccardo Montagnin <[email protected]>
(cherry picked from commit cac1c3d)

x/restaking/keeper/grpc_query.go

@@ -172,24 +172,30 @@ func (k Querier) PoolDelegations(ctx context.Context, req *types.QueryPoolDelega
 
 	// Get the pool delegations store
 	store := k.storeService.OpenKVStore(ctx)
-	delegationsStore := prefix.NewStore(runtime.KVStoreAdapter(store), types.PoolDelegationPrefix)
+	keyPrefix := types.DelegationsByPoolIDStorePrefix(req.PoolId)
+	indexStore := prefix.NewStore(runtime.KVStoreAdapter(store), keyPrefix)
 
 	// Query the pool delegations for the given pool id
-	delegations, pageRes, err := query.GenericFilteredPaginate(k.cdc, delegationsStore, req.Pagination, func(key []byte, delegation *types.Delegation) (*types.Delegation, error) {
-		if delegation.TargetID != req.PoolId {
-			return nil, nil
+	var delegations []types.Delegation
+	pageRes, err := query.Paginate(indexStore, req.Pagination, func(key, _ []byte) error {
+		_, userAddress := types.ParseDelegationByPoolIDStoreKey(append(keyPrefix, key...))
+		delegation, found, err := k.GetPoolDelegation(ctx, req.PoolId, userAddress)
+		if err != nil {
+			return err
 		}
-		return delegation, nil
-	}, func() *types.Delegation {
-		return &types.Delegation{}
+		if !found {
+			return types.ErrDelegationNotFound
+		}
+		delegations = append(delegations, delegation)
+		return nil
 	})
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
 
 	poolDelegations := make([]types.DelegationResponse, len(delegations))
 	for i, delegation := range delegations {
-		response, err := PoolDelegationToPoolDelegationResponse(ctx, k.Keeper, *delegation)
+		response, err := PoolDelegationToPoolDelegationResponse(ctx, k.Keeper, delegation)
 		if err != nil {
 			return nil, status.Error(codes.Internal, err.Error())
 		}
@@ -311,24 +317,30 @@ func (k Querier) OperatorDelegations(ctx context.Context, req *types.QueryOperat
 
 	// Get the operator delegations store
 	store := k.storeService.OpenKVStore(ctx)
-	delegationsStore := prefix.NewStore(runtime.KVStoreAdapter(store), types.OperatorDelegationPrefix)
+	keyPrefix := types.DelegationsByOperatorIDStorePrefix(req.OperatorId)
+	indexStore := prefix.NewStore(runtime.KVStoreAdapter(store), keyPrefix)
 
-	// Query the operator delegations for the given pool id
-	delegations, pageRes, err := query.GenericFilteredPaginate(k.cdc, delegationsStore, req.Pagination, func(key []byte, delegation *types.Delegation) (*types.Delegation, error) {
-		if delegation.TargetID != req.OperatorId {
-			return nil, nil
+	// Query the operator delegations for the given operator id
+	var delegations []types.Delegation
+	pageRes, err := query.Paginate(indexStore, req.Pagination, func(key, _ []byte) error {
+		_, userAddress := types.ParseDelegationByOperatorIDStoreKey(append(keyPrefix, key...))
+		delegation, found, err := k.GetOperatorDelegation(ctx, req.OperatorId, userAddress)
+		if err != nil {
+			return err
+		}
+		if !found {
+			return types.ErrDelegationNotFound
 		}
-		return delegation, nil
-	}, func() *types.Delegation {
-		return &types.Delegation{}
+		delegations = append(delegations, delegation)
+		return nil
 	})
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
 
 	operatorDelegations := make([]types.DelegationResponse, len(delegations))
 	for i, delegation := range delegations {
-		response, err := OperatorDelegationToOperatorDelegationResponse(ctx, k.Keeper, *delegation)
+		response, err := OperatorDelegationToOperatorDelegationResponse(ctx, k.Keeper, delegation)
 		if err != nil {
 			return nil, status.Error(codes.Internal, err.Error())
 		}
@@ -450,24 +462,29 @@ func (k Querier) ServiceDelegations(ctx context.Context, req *types.QueryService
 
 	// Get the service delegations store
 	store := k.storeService.OpenKVStore(ctx)
-	delegationsStore := prefix.NewStore(runtime.KVStoreAdapter(store), types.ServiceDelegationPrefix)
+	keyPrefix := types.DelegationsByServiceIDStorePrefix(req.ServiceId)
+	indexStore := prefix.NewStore(runtime.KVStoreAdapter(store), keyPrefix)
 
-	// Query the service delegations for the given pool id
-	delegations, pageRes, err := query.GenericFilteredPaginate(k.cdc, delegationsStore, req.Pagination, func(key []byte, delegation *types.Delegation) (*types.Delegation, error) {
-		if delegation.TargetID != req.ServiceId {
-			return nil, nil
+	// Query the service delegations for the given service id
+	var delegations []types.Delegation
+	pageRes, err := query.Paginate(indexStore, req.Pagination, func(key, _ []byte) error {
+		_, userAddress := types.ParseDelegationByServiceIDStoreKey(append(keyPrefix, key...))
+		delegation, found, err := k.GetServiceDelegation(ctx, req.ServiceId, userAddress)
+		if err != nil {
+			return err
 		}
-		return delegation, nil
-	}, func() *types.Delegation {
-		return &types.Delegation{}
+		if !found {
+			return types.ErrDelegationNotFound
+		}
+		delegations = append(delegations, delegation)
+		return nil
 	})
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
-
 	serviceDelegationResponses := make([]types.DelegationResponse, len(delegations))
 	for i, delegation := range delegations {
-		response, err := ServiceDelegationToServiceDelegationResponse(ctx, k.Keeper, *delegation)
+		response, err := ServiceDelegationToServiceDelegationResponse(ctx, k.Keeper, delegation)
 		if err != nil {
 			return nil, status.Error(codes.Internal, err.Error())
 		}

x/restaking/types/keys.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"encoding/binary"
 	"time"
 
 	"cosmossdk.io/errors"
@@ -31,13 +32,13 @@ var (
 	PoolDelegationsByPoolIDPrefix = []byte{0xa2}
 	PoolUnbondingDelegationPrefix = []byte{0xa3}
 
-	OperatorDelegationPrefix          = []byte{0xb1}
-	OperatorDelegationByOperatorID    = []byte{0xb2}
-	OperatorUnbondingDelegationPrefix = []byte{0xb3}
+	OperatorDelegationPrefix              = []byte{0xb1}
+	OperatorDelegationsByOperatorIDPrefix = []byte{0xb2}
+	OperatorUnbondingDelegationPrefix     = []byte{0xb3}
 
-	ServiceDelegationPrefix            = []byte{0xc1}
-	ServiceDelegationByServiceIDPrefix = []byte{0xc2}
-	ServiceUnbondingDelegationPrefix   = []byte{0xc3}
+	ServiceDelegationPrefix             = []byte{0xc1}
+	ServiceDelegationsByServiceIDPrefix = []byte{0xc2}
+	ServiceUnbondingDelegationPrefix    = []byte{0xc3}
 
 	UnbondingQueueKey = []byte{0xd1}
 
@@ -92,6 +93,13 @@ func DelegationByPoolIDStoreKey(poolID uint32, delegatorAddress string) []byte {
 	return append(DelegationsByPoolIDStorePrefix(poolID), []byte(delegatorAddress)...)
 }
 
+// ParseDelegationByPoolIDStoreKey returns the poolID and delegator address from the given key
+func ParseDelegationByPoolIDStoreKey(key []byte) (poolID uint32, delegatorAddress string) {
+	poolID = binary.BigEndian.Uint32(key[1:5])
+	delegatorAddress = string(key[5:])
+	return
+}
+
 // PoolUnbondingDelegationsStorePrefix returns the prefix used to store all the unbonding delegations to a given pool
 func PoolUnbondingDelegationsStorePrefix(delegatorAddress string) []byte {
 	return append(PoolUnbondingDelegationPrefix, []byte(delegatorAddress)...)
@@ -116,14 +124,21 @@ func UserOperatorDelegationStoreKey(delegator string, operatorID uint32) []byte
 
 // DelegationsByOperatorIDStorePrefix returns the prefix used to store the delegations to a given operator
 func DelegationsByOperatorIDStorePrefix(operatorID uint32) []byte {
-	return append(OperatorDelegationByOperatorID, operatorstypes.GetOperatorIDBytes(operatorID)...)
+	return append(OperatorDelegationsByOperatorIDPrefix, operatorstypes.GetOperatorIDBytes(operatorID)...)
 }
 
 // DelegationByOperatorIDStoreKey returns the key used to store the operator -> user delegation association
 func DelegationByOperatorIDStoreKey(operatorID uint32, delegatorAddress string) []byte {
 	return append(DelegationsByOperatorIDStorePrefix(operatorID), []byte(delegatorAddress)...)
 }
 
+// ParseDelegationByOperatorIDStoreKey returns the operatorID and delegator address from the given key
+func ParseDelegationByOperatorIDStoreKey(key []byte) (operatorID uint32, delegatorAddress string) {
+	operatorID = binary.BigEndian.Uint32(key[1:5])
+	delegatorAddress = string(key[5:])
+	return
+}
+
 // OperatorUnbondingDelegationsStorePrefix returns the prefix used to store all the unbonding delegations to a given pool
 func OperatorUnbondingDelegationsStorePrefix(delegatorAddress string) []byte {
 	return append(OperatorUnbondingDelegationPrefix, []byte(delegatorAddress)...)
@@ -148,14 +163,21 @@ func UserServiceDelegationStoreKey(delegator string, serviceID uint32) []byte {
 
 // DelegationsByServiceIDStorePrefix returns the prefix used to store the delegations to a given service
 func DelegationsByServiceIDStorePrefix(serviceID uint32) []byte {
-	return append(ServiceDelegationByServiceIDPrefix, servicestypes.GetServiceIDBytes(serviceID)...)
+	return append(ServiceDelegationsByServiceIDPrefix, servicestypes.GetServiceIDBytes(serviceID)...)
 }
 
 // DelegationByServiceIDStoreKey returns the key used to store the service -> user delegation association
 func DelegationByServiceIDStoreKey(serviceID uint32, delegatorAddress string) []byte {
 	return append(DelegationsByServiceIDStorePrefix(serviceID), []byte(delegatorAddress)...)
 }
 
+// ParseDelegationByServiceIDStoreKey returns the serviceID and delegator address from the given key
+func ParseDelegationByServiceIDStoreKey(key []byte) (serviceID uint32, delegatorAddress string) {
+	serviceID = binary.BigEndian.Uint32(key[1:5])
+	delegatorAddress = string(key[5:])
+	return
+}
+
 // ServiceUnbondingDelegationsStorePrefix returns the prefix used to store all the unbonding delegations to a given pool
 func ServiceUnbondingDelegationsStorePrefix(delegatorAddress string) []byte {
 	return append(ServiceUnbondingDelegationPrefix, []byte(delegatorAddress)...)
@@ -171,13 +193,10 @@ func GetDelegationKeyBuilders(delegation Delegation) (DelegationKeyBuilder, Dele
 	switch delegation.Type {
 	case DELEGATION_TYPE_POOL:
 		return UserPoolDelegationStoreKey, DelegationByPoolIDStoreKey, nil
-
 	case DELEGATION_TYPE_OPERATOR:
 		return UserOperatorDelegationStoreKey, DelegationByOperatorIDStoreKey, nil
-
 	case DELEGATION_TYPE_SERVICE:
 		return UserServiceDelegationStoreKey, DelegationByServiceIDStoreKey, nil
-
 	default:
 		return nil, nil, errors.Wrapf(ErrInvalidDelegationType, "invalid delegation type: %v", delegation.Type)
 	}