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),

custom/staking/keeper/keeper.go

@@ -14,7 +14,7 @@ type Keeper struct {
 	cdc               codec.BinaryCodec
 	acck              accountkeeper.AccountKeeper
 	mintkeeper        *mintkeeper.Keeper
-	stakingmiddleware *stakingmiddleware.Keeper
+	Stakingmiddleware *stakingmiddleware.Keeper
 	authority         string
 }
 
@@ -46,7 +46,7 @@ func NewKeeper(
 		acck:              acck,
 		authority:         authority,
 		mintkeeper:        mintkeeper,
-		stakingmiddleware: stakingmiddleware,
+		Stakingmiddleware: stakingmiddleware,
 		cdc:               cdc,
 	}
 	return keeper

custom/staking/keeper/msg_server.go

@@ -3,8 +3,8 @@
 import (
 	"context"
 
-	"cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 )
@@ -29,23 +29,26 @@
 }
 
 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 {
+	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,
+		)
+	}
+
+	// k.mintkeeper.SetLastTotalPower(ctx, math.Int{})
+	// k.stakingmiddleware.SetLastTotalPower(ctx, math.Int{})
+
+	// delegations := k.Stakingmiddleware.DequeueAllDelegation(ctx)
+	// if len(delegations) > 2 {
 	// 	return nil, sdkerrors.Wrapf(
-	// 		sdkerrors.ErrInvalidRequest, "invalid coin denomination: got %s, expected %s", msg.Amount.Denom, bondDenom,
+	// 		sdkerrors.ErrInvalidRequest, "should always be less then X : got %s, expected %s", len(delegations), 1,
 	// 	)
 	// }
 
-	// 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{})
+	k.Stakingmiddleware.SetDelegation(ctx, msg.DelegatorAddress, msg.ValidatorAddress, msg.Amount.Denom, msg.Amount.Amount)
 
 	return &types.MsgDelegateResponse{}, nil
 	// return nil, fmt.Errorf("My custom error: Nikita")

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,60 @@
 // 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) BeginBlock(ctx sdk.Context, _abc abcitype.RequestBeginBlock) {
+// 	//Define the logic around the batching.
+// 	am.AppModule.BeginBlock(ctx, _abc)
+// }
+
+func (am AppModule) EndBlock(ctx sdk.Context, _abc abcitype.RequestEndBlock) []abcitype.ValidatorUpdate {
+	//Define the logic around the batching.
+	//TODO!!!
+
+	println("EndBlock Custom Staking Module")
+
+	delegations := am.keeper.Stakingmiddleware.DequeueAllDelegation(ctx)
+	println("Delegations: ", delegations)
+	println("Delegations len: ", len(delegations))
+	//for delegations print the delegator address and the validator address
+	for _, delegation := range delegations {
+		println("Delegator Address: ", delegation.DelegatorAddress)
+		println("Validator Address: ", delegation.ValidatorAddress)
+		fmt.Println("Amount", delegation.Amount.Amount)
+
+		msgDelegate := stakingtypes.MsgDelegate{DelegatorAddress: delegation.DelegatorAddress, ValidatorAddress: delegation.ValidatorAddress, Amount: delegation.Amount}
+		_, err := am.msgServer.Delegate(ctx, &msgDelegate)
+		if err != nil {
+			println("Error for Delegator Address: ", delegation.DelegatorAddress)
+		}
+
+	}
+
+	ctx.EventManager().EmitEvent(
+		sdk.NewEvent(
+			"DequeueAllDelegation",
+			// sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
+			// sdk.NewAttribute(types.AttributeKeyValidator, dvPair.ValidatorAddress),
+			// sdk.NewAttribute(types.AttributeKeyDelegator, dvPair.DelegatorAddress),
+		),
+	)
+
+	return am.AppModule.EndBlock(ctx, _abc)
+}

proto/centauri/stakingmiddleware/v1beta1/genesis.proto

@@ -2,6 +2,8 @@ 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";
@@ -14,4 +16,7 @@ message GenesisState {
     (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
     (gogoproto.nullable)   = false
   ];
+
+  // delegations defines the delegations active at genesis.
+  repeated Delegation delegations = 2 [(gogoproto.nullable) = false, (amino.dont_omitempty) = true];  
 }

proto/centauri/stakingmiddleware/v1beta1/stakingmiddleware.proto

@@ -0,0 +1,26 @@
+syntax = "proto3";
+package centauri.stakingmiddleware.v1beta1;
+
+option go_package = "x/stakingmiddleware/types";
+
+import "gogoproto/gogo.proto";
+import "cosmos_proto/cosmos.proto";
+import "cosmos/base/v1beta1/coin.proto";
+
+import "amino/amino.proto";
+import "cosmos/msg/v1/msg.proto";
+
+
+// MsgDelegate defines a SDK message for performing a delegation of coins
+// from a delegator to a validator.
+message Delegation {
+  option (cosmos.msg.v1.signer) = "delegator_address";
+  option (amino.name)           = "cosmos-sdk/MsgDelegate";
+
+  option (gogoproto.equal)           = false;
+  option (gogoproto.goproto_getters) = false;
+
+  string                   delegator_address = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+  string                   validator_address = 2 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+  cosmos.base.v1beta1.Coin amount            = 3 [(gogoproto.nullable) = false, (amino.dont_omitempty) = true];
+}

x/stakingmiddleware/keeper/keeper.go

@@ -1,14 +1,14 @@
 package keeper
 
 import (
-	"cosmossdk.io/math"
 	"github.com/cometbft/cometbft/libs/log"
-	"github.com/notional-labs/composable/v6/x/mint/types"
+	"github.com/notional-labs/composable/v6/x/stakingmiddleware/types"
 
 	"github.com/cosmos/cosmos-sdk/codec"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	// "github.com/notional-labs/composable/v6/x/mint/types"
+
+	sdkmath "cosmossdk.io/math"
 )
 
 // Keeper of the staking middleware store
@@ -80,22 +80,90 @@
 // }
 
 // SetLastTotalPower Set the last total validator power.
-func (k Keeper) SetLastTotalPower(ctx sdk.Context, power math.Int) {
+func (k Keeper) SetLastTotalPower(ctx sdk.Context, power sdkmath.Int) {
 	store := ctx.KVStore(k.storeKey)
 	bz := k.cdc.MustMarshal(&sdk.IntProto{Int: power})
 	store.Set(types.DelegationKey, bz)
 }
 
-func (k Keeper) GetLastTotalPower(ctx sdk.Context) math.Int {
+func (k Keeper) GetLastTotalPower(ctx sdk.Context) sdkmath.Int {
 	store := ctx.KVStore(k.storeKey)
 	bz := store.Get(types.DelegationKey)
 
 	if bz == nil {
-		return math.ZeroInt()
+		return sdkmath.ZeroInt()
 	}
 
 	ip := sdk.IntProto{}
 	k.cdc.MustUnmarshal(bz, &ip)
 
 	return ip.Int
 }
+
+func (k Keeper) SetDelegation(ctx sdk.Context, sourceDelegatorAddress, validatorAddress, denom string, amount sdkmath.Int) {
+	delegation := types.Delegation{DelegatorAddress: sourceDelegatorAddress, ValidatorAddress: validatorAddress, Amount: sdk.NewCoin(denom, amount)}
+	delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+
+	store := ctx.KVStore(k.storeKey)
+	b := k.cdc.MustMarshal(&delegation)
+	store.Set(types.GetDelegationKey(delegatorAddress, GetValidatorAddr(delegation)), b)
+}
+
+func (k Keeper) IterateDelegations(ctx sdk.Context, fn func(index int64, ubd types.Delegation) (stop bool)) {
+	store := ctx.KVStore(k.storeKey)
+
+	iterator := sdk.KVStorePrefixIterator(store, types.DelegationKey)
+	defer iterator.Close()
+
+	for i := int64(0); iterator.Valid(); iterator.Next() {
+		ubd := MustUnmarshalUBD(k.cdc, iterator.Value())
+		if stop := fn(i, ubd); stop {
+			break
+		}
+		i++
+	}
+}
+
+// DequeueAllMatureUBDQueue returns a concatenated list of all the timeslices inclusively previous to
+// currTime, and deletes the timeslices from the queue.
+func (k Keeper) DequeueAllDelegation(ctx sdk.Context) (delegations []types.Delegation) {
+	store := ctx.KVStore(k.storeKey)
+
+	// gets an iterator for all timeslices from time 0 until the current Blockheader time
+	delegationIterator := sdk.KVStorePrefixIterator(store, types.DelegationKey)
+	defer delegationIterator.Close()
+
+	for ; delegationIterator.Valid(); delegationIterator.Next() {
+		delegation := types.Delegation{}
+		value := delegationIterator.Value()
+		k.cdc.MustUnmarshal(value, &delegation)
+
+		delegations = append(delegations, delegation)
+
+		store.Delete(delegationIterator.Key())
+	}
+
+	return delegations
+}
+
+func GetValidatorAddr(d types.Delegation) sdk.ValAddress {
+	addr, err := sdk.ValAddressFromBech32(d.ValidatorAddress)
+	if err != nil {
+		panic(err)
+	}
+	return addr
+}
+
+func UnmarshalBD(cdc codec.BinaryCodec, value []byte) (ubd types.Delegation, err error) {
+	err = cdc.Unmarshal(value, &ubd)
+	return ubd, err
+}
+
+func MustUnmarshalUBD(cdc codec.BinaryCodec, value []byte) types.Delegation {
+	ubd, err := UnmarshalBD(cdc, value)
+	if err != nil {
+		panic(err)
+	}
+
+	return ubd
+}

x/stakingmiddleware/keeper/msg_server.go

@@ -22,6 +22,5 @@ func NewMsgServerImpl(k Keeper) types.MsgServer {
 
 // UpdateParams updates the params.
 func (ms msgServer) SetPower(goCtx context.Context, req *types.MsgSetPower) (*types.MsgSetPowerResponse, error) {
-
 	return &types.MsgSetPowerResponse{}, nil
 }

x/stakingmiddleware/module.go

@@ -20,7 +20,6 @@ import (
 	"github.com/notional-labs/composable/v6/x/stakingmiddleware/client/cli"
 	"github.com/notional-labs/composable/v6/x/stakingmiddleware/keeper"
 	"github.com/notional-labs/composable/v6/x/stakingmiddleware/types"
-	middlewarestaking "github.com/notional-labs/composable/v6/x/stakingmiddleware/types"
 )
 
 var (
@@ -38,7 +37,7 @@ var _ module.AppModuleBasic = AppModuleBasic{}
 
 // Name returns the mint module's name.
 func (AppModuleBasic) Name() string {
-	return middlewarestaking.ModuleName
+	return types.ModuleName
 }
 
 // RegisterLegacyAminoCodec registers the mint module's types on the given LegacyAmino codec.