add delegation store to custom staking middleware

app/app.go

@@ -361,7 +361,7 @@ func NewComposableApp(
 		mint.NewAppModule(appCodec, app.MintKeeper, app.AccountKeeper, nil),
 		slashing.NewAppModule(appCodec, app.SlashingKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper, app.GetSubspace(slashingtypes.ModuleName)),
 		distr.NewAppModule(appCodec, app.DistrKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper, app.GetSubspace(distrtypes.ModuleName)),
-		customstaking.NewAppModule(appCodec, *&app.CustomStakingKeeper, app.AccountKeeper, app.BankKeeper, app.GetSubspace(stakingtypes.ModuleName)),
+		customstaking.NewAppModule(appCodec, app.CustomStakingKeeper, app.AccountKeeper, app.BankKeeper, app.GetSubspace(stakingtypes.ModuleName)),
 		stakingmiddleware.NewAppModule(appCodec, app.StakingMiddlewareKeeper),
 		authzmodule.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
 		upgrade.NewAppModule(app.UpgradeKeeper),

app/keepers/keepers.go

@@ -196,7 +196,7 @@ func (appKeepers *AppKeepers) InitNormalKeepers(
 	appKeepers.StakingMiddlewareKeeper = stakingmiddleware.NewKeeper(appCodec, appKeepers.keys[stakingmiddlewaretypes.StoreKey], authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
 	appKeepers.CustomStakingKeeper = customstaking.NewKeeper(
-		appCodec /*appKeepers.keys[stakingtypes.StoreKey],*/, *appKeepers.StakingKeeper, appKeepers.AccountKeeper, &appKeepers.MintKeeper, &appKeepers.StakingMiddlewareKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
+		appCodec /*appKeepers.keys[stakingtypes.StoreKey],*/, *appKeepers.StakingKeeper, appKeepers.AccountKeeper, &appKeepers.StakingMiddlewareKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 
 	appKeepers.DistrKeeper = distrkeeper.NewKeeper(

custom/staking/abci.go

@@ -0,0 +1,22 @@
+package bank
+
+import (
+	"time"
+
+	abci "github.com/cometbft/cometbft/abci/types"
+
+	"github.com/cosmos/cosmos-sdk/telemetry"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	// "github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+
+	customstakingkeeper "github.com/notional-labs/composable/v6/custom/staking/keeper"
+)
+
+// Called every block, update validator set
+func EndBlocker(ctx sdk.Context, k *customstakingkeeper.Keeper) []abci.ValidatorUpdate {
+	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyEndBlocker)
+
+	return k.BlockValidatorUpdates(ctx, ctx.BlockHeight())
+}

custom/staking/keeper/keeper.go

@@ -1,65 +1,125 @@
 package keeper
 
 import (
+	abcicometbft "github.com/cometbft/cometbft/abci/types"
 	"github.com/cosmos/cosmos-sdk/codec"
 	accountkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
 
-	mintkeeper "github.com/notional-labs/composable/v6/x/mint/keeper"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	stakingmiddleware "github.com/notional-labs/composable/v6/x/stakingmiddleware/keeper"
 )
 
 type Keeper struct {
 	stakingkeeper.Keeper
 	cdc               codec.BinaryCodec
 	acck              accountkeeper.AccountKeeper
-	mintkeeper        *mintkeeper.Keeper
-	stakingmiddleware *stakingmiddleware.Keeper
+	Stakingmiddleware *stakingmiddleware.Keeper
 	authority         string
 }
 
-// func NewBaseKeeper(
-// 	cdc codec.BinaryCodec,
-// 	key storetypes.StoreKey,
-// 	ak types.AccountKeeper,
-// 	acck accountkeeper.AccountKeeper,
-// 	bk bankkeeper.Keeper,
-// 	authority string,
-// ) Keeper {
-// 	keeper := Keeper{
-// 		Keeper: *stakingkeeper.NewKeeper(cdc, key, ak, bk, authority),
-// 		acck:   acck,
-// 	}
-// 	return keeper
-// }
+func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, hight int64) []abcicometbft.ValidatorUpdate {
+	// Calculate validator set changes.
+	//
+	// NOTE: ApplyAndReturnValidatorSetUpdates has to come before
+	// UnbondAllMatureValidatorQueue.
+	// This fixes a bug when the unbonding period is instant (is the case in
+	// some of the tests). The test expected the validator to be completely
+	// unbonded after the Endblocker (go from Bonded -> Unbonding during
+	// ApplyAndReturnValidatorSetUpdates and then Unbonding -> Unbonded during
+	// UnbondAllMatureValidatorQueue).
+	println("BlockValidatorUpdates Custom Staking Module")
+	params := k.Stakingmiddleware.GetParams(ctx)
+	println("BlocksPerEpoch: ", params.BlocksPerEpoch)
+	should_execute_batch := (hight % int64(params.BlocksPerEpoch)) == 0
+	var validatorUpdates []abcicometbft.ValidatorUpdate
+	if should_execute_batch {
+		println("Should Execute Batch: ", hight)
+		v, err := k.ApplyAndReturnValidatorSetUpdates(ctx)
+		if err != nil {
+			panic(err)
+		}
+		validatorUpdates = v
+	}
+
+	// unbond all mature validators from the unbonding queue
+	k.UnbondAllMatureValidators(ctx)
+
+	// Remove all mature unbonding delegations from the ubd queue.
+	matureUnbonds := k.DequeueAllMatureUBDQueue(ctx, ctx.BlockHeader().Time)
+	for _, dvPair := range matureUnbonds {
+		addr, err := sdk.ValAddressFromBech32(dvPair.ValidatorAddress)
+		if err != nil {
+			panic(err)
+		}
+		delegatorAddress := sdk.MustAccAddressFromBech32(dvPair.DelegatorAddress)
+
+		balances, err := k.CompleteUnbonding(ctx, delegatorAddress, addr)
+		if err != nil {
+			continue
+		}
+
+		ctx.EventManager().EmitEvent(
+			sdk.NewEvent(
+				types.EventTypeCompleteUnbonding,
+				sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
+				sdk.NewAttribute(types.AttributeKeyValidator, dvPair.ValidatorAddress),
+				sdk.NewAttribute(types.AttributeKeyDelegator, dvPair.DelegatorAddress),
+			),
+		)
+	}
+
+	// Remove all mature redelegations from the red queue.
+	matureRedelegations := k.DequeueAllMatureRedelegationQueue(ctx, ctx.BlockHeader().Time)
+	for _, dvvTriplet := range matureRedelegations {
+		valSrcAddr, err := sdk.ValAddressFromBech32(dvvTriplet.ValidatorSrcAddress)
+		if err != nil {
+			panic(err)
+		}
+		valDstAddr, err := sdk.ValAddressFromBech32(dvvTriplet.ValidatorDstAddress)
+		if err != nil {
+			panic(err)
+		}
+		delegatorAddress := sdk.MustAccAddressFromBech32(dvvTriplet.DelegatorAddress)
+
+		balances, err := k.CompleteRedelegation(
+			ctx,
+			delegatorAddress,
+			valSrcAddr,
+			valDstAddr,
+		)
+		if err != nil {
+			continue
+		}
+
+		ctx.EventManager().EmitEvent(
+			sdk.NewEvent(
+				types.EventTypeCompleteRedelegation,
+				sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
+				sdk.NewAttribute(types.AttributeKeyDelegator, dvvTriplet.DelegatorAddress),
+				sdk.NewAttribute(types.AttributeKeySrcValidator, dvvTriplet.ValidatorSrcAddress),
+				sdk.NewAttribute(types.AttributeKeyDstValidator, dvvTriplet.ValidatorDstAddress),
+			),
+		)
+	}
+
+	return validatorUpdates
+}
 
 func NewKeeper(
 	cdc codec.BinaryCodec,
 	staking stakingkeeper.Keeper,
 	acck accountkeeper.AccountKeeper,
-	mintkeeper *mintkeeper.Keeper,
 	stakingmiddleware *stakingmiddleware.Keeper,
 	authority string,
 ) Keeper {
 	keeper := Keeper{
 		Keeper:            staking,
 		acck:              acck,
 		authority:         authority,
-		mintkeeper:        mintkeeper,
-		stakingmiddleware: stakingmiddleware,
+		Stakingmiddleware: stakingmiddleware,
 		cdc:               cdc,
 	}
 	return keeper
 }
-
-// func (k *Keeper) RegisterKeepers(akk banktypes.StakingKeeper) {
-// 	k.acck = sk
-// }
-
-// func (k Keeper) StoreDelegation(ctx sdk.Context, delegation types.Delegation) {
-// 	delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
-
-// 	store := ctx.KVStore(k.storeKey)
-// 	b := types.MustMarshalDelegation(k.cdc, delegation)
-// 	store.Set(customstakingtypes.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()), b)
-// }

custom/staking/keeper/msg_server.go

@@ -3,8 +3,6 @@ package keeper
 import (
 	"context"
 
-	"cosmossdk.io/math"
-	sdk "github.com/cosmos/cosmos-sdk/types"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 )
@@ -29,27 +27,7 @@ func (k msgServer) EditValidator(goCtx context.Context, msg *types.MsgEditValida
 }
 
 func (k msgServer) Delegate(goCtx context.Context, msg *types.MsgDelegate) (*types.MsgDelegateResponse, error) {
-
-	ctx := sdk.UnwrapSDKContext(goCtx)
-
-	// bondDenom := k.BondDenom(ctx)
-	// if msg.Amount.Denom != bondDenom {
-	// 	return nil, sdkerrors.Wrapf(
-	// 		sdkerrors.ErrInvalidRequest, "invalid coin denomination: got %s, expected %s", msg.Amount.Denom, bondDenom,
-	// 	)
-	// }
-
-	// delegation := types.Delegation{
-	// 	DelegatorAddress: msg.DelegatorAddress,
-	// 	ValidatorAddress: msg.ValidatorAddress,
-	// 	Shares:           msg.Amount.Amount.ToLegacyDec(),
-	// }
-	k.mintkeeper.SetLastTotalPower(ctx, math.Int{})
-	k.stakingmiddleware.SetLastTotalPower(ctx, math.Int{})
-
-	return &types.MsgDelegateResponse{}, nil
-	// return nil, fmt.Errorf("My custom error: Nikita")
-	// return k.msgServer.Delegate(goCtx, msg)
+	return k.msgServer.Delegate(goCtx, msg)
 }
 
 func (k msgServer) BeginRedelegate(goCtx context.Context, msg *types.MsgBeginRedelegate) (*types.MsgBeginRedelegateResponse, error) {

custom/staking/module.go

@@ -3,32 +3,35 @@
 import (
 	"fmt"
 
+	abcitype "github.com/cometbft/cometbft/abci/types"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	stakingmodule "github.com/cosmos/cosmos-sdk/x/staking"
 	"github.com/cosmos/cosmos-sdk/x/staking/exported"
-	"github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
-	"github.com/cosmos/cosmos-sdk/x/staking/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	// custombankkeeper "github.com/notional-labs/composable/v6/custom/bank/keeper"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	customstakingkeeper "github.com/notional-labs/composable/v6/custom/staking/keeper"
 )
 
 // AppModule wraps around the bank module and the bank keeper to return the right total supply
 type AppModule struct {
 	stakingmodule.AppModule
-	keeper   customstakingkeeper.Keeper
-	subspace exported.Subspace
+	keeper    customstakingkeeper.Keeper
+	subspace  exported.Subspace
+	msgServer stakingtypes.MsgServer
 }
 
 // NewAppModule creates a new AppModule object
-func NewAppModule(cdc codec.Codec, keeper customstakingkeeper.Keeper, accountKeeper types.AccountKeeper, bankKeeper types.BankKeeper, ss exported.Subspace) AppModule {
+func NewAppModule(cdc codec.Codec, keeper customstakingkeeper.Keeper, accountKeeper stakingtypes.AccountKeeper, bankKeeper stakingtypes.BankKeeper, ss exported.Subspace) AppModule {
 	stakingModule := stakingmodule.NewAppModule(cdc, &keeper.Keeper, accountKeeper, bankKeeper, ss)
 	return AppModule{
 		AppModule: stakingModule,
 		keeper:    keeper,
 		subspace:  ss,
+		msgServer: stakingkeeper.NewMsgServerImpl(&keeper.Keeper),
 	}
 }
 
@@ -37,20 +40,24 @@
 // when trying to force this custom keeper into a bankkeeper.BaseKeeper
 func (am AppModule) RegisterServices(cfg module.Configurator) {
 	// types.RegisterMsgServer(cfg.MsgServer(), keeper.NewMsgServerImpl(&am.keeper))
-	types.RegisterMsgServer(cfg.MsgServer(), customstakingkeeper.NewMsgServerImpl(am.keeper.Keeper, am.keeper))
-	querier := keeper.Querier{Keeper: &am.keeper.Keeper}
-	types.RegisterQueryServer(cfg.QueryServer(), querier)
+	stakingtypes.RegisterMsgServer(cfg.MsgServer(), customstakingkeeper.NewMsgServerImpl(am.keeper.Keeper, am.keeper))
+	querier := stakingkeeper.Querier{Keeper: &am.keeper.Keeper}
+	stakingtypes.RegisterQueryServer(cfg.QueryServer(), querier)
 
 	m := stakingkeeper.NewMigrator(&am.keeper.Keeper, am.subspace)
-	if err := cfg.RegisterMigration(types.ModuleName, 1, m.Migrate1to2); err != nil {
+	if err := cfg.RegisterMigration(stakingtypes.ModuleName, 1, m.Migrate1to2); err != nil {
 		panic(fmt.Sprintf("failed to migrate x/staking from version 1 to 2: %v", err))
 	}
 
-	if err := cfg.RegisterMigration(types.ModuleName, 2, m.Migrate2to3); err != nil {
+	if err := cfg.RegisterMigration(stakingtypes.ModuleName, 2, m.Migrate2to3); err != nil {
 		panic(fmt.Sprintf("failed to migrate x/staking from version 2 to 3: %v", err))
 	}
 
-	if err := cfg.RegisterMigration(types.ModuleName, 3, m.Migrate3to4); err != nil {
+	if err := cfg.RegisterMigration(stakingtypes.ModuleName, 3, m.Migrate3to4); err != nil {
 		panic(fmt.Sprintf("failed to migrate x/staking from version 3 to 4: %v", err))
 	}
 }
+
+func (am AppModule) EndBlock(ctx sdk.Context, _abc abcitype.RequestEndBlock) []abcitype.ValidatorUpdate {
+	return EndBlocker(ctx, &am.keeper)
+}

proto/centauri/mint/v1beta1/genesis.proto

@@ -19,11 +19,4 @@ message GenesisState {
     (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coin",
     (gogoproto.nullable) = false
   ];
-
-  // last_total_power tracks the total amounts of bonded tokens recorded during
-  // the previous end block.
-  bytes last_total_power = 4 [
-    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
-    (gogoproto.nullable)   = false
-  ];
 }

proto/centauri/stakingmiddleware/v1beta1/genesis.proto

@@ -2,16 +2,13 @@ syntax = "proto3";
 package centauri.stakingmiddleware.v1beta1;
 
 import "gogoproto/gogo.proto";
+import "centauri/stakingmiddleware/v1beta1/stakingmiddleware.proto";
+import "amino/amino.proto";
 
 
 option go_package = "x/stakingmiddleware/types";
 
 // GenesisState defines the stakingmiddleware module's genesis state.
 message GenesisState {
-  // last_total_power tracks the total amounts of bonded tokens recorded during
-  // the previous end block.
-  bytes last_total_power = 1 [
-    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
-    (gogoproto.nullable)   = false
-  ];
+  Params params = 1 [ (gogoproto.nullable) = false ];
 }

proto/centauri/stakingmiddleware/v1beta1/query.proto

@@ -3,20 +3,23 @@ package centauri.stakingmiddleware.v1beta1;
 
 import "gogoproto/gogo.proto";
 import "google/api/annotations.proto";
+import "centauri/stakingmiddleware/v1beta1/stakingmiddleware.proto";
 
 option go_package = "x/stakingmiddleware/types";
 
 // Query provides defines the gRPC querier service.
 service Query {
   // Params returns the total set of minting parameters.
-  rpc Power(QueryPowerRequest) returns (QueryPowerResponse) {
-    option (google.api.http).get = "/cosmos/stakingmiddleware/v1beta1/params";
+  rpc Params(QueryParamsRequest) returns (QueryParamsResponse) {
+    option (google.api.http).get = "/cosmos/mint/v1beta1/params";
   }
 }
 
 // QueryParamsRequest is the request type for the Query/Params RPC method.
-message QueryPowerRequest {}
+message QueryParamsRequest {}
 
 // QueryParamsResponse is the response type for the Query/Params RPC method.
-message QueryPowerResponse {
+message QueryParamsResponse {
+  // params defines the parameters of the module.
+  Params params = 1 [ (gogoproto.nullable) = false ];
 }