refine some RPCs regarding asset and delegation

precompiles/assets/query.go

@@ -167,11 +167,11 @@ func (p Precompile) GetStakerBalanceByToken(
 		ClientChainID:      clientChainID,
 		StakerAddress:      stakerAddress,
 		TokenID:            assetAddress,
-		Balance:            balance.Balance,
-		Withdrawable:       balance.Withdrawable,
-		Delegated:          balance.Delegated,
-		PendingUndelegated: balance.PendingUndelegated,
-		TotalDeposited:     balance.TotalDeposited,
+		Balance:            balance.Balance.BigInt(),
+		Withdrawable:       balance.Withdrawable.BigInt(),
+		Delegated:          balance.Delegated.BigInt(),
+		PendingUndelegated: balance.PendingUndelegated.BigInt(),
+		TotalDeposited:     balance.TotalDeposited.BigInt(),
 	}
 
 	return method.Outputs.Pack(true, result)

precompiles/avs/query_test.go

@@ -61,8 +61,7 @@ var baseTestCases = []avsTestCases{
 
 func (suite *AVSManagerPrecompileSuite) TestGetOptedInOperatorAccAddrs() {
 	method := suite.precompile.Methods[avsManagerPrecompile.MethodGetOptInOperators]
-	operatorAddress, avsAddress, slashContract :=
-		sdk.AccAddress(utiltx.GenerateAddress().Bytes()).String(), suite.Address, "0xDF907c29719154eb9872f021d21CAE6E5025d7aB"
+	operatorAddress, avsAddress, slashContract := sdk.AccAddress(utiltx.GenerateAddress().Bytes()).String(), suite.Address, "0xDF907c29719154eb9872f021d21CAE6E5025d7aB"
 
 	operatorOptIn := func() {
 		optedInfo := &types.OptedInfo{

proto/imuachain/assets/v1/query.proto

@@ -99,6 +99,62 @@ message QueryParamsResponse {
   Params params = 1;
 }
 
+// StakerBalance is a struct to describe the balance of a staker for a specific asset
+// balance = withdrawable + delegated + pendingUndelegated
+// pendingUndelegated is the amount of the asset that is during unbonding period and not yet withdrawable
+// it would finally be withdrawable after the unbonding period, but the final amount may be less than the pendingUndelegated
+// because of the penalty during the unbonding period
+message StakerBalance {
+  // staker_id is the staker id for which the query is made.
+  string staker_id = 1;
+  // asset_id is the asset for which the query is made.
+  string asset_id = 2;
+  // balance = withdrawable + delegated + pendingUndelegated
+  string balance = 3 [
+    (cosmos_proto.scalar) = "cosmos.Int",
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.nullable) = false
+  ];
+  // withdrawable is the amount that can be withdrawn
+  string withdrawable = 4 [
+    (cosmos_proto.scalar) = "cosmos.Int",
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.nullable) = false
+  ];
+  // delegated is the amount that has been delegated to the operators
+  string delegated = 5 [
+    (cosmos_proto.scalar) = "cosmos.Int",
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.nullable) = false
+  ];
+  // pending_undelegated is the amount that has been undelegated but not been completed
+  string pending_undelegated = 6 [
+    (cosmos_proto.scalar) = "cosmos.Int",
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.nullable) = false
+  ];
+  // total_deposited is the total deposit amount, compared to the balance, it might include slashed amount
+  string total_deposited = 7 [
+    (cosmos_proto.scalar) = "cosmos.Int",
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.nullable) = false
+  ];
+}
+
+// QueryStakerBalanceRequest is the request type for the Query/QueryStakerBalance RPC method.
+message QueryStakerBalanceRequest {
+  // staker_id is the staker id for which the query is made.
+  string staker_id = 1;
+  // asset_id is the asset for which the query is made.
+  string asset_id = 2;
+}
+
+// QueryStakerBalanceResponse is the response type for the Query/QueryStakerBalance RPC method.
+message QueryStakerBalanceResponse {
+  // staker_balance is the response for QueryStakerBalanceRequest
+  StakerBalance staker_balance = 1;
+}
+
 // Query defines the gRPC query service for the assets module.
 service Query {
   // Params retrieves the assets module params
@@ -146,4 +202,9 @@ service Query {
     option (cosmos.query.v1.module_query_safe) = true;
     option (google.api.http).get = "/imuachain/assets/v1/operator_asset/{operator_addr}/{asset_id}";
   }
+  // QueryStakerBalance queries the staker balance for specified asset
+  rpc QueryStakerBalance(QueryStakerBalanceRequest) returns (QueryStakerBalanceResponse) {
+    option (cosmos.query.v1.module_query_safe) = true;
+    option (google.api.http).get = "/imuachain/assets/v1/staker_balance/{staker_id}/{asset_id}";
+  }
 }

proto/imuachain/delegation/v1/query.proto

@@ -51,11 +51,31 @@ message DelegationAmounts {
   ];
 }
 
+// SingleDelegationInfo includes the share, amount and pending undelegation amount
+message SingleDelegationInfo {
+  // delegation_amounts is the delegation info recorded in the KVStore
+  DelegationAmounts delegation_amounts = 1 ;
+  // max_undelegatable_amount is the maximum amount that can be undelegated
+  string max_undelegatable_amount = 2 [
+    (cosmos_proto.scalar) = "cosmos.Int",
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.nullable) = false
+  ];
+}
+
+// DelegationInfoAndOperator includes the SingleDelegationInfo and operator address
+message DelegationInfoAndOperator {
+  // operator is the operator address of this delegation
+  string operator = 1;
+  // delegation_info is the delegation information
+  SingleDelegationInfo delegation_info = 2;
+}
+
 // QueryDelegationInfoResponse is the response for delegations by staker id and
 // asset id.
 message QueryDelegationInfoResponse {
   // delegation_infos is the delegation information for each operator.
-  map<string, DelegationAmounts> delegation_infos = 1;
+  repeated DelegationInfoAndOperator delegation_infos = 1;
 }
 
 // SingleDelegationInfoReq is the request to obtain the single delegation information
@@ -71,14 +91,8 @@ message SingleDelegationInfoReq {
 
 // SingleDelegationInfoResponse is the response to QuerySingleDelegationInfo
 message SingleDelegationInfoResponse {
-  // delegation_amounts is the delegation info recorded in the KVStore
-  DelegationAmounts delegation_amounts = 1;
-  // max_undelegatable_amount is the maximum amount that can be undelegated
-  string max_undelegatable_amount = 2 [
-    (cosmos_proto.scalar) = "cosmos.Int",
-    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
-    (gogoproto.nullable) = false
-  ];
+  // single_delegation_info
+  SingleDelegationInfo single_delegation_info = 1 ;
 }
 
 // UndelegationHoldCountReq is the request to obtain the undelegation hold count.
@@ -221,5 +235,4 @@ service Query {
     option (cosmos.query.v1.module_query_safe) = true;
     option (google.api.http).get = "/imuachain/delegation/v1/delegated_stakers/{operator}/{asset_id}";
   }
-
 }

testutil/batch/tx_check.go

@@ -55,7 +55,7 @@ func (m *Manager) QueryDelegatedAmount(clientChainLzID uint64, stakerAddr, asset
 	if err != nil {
 		return sdkmath.ZeroInt(), err
 	}
-	return delegationInfo.MaxUndelegatableAmount, nil
+	return delegationInfo.SingleDelegationInfo.MaxUndelegatableAmount, nil
 }
 
 func (m *Manager) PrecompileTxOnChainCheck(batchID uint, msgType string) error {

x/assets/keeper/grpc_query.go

@@ -110,3 +110,14 @@ func (k Keeper) QueOperatorSpecifiedAssetAmount(ctx context.Context, req *assets
 	}
 	return k.GetOperatorSpecifiedAssetInfo(c, addr, strings.ToLower(req.AssetId))
 }
+
+func (k Keeper) QueryStakerBalance(ctx context.Context, req *assetstype.QueryStakerBalanceRequest) (*assetstype.QueryStakerBalanceResponse, error) {
+	c := sdk.UnwrapSDKContext(ctx)
+	stakerBalance, err := k.GetStakerBalanceByAsset(c, strings.ToLower(req.StakerId), strings.ToLower(req.AssetId))
+	if err != nil {
+		return nil, err
+	}
+	return &assetstype.QueryStakerBalanceResponse{
+		StakerBalance: &stakerBalance,
+	}, nil
+}

x/assets/keeper/staker_asset.go

@@ -3,14 +3,12 @@ package keeper
 import (
 	"fmt"
 
-	"github.com/ethereum/go-ethereum/common/hexutil"
-	assetstype "github.com/imua-xyz/imuachain/x/assets/types"
-	delegationkeeper "github.com/imua-xyz/imuachain/x/delegation/keeper"
-
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/math"
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/ethereum/go-ethereum/common/hexutil"
+	assetstype "github.com/imua-xyz/imuachain/x/assets/types"
 )
 
 // AllDeposits
@@ -99,21 +97,13 @@ func (k Keeper) GetStakerSpecifiedAssetInfo(ctx sdk.Context, stakerID string, as
 		if err != nil {
 			return nil, errorsmod.Wrap(err, "failed to GetDelegationInfo")
 		}
-		for operator, record := range delegationInfoRecords.DelegationInfos {
-			operatorAssetInfo, err := k.GetOperatorSpecifiedAssetInfo(ctx, sdk.MustAccAddressFromBech32(operator), assetID)
-			if err != nil {
-				return nil, errorsmod.Wrap(err, "failed to GetOperatorSpecifiedAssetInfo")
-			}
-			// the `undelegatableTokens` are currently delegated tokens. they are post-slashing, if any is applied.
+		for _, delegationInfo := range delegationInfoRecords.DelegationInfos {
+			// the `MaxUndelegatableAmount` are currently delegated tokens. they are post-slashing, if any is applied.
 			// this is because slashing is applied to an operator's total amount, of which, the share of a staker is kept
 			// unchanged.
-			undelegatableTokens, err := delegationkeeper.TokensFromShares(record.UndelegatableShare, operatorAssetInfo.TotalShare, operatorAssetInfo.TotalAmount)
-			if err != nil {
-				return nil, errorsmod.Wrap(err, "failed to get shares from token")
-			}
 			// this amount is post-slashing, as explained above.
-			info.TotalDepositAmount = info.TotalDepositAmount.Add(undelegatableTokens).Add(record.WaitUndelegationAmount)
-			info.PendingUndelegationAmount = info.PendingUndelegationAmount.Add(record.WaitUndelegationAmount)
+			info.TotalDepositAmount = info.TotalDepositAmount.Add(delegationInfo.DelegationInfo.MaxUndelegatableAmount).Add(delegationInfo.DelegationInfo.DelegationAmounts.WaitUndelegationAmount)
+			info.PendingUndelegationAmount = info.PendingUndelegationAmount.Add(delegationInfo.DelegationInfo.DelegationAmounts.WaitUndelegationAmount)
 		}
 		return info, nil
 	}
@@ -216,13 +206,13 @@ func (k Keeper) GetStakerBalanceByAsset(ctx sdk.Context, stakerID string, assetI
 	totalBalance := stakerAssetInfo.WithdrawableAmount.Add(stakerAssetInfo.PendingUndelegationAmount).Add(delegatedAmount)
 
 	balance = assetstype.StakerBalance{
-		StakerID:           stakerID,
-		AssetID:            assetID,
-		Balance:            totalBalance.BigInt(),
-		Withdrawable:       stakerAssetInfo.WithdrawableAmount.BigInt(),
-		Delegated:          delegatedAmount.BigInt(),
-		PendingUndelegated: stakerAssetInfo.PendingUndelegationAmount.BigInt(),
-		TotalDeposited:     stakerAssetInfo.TotalDepositAmount.BigInt(),
+		StakerId:           stakerID,
+		AssetId:            assetID,
+		Balance:            totalBalance,
+		Withdrawable:       stakerAssetInfo.WithdrawableAmount,
+		Delegated:          delegatedAmount,
+		PendingUndelegated: stakerAssetInfo.PendingUndelegationAmount,
+		TotalDeposited:     stakerAssetInfo.TotalDepositAmount,
 	}
 
 	return balance, nil

x/assets/types/general.go

@@ -2,7 +2,6 @@ package types
 
 import (
 	"fmt"
-	"math/big"
 	"strings"
 
 	"github.com/imua-xyz/imuachain/utils"
@@ -98,15 +97,15 @@ type CreateQueryContext func(height int64, prove bool) (sdk.Context, error)
 // pendingUndelegated is the amount of the asset that is during unbonding period and not yet withdrawable
 // it would finally be withdrawable after the unbonding period, but the final amount may be less than the pendingUndelegated
 // because of the penalty during the unbonding period
-type StakerBalance struct {
+/*type StakerBalance struct {
 	StakerID           string
 	AssetID            string
 	Balance            *big.Int
 	Withdrawable       *big.Int
 	Delegated          *big.Int
 	PendingUndelegated *big.Int
 	TotalDeposited     *big.Int
-}
+}*/
 
 // GetStakerIDAndAssetID stakerID = stakerAddress+'_'+clientChainLzID,assetID =
 // assetAddress+'_'+clientChainLzID