PLT-980: Add validator-indexed delegation store (phase 1)

sei-cosmos/x/staking/keeper/delegation.go

@@ -98,19 +98,27 @@ func (k Keeper) GetDelegatorDelegations(ctx sdk.Context, delegator sdk.AccAddres
 // SetDelegation sets a delegation.
 func (k Keeper) SetDelegation(ctx sdk.Context, delegation types.Delegation) {
 	delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+	valAddr := delegation.GetValidatorAddr()
 
 	store := ctx.KVStore(k.storeKey)
 	b := types.MustMarshalDelegation(k.cdc, delegation)
-	store.Set(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()), b)
+	store.Set(types.GetDelegationKey(delegatorAddress, valAddr), b)
+	if delegationByValIndexActive(ctx) {
+		store.Set(types.GetDelegationByValIndexKey(delegatorAddress, valAddr), []byte{}) // index, store empty bytes
+	}
 }
 
 // RemoveDelegation removes a delegation.
 func (k Keeper) RemoveDelegation(ctx sdk.Context, delegation types.Delegation) {
 	delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+	valAddr := delegation.GetValidatorAddr()
 
-	k.BeforeDelegationRemoved(ctx, delegatorAddress, delegation.GetValidatorAddr())
+	k.BeforeDelegationRemoved(ctx, delegatorAddress, valAddr)
 	store := ctx.KVStore(k.storeKey)
-	store.Delete(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()))
+	store.Delete(types.GetDelegationKey(delegatorAddress, valAddr))
+	if delegationByValIndexActive(ctx) {
+		store.Delete(types.GetDelegationByValIndexKey(delegatorAddress, valAddr))
+	}
 }
 
 // GetUnbondingDelegations returns a given amount of all the delegator unbonding-delegations.

sei-cosmos/x/staking/keeper/delegation_index.go

@@ -0,0 +1,109 @@
+package keeper
+
+import (
+	"time"
+
+	"golang.org/x/mod/semver"
+
+	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
+)
+
+// DelegationByValIndexUpgrade is the upgrade at which SetDelegation and
+// RemoveDelegation begin maintaining the validator-indexed delegation store.
+const DelegationByValIndexUpgrade = "v6.7"
+
+// delegationByValIndexActive reports whether SetDelegation and RemoveDelegation
+// should maintain the validator-indexed delegation store.
+//
+// Live execution always does: the current binary only ever executes at/after
+// the upgrade that ships this behavior, so a non-tracing block is never a
+// pre-upgrade block. The only place pre-upgrade behavior must be reproduced is
+// when re-tracing a historical block, where the era is signaled through
+// ClosestUpgradeName (see app.RPCContextProvider).
+func delegationByValIndexActive(ctx sdk.Context) bool {
+	if !ctx.IsTracing() {
+		return true
+	}
+	return semver.Compare(ctx.ClosestUpgradeName(), DelegationByValIndexUpgrade) >= 0
+}
+
+// BackfillDelegationByValIndexResult reports the outcome of a delegation index backfill.
+type BackfillDelegationByValIndexResult struct {
+	TotalDelegations int
+	IndexWritten     int
+	AlreadyIndexed   int
+	DryRun           bool
+	Elapsed          time.Duration
+}
+
+// BackfillDelegationByValIndexProgress is a progress snapshot during backfill.
+type BackfillDelegationByValIndexProgress struct {
+	TotalDelegations int
+	IndexWritten     int
+	AlreadyIndexed   int
+	Elapsed          time.Duration
+}
+
+// BackfillProgress reports incremental backfill progress. Nil disables callbacks.
+type BackfillProgress func(BackfillDelegationByValIndexProgress)
+
+const (
+	backfillProgressDelegationInterval = 100_000
+	backfillProgressMinInterval        = 10 * time.Second
+)
+
+// BackfillDelegationByValIndex writes validator-indexed delegation keys for existing
+// delegations. When dryRun is true, delegations are counted but no store writes occur.
+func (k Keeper) BackfillDelegationByValIndex(ctx sdk.Context, dryRun bool, progress BackfillProgress) BackfillDelegationByValIndexResult {
+	start := time.Now()
+	result := BackfillDelegationByValIndexResult{DryRun: dryRun}
+	lastProgressReport := start
+
+	reportProgress := func() {
+		if progress == nil {
+			return
+		}
+		progress(BackfillDelegationByValIndexProgress{
+			TotalDelegations: result.TotalDelegations,
+			IndexWritten:     result.IndexWritten,
+			AlreadyIndexed:   result.AlreadyIndexed,
+			Elapsed:          time.Since(start),
+		})
+		lastProgressReport = time.Now()
+	}
+
+	store := ctx.KVStore(k.storeKey)
+	iterator := sdk.KVStorePrefixIterator(store, types.DelegationKey)
+	defer func() { _ = iterator.Close() }()
+
+	for ; iterator.Valid(); iterator.Next() {
+		delegation := types.MustUnmarshalDelegation(k.cdc, iterator.Value())
+		result.TotalDelegations++
+
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		indexKey := types.GetDelegationByValIndexKey(delegatorAddress, valAddr)
+		if store.Has(indexKey) {
+			result.AlreadyIndexed++
+		} else {
+			if !dryRun {
+				store.Set(indexKey, []byte{})
+			}
+			result.IndexWritten++
+		}
+
+		if progress == nil {
+			continue
+		}
+		now := time.Now()
+		if result.TotalDelegations == 1 ||
+			result.TotalDelegations%backfillProgressDelegationInterval == 0 ||
+			now.Sub(lastProgressReport) >= backfillProgressMinInterval {
+			reportProgress()
+		}
+	}
+
+	result.Elapsed = time.Since(start)
+	return result
+}

sei-cosmos/x/staking/keeper/delegation_index_test.go

@@ -0,0 +1,111 @@
+package keeper_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
+)
+
+func TestDelegationByValIndexDualWrite(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+
+	addrDels, valAddrs := generateAddresses(app, ctx, 1)
+	delegation := types.NewDelegation(addrDels[0], valAddrs[0], sdk.NewDec(1))
+
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	indexKey := types.GetDelegationByValIndexKey(addrDels[0], valAddrs[0])
+
+	require.False(t, store.Has(indexKey))
+
+	app.StakingKeeper.SetDelegation(ctx, delegation)
+	require.True(t, store.Has(indexKey))
+	require.NotNil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+
+	app.StakingKeeper.RemoveDelegation(ctx, delegation)
+	require.False(t, store.Has(indexKey))
+	require.Nil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+}
+
+func TestDelegationByValIndexTracingPreUpgradeNoDualWrite(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+	ctx = ctx.WithIsTracing(true).WithClosestUpgradeName("v6.6")
+
+	addrDels, valAddrs := generateAddresses(app, ctx, 1)
+	delegation := types.NewDelegation(addrDels[0], valAddrs[0], sdk.NewDec(1))
+
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	indexKey := types.GetDelegationByValIndexKey(addrDels[0], valAddrs[0])
+
+	app.StakingKeeper.SetDelegation(ctx, delegation)
+	require.False(t, store.Has(indexKey))
+	require.NotNil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+
+	app.StakingKeeper.RemoveDelegation(ctx, delegation)
+	require.False(t, store.Has(indexKey))
+	require.Nil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+}
+
+func TestDelegationByValIndexTracingPostUpgradeDualWrite(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+	ctx = ctx.WithIsTracing(true).WithClosestUpgradeName("v6.7")
+
+	addrDels, valAddrs := generateAddresses(app, ctx, 1)
+	delegation := types.NewDelegation(addrDels[0], valAddrs[0], sdk.NewDec(1))
+
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	indexKey := types.GetDelegationByValIndexKey(addrDels[0], valAddrs[0])
+
+	app.StakingKeeper.SetDelegation(ctx, delegation)
+	require.True(t, store.Has(indexKey))
+}
+
+func TestBackfillDelegationByValIndex(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+
+	addrDels, valAddrs := generateAddresses(app, ctx, 2)
+	delegations := []types.Delegation{
+		types.NewDelegation(addrDels[0], valAddrs[0], sdk.NewDec(1)),
+		types.NewDelegation(addrDels[0], valAddrs[1], sdk.NewDec(2)),
+		types.NewDelegation(addrDels[1], valAddrs[0], sdk.NewDec(3)),
+	}
+
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	for _, delegation := range delegations {
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		store.Set(
+			types.GetDelegationKey(delegatorAddress, valAddr),
+			types.MustMarshalDelegation(app.AppCodec(), delegation),
+		)
+	}
+
+	dryRunResult := app.StakingKeeper.BackfillDelegationByValIndex(ctx, true, nil)
+	require.Equal(t, 3, dryRunResult.TotalDelegations)
+	require.Equal(t, 3, dryRunResult.IndexWritten)
+	require.Equal(t, 0, dryRunResult.AlreadyIndexed)
+	require.True(t, dryRunResult.DryRun)
+	for _, delegation := range delegations {
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		require.False(t, store.Has(types.GetDelegationByValIndexKey(delegatorAddress, valAddr)))
+	}
+
+	writeResult := app.StakingKeeper.BackfillDelegationByValIndex(ctx, false, nil)
+	require.Equal(t, 3, writeResult.TotalDelegations)
+	require.Equal(t, 3, writeResult.IndexWritten)
+	require.Equal(t, 0, writeResult.AlreadyIndexed)
+	require.False(t, writeResult.DryRun)
+	for _, delegation := range delegations {
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		require.True(t, store.Has(types.GetDelegationByValIndexKey(delegatorAddress, valAddr)))
+	}
+
+	repeatResult := app.StakingKeeper.BackfillDelegationByValIndex(ctx, false, nil)
+	require.Equal(t, 3, repeatResult.TotalDelegations)
+	require.Equal(t, 0, repeatResult.IndexWritten)
+	require.Equal(t, 3, repeatResult.AlreadyIndexed)
+}

sei-cosmos/x/staking/simulation/decoder.go

@@ -41,6 +41,8 @@ func NewDecodeStore(cdc codec.Codec) func(kvA, kvB kv.Pair) string {
 			cdc.MustUnmarshal(kvB.Value, &delegationB)
 
 			return fmt.Sprintf("%v\n%v", delegationA, delegationB)
+		case bytes.Equal(kvA.Key[:1], types.DelegationByValIndexKey):
+			return fmt.Sprintf("%X\n%X", kvA.Key, kvB.Key)
 		case bytes.Equal(kvA.Key[:1], types.UnbondingDelegationKey),
 			bytes.Equal(kvA.Key[:1], types.UnbondingDelegationByValIndexKey):
 			var ubdA, ubdB types.UnbondingDelegation

sei-cosmos/x/staking/types/keys.go

@@ -42,6 +42,7 @@ var (
 	RedelegationKey                  = []byte{0x34} // key for a redelegation
 	RedelegationByValSrcIndexKey     = []byte{0x35} // prefix for each key for an redelegation, by source validator operator
 	RedelegationByValDstIndexKey     = []byte{0x36} // prefix for each key for an redelegation, by destination validator operator
+	DelegationByValIndexKey          = []byte{0x37} // prefix for each key for a delegation, by validator operator
 
 	UnbondingQueueKey    = []byte{0x41} // prefix for the timestamps in unbonding queue
 	RedelegationQueueKey = []byte{0x42} // prefix for the timestamps in redelegations queue
@@ -198,6 +199,31 @@ func GetDelegationsKey(delAddr sdk.AccAddress) []byte {
 	return append(DelegationKey, address.MustLengthPrefix(delAddr)...)
 }
 
+// GetDelegationByValIndexKey creates the index-key for a delegation, stored by validator-index.
+// VALUE: none (key rearrangement used)
+func GetDelegationByValIndexKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
+	return append(GetDelegationsByValIndexKey(valAddr), address.MustLengthPrefix(delAddr)...)
+}
+
+// GetDelegationKeyFromValIndexKey rearranges the ValIndexKey to get the DelegationKey.
+func GetDelegationKeyFromValIndexKey(indexKey []byte) []byte {
+	kv.AssertKeyAtLeastLength(indexKey, 2)
+	addrs := indexKey[1:] // remove prefix bytes
+
+	valAddrLen := addrs[0]
+	kv.AssertKeyAtLeastLength(addrs, 2+int(valAddrLen))
+	valAddr := addrs[1 : 1+valAddrLen]
+	kv.AssertKeyAtLeastLength(addrs, 3+int(valAddrLen))
+	delAddr := addrs[valAddrLen+2:]
+
+	return GetDelegationKey(delAddr, valAddr)
+}
+
+// GetDelegationsByValIndexKey creates the prefix keyspace for delegation indexes by validator.
+func GetDelegationsByValIndexKey(valAddr sdk.ValAddress) []byte {
+	return append(DelegationByValIndexKey, address.MustLengthPrefix(valAddr)...)
+}
+
 // GetUBDKey creates the key for an unbonding delegation by delegator and validator addr
 // VALUE: staking/UnbondingDelegation
 func GetUBDKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {

sei-cosmos/x/staking/types/keys_test.go

@@ -49,6 +49,16 @@ func TestGetValidatorPowerRank(t *testing.T) {
 	}
 }
 
+func TestGetDelegationKeyFromValIndexKey(t *testing.T) {
+	delAddr := sdk.AccAddress(keysAddr1)
+	valAddr := sdk.ValAddress(keysAddr2)
+
+	delegationKey := types.GetDelegationKey(delAddr, valAddr)
+	indexKey := types.GetDelegationByValIndexKey(delAddr, valAddr)
+
+	require.Equal(t, delegationKey, types.GetDelegationKeyFromValIndexKey(indexKey))
+}
+
 func TestGetREDByValDstIndexKey(t *testing.T) {
 	tests := []struct {
 		delAddr    sdk.AccAddress

tools/cmd.go

@@ -3,6 +3,7 @@ package tools
 import (
 	"github.com/spf13/cobra"
 
+	stakingcmd "github.com/sei-protocol/sei-chain/tools/staking/cmd"
 	scanner "github.com/sei-protocol/sei-chain/tools/tx-scanner/cmd"
 )
 
@@ -12,5 +13,6 @@ func ToolCmd() *cobra.Command {
 		Short: "A set of useful tools for sei chain",
 	}
 	toolsCmd.AddCommand(scanner.ScanCmd())
+	toolsCmd.AddCommand(stakingcmd.StakingCmd())
 	return toolsCmd
 }