Staking precompile delegation query (sei-protocol#1770)

* working draft

* - hardhat test
- basic unit test

* formatting

* uncomment all hardhat tests

* more testing

* - add unassociated address test
- remove redundant want clause

* - flatten abi

* update test

* remove extra comment

* add missing change

app/app.go

@@ -690,6 +690,7 @@ func New(
 			wasmkeeper.NewDefaultPermissionKeeper(app.WasmKeeper),
 			app.WasmKeeper,
 			stakingkeeper.NewMsgServerImpl(app.StakingKeeper),
+			stakingkeeper.Querier{Keeper: app.StakingKeeper},
 			app.GovKeeper,
 			app.DistrKeeper,
 			app.OracleKeeper,

contracts/test/EVMPrecompileTest.js

@@ -90,7 +90,21 @@ describe("EVM Precompile Tester", function () {
             });
             const receipt = await delegate.wait();
             expect(receipt.status).to.equal(1);
-            // TODO: Add staking query precompile here
+
+            const delegation = await staking.delegation(accounts[0].evmAddress, validatorAddr);
+            expect(delegation).to.not.be.null;
+            expect(delegation[0][0]).to.equal(10000n);
+
+            const undelegate = await staking.undelegate(validatorAddr, delegation[0][0]);
+            const undelegateReceipt = await undelegate.wait();
+            expect(undelegateReceipt.status).to.equal(1);
+
+            try {
+                await staking.delegation(accounts[0].evmAddress, validatorAddr);
+                expect.fail("Expected an error here since we undelegated the amount and delegation should not exist anymore.");
+            } catch (error) {
+                expect(error).to.have.property('message').that.includes('execution reverted');
+            }
         });
     });
 

precompiles/common/expected_keepers.go

@@ -81,6 +81,10 @@ type StakingKeeper interface {
 	Undelegate(goCtx context.Context, msg *stakingtypes.MsgUndelegate) (*stakingtypes.MsgUndelegateResponse, error)
 }
 
+type StakingQuerier interface {
+	Delegation(c context.Context, req *stakingtypes.QueryDelegationRequest) (*stakingtypes.QueryDelegationResponse, error)
+}
+
 type GovKeeper interface {
 	AddVote(ctx sdk.Context, proposalID uint64, voterAddr sdk.AccAddress, options govtypes.WeightedVoteOptions) error
 	AddDeposit(ctx sdk.Context, proposalID uint64, depositorAddr sdk.AccAddress, depositAmount sdk.Coins) (bool, error)

precompiles/common/precompiles.go

@@ -16,6 +16,7 @@ import (
 	"github.com/sei-protocol/sei-chain/utils"
 	"github.com/sei-protocol/sei-chain/utils/metrics"
 	"github.com/sei-protocol/sei-chain/x/evm/state"
+	"github.com/sei-protocol/sei-chain/x/evm/types"
 )
 
 const UnknownMethodCallGas uint64 = 3000
@@ -263,3 +264,19 @@ func MustGetABI(f embed.FS, filename string) abi.ABI {
 	}
 	return newAbi
 }
+
+func GetSeiAddressByEvmAddress(ctx sdk.Context, evmAddress common.Address, evmKeeper EVMKeeper) (sdk.AccAddress, error) {
+	seiAddr, associated := evmKeeper.GetSeiAddress(ctx, evmAddress)
+	if !associated {
+		return nil, types.NewAssociationMissingErr(evmAddress.Hex())
+	}
+	return seiAddr, nil
+}
+
+func GetSeiAddressFromArg(ctx sdk.Context, arg interface{}, evmKeeper EVMKeeper) (sdk.AccAddress, error) {
+	addr := arg.(common.Address)
+	if addr == (common.Address{}) {
+		return nil, errors.New("invalid addr")
+	}
+	return GetSeiAddressByEvmAddress(ctx, addr, evmKeeper)
+}

precompiles/setup.go

@@ -45,6 +45,7 @@ func InitializePrecompiles(
 	wasmdKeeper common.WasmdKeeper,
 	wasmdViewKeeper common.WasmdViewKeeper,
 	stakingKeeper common.StakingKeeper,
+	stakingQuerier common.StakingQuerier,
 	govKeeper common.GovKeeper,
 	distrKeeper common.DistributionKeeper,
 	oracleKeeper common.OracleKeeper,
@@ -75,7 +76,7 @@ func InitializePrecompiles(
 	if err != nil {
 		return err
 	}
-	stakingp, err := staking.NewPrecompile(stakingKeeper, evmKeeper, bankKeeper)
+	stakingp, err := staking.NewPrecompile(stakingKeeper, stakingQuerier, evmKeeper, bankKeeper)
 	if err != nil {
 		return err
 	}
@@ -134,7 +135,7 @@ func InitializePrecompiles(
 func GetPrecompileInfo(name string) PrecompileInfo {
 	if !Initialized {
 		// Precompile Info does not require any keeper state
-		_ = InitializePrecompiles(true, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
+		_ = InitializePrecompiles(true, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
 	}
 	i, ok := PrecompileNamesToInfo[name]
 	if !ok {

precompiles/staking/Staking.sol

@@ -23,4 +23,27 @@ interface IStaking {
         string memory valAddress,
         uint256 amount
     ) external returns (bool success);
+
+    // Queries
+    function delegation(
+        address delegator,
+        string memory valAddress
+    ) external view returns (Delegation delegation);
+
+    struct Delegation {
+        Balance balance;
+        DelegationDetails delegation;
+    }
+
+    struct Balance {
+        uint256 amount;
+        string denom;
+    }
+
+    struct DelegationDetails {
+        string delegator_address;
+        uint256 shares;
+        uint256 decimals;
+        string validator_address;
+    }
 }

precompiles/staking/abi.json

@@ -1 +1 @@
-[{"inputs":[{"internalType":"string","name":"valAddress","type":"string"}],"name":"delegate","outputs":[{"internalType":"bool","name":"success","type":"bool"}],"stateMutability":"payable","type":"function"},{"inputs":[{"internalType":"string","name":"srcAddress","type":"string"},{"internalType":"string","name":"dstAddress","type":"string"},{"internalType":"uint256","name":"amount","type":"uint256"}],"name":"redelegate","outputs":[{"internalType":"bool","name":"success","type":"bool"}],"stateMutability":"nonpayable","type":"function"},{"inputs":[{"internalType":"string","name":"valAddress","type":"string"},{"internalType":"uint256","name":"amount","type":"uint256"}],"name":"undelegate","outputs":[{"internalType":"bool","name":"success","type":"bool"}],"stateMutability":"nonpayable","type":"function"}]
+[{"inputs":[{"internalType":"string","name":"valAddress","type":"string"}],"name":"delegate","outputs":[{"internalType":"bool","name":"success","type":"bool"}],"stateMutability":"payable","type":"function"},{"inputs":[{"internalType":"string","name":"srcAddress","type":"string"},{"internalType":"string","name":"dstAddress","type":"string"},{"internalType":"uint256","name":"amount","type":"uint256"}],"name":"redelegate","outputs":[{"internalType":"bool","name":"success","type":"bool"}],"stateMutability":"nonpayable","type":"function"},{"inputs":[{"internalType":"string","name":"valAddress","type":"string"},{"internalType":"uint256","name":"amount","type":"uint256"}],"name":"undelegate","outputs":[{"internalType":"bool","name":"success","type":"bool"}],"stateMutability":"nonpayable","type":"function"},{"inputs":[{"internalType":"address","name":"delegator","type":"address"},{"internalType":"string","name":"valAddress","type":"string"}],"name":"delegation","outputs":[{"components":[{"components":[{"internalType":"uint256","name":"amount","type":"uint256"},{"internalType":"string","name":"denom","type":"string"}],"internalType":"struct Balance","name":"balance","type":"tuple"},{"components":[{"internalType":"string","name":"delegator_address","type":"string"},{"internalType":"uint256","name":"shares","type":"uint256"},{"internalType":"uint256","name":"decimals","type":"uint256"},{"internalType":"string","name":"validator_address","type":"string"}],"internalType":"struct DelegationDetails","name":"delegation","type":"tuple"}],"internalType":"struct Delegation","name":"delegation","type":"tuple"}],"stateMutability":"view","type":"function"}]

precompiles/staking/staking.go

@@ -19,6 +19,7 @@ const (
 	DelegateMethod   = "delegate"
 	RedelegateMethod = "redelegate"
 	UndelegateMethod = "undelegate"
+	DelegationMethod = "delegation"
 )
 
 const (
@@ -31,24 +32,27 @@ const (
 var f embed.FS
 
 type PrecompileExecutor struct {
-	stakingKeeper pcommon.StakingKeeper
-	evmKeeper     pcommon.EVMKeeper
-	bankKeeper    pcommon.BankKeeper
-	address       common.Address
+	stakingKeeper  pcommon.StakingKeeper
+	stakingQuerier pcommon.StakingQuerier
+	evmKeeper      pcommon.EVMKeeper
+	bankKeeper     pcommon.BankKeeper
+	address        common.Address
 
 	DelegateID   []byte
 	RedelegateID []byte
 	UndelegateID []byte
+	DelegationID []byte
 }
 
-func NewPrecompile(stakingKeeper pcommon.StakingKeeper, evmKeeper pcommon.EVMKeeper, bankKeeper pcommon.BankKeeper) (*pcommon.Precompile, error) {
+func NewPrecompile(stakingKeeper pcommon.StakingKeeper, stakingQuerier pcommon.StakingQuerier, evmKeeper pcommon.EVMKeeper, bankKeeper pcommon.BankKeeper) (*pcommon.Precompile, error) {
 	newAbi := pcommon.MustGetABI(f, "abi.json")
 
 	p := &PrecompileExecutor{
-		stakingKeeper: stakingKeeper,
-		evmKeeper:     evmKeeper,
-		bankKeeper:    bankKeeper,
-		address:       common.HexToAddress(StakingAddress),
+		stakingKeeper:  stakingKeeper,
+		stakingQuerier: stakingQuerier,
+		evmKeeper:      evmKeeper,
+		bankKeeper:     bankKeeper,
+		address:        common.HexToAddress(StakingAddress),
 	}
 
 	for name, m := range newAbi.Methods {
@@ -59,6 +63,8 @@ func NewPrecompile(stakingKeeper pcommon.StakingKeeper, evmKeeper pcommon.EVMKee
 			p.RedelegateID = m.ID
 		case UndelegateMethod:
 			p.UndelegateID = m.ID
+		case DelegationMethod:
+			p.DelegationID = m.ID
 		}
 	}
 
@@ -80,20 +86,27 @@ func (p PrecompileExecutor) RequiredGas(input []byte, method *abi.Method) uint64
 }
 
 func (p PrecompileExecutor) Execute(ctx sdk.Context, method *abi.Method, caller common.Address, callingContract common.Address, args []interface{}, value *big.Int, readOnly bool, evm *vm.EVM) (bz []byte, err error) {
-	if readOnly {
-		return nil, errors.New("cannot call staking precompile from staticcall")
-	}
 	if caller.Cmp(callingContract) != 0 {
 		return nil, errors.New("cannot delegatecall staking")
 	}
-
 	switch method.Name {
 	case DelegateMethod:
+		if readOnly {
+			return nil, errors.New("cannot call staking precompile from staticcall")
+		}
 		return p.delegate(ctx, method, caller, args, value)
 	case RedelegateMethod:
+		if readOnly {
+			return nil, errors.New("cannot call staking precompile from staticcall")
+		}
 		return p.redelegate(ctx, method, caller, args, value)
 	case UndelegateMethod:
+		if readOnly {
+			return nil, errors.New("cannot call staking precompile from staticcall")
+		}
 		return p.undelegate(ctx, method, caller, args, value)
+	case DelegationMethod:
+		return p.delegation(ctx, method, args, value)
 	}
 	return
 }
@@ -179,3 +192,61 @@ func (p PrecompileExecutor) undelegate(ctx sdk.Context, method *abi.Method, call
 	}
 	return method.Outputs.Pack(true)
 }
+
+type Delegation struct {
+	Balance    Balance
+	Delegation DelegationDetails
+}
+
+type Balance struct {
+	Amount *big.Int
+	Denom  string
+}
+
+type DelegationDetails struct {
+	DelegatorAddress string
+	Shares           *big.Int
+	Decimals         *big.Int
+	ValidatorAddress string
+}
+
+func (p PrecompileExecutor) delegation(ctx sdk.Context, method *abi.Method, args []interface{}, value *big.Int) ([]byte, error) {
+	if err := pcommon.ValidateNonPayable(value); err != nil {
+		return nil, err
+	}
+
+	if err := pcommon.ValidateArgsLength(args, 2); err != nil {
+		return nil, err
+	}
+
+	seiDelegatorAddress, err := pcommon.GetSeiAddressFromArg(ctx, args[0], p.evmKeeper)
+	if err != nil {
+		return nil, err
+	}
+
+	validatorBech32 := args[1].(string)
+	delegationRequest := &stakingtypes.QueryDelegationRequest{
+		DelegatorAddr: seiDelegatorAddress.String(),
+		ValidatorAddr: validatorBech32,
+	}
+
+	delegationResponse, err := p.stakingQuerier.Delegation(sdk.WrapSDKContext(ctx), delegationRequest)
+	if err != nil {
+		return nil, err
+	}
+
+	delegation := Delegation{
+		Balance: Balance{
+			Amount: delegationResponse.GetDelegationResponse().GetBalance().Amount.BigInt(),
+			Denom:  delegationResponse.GetDelegationResponse().GetBalance().Denom,
+		},
+		Delegation: DelegationDetails{
+			DelegatorAddress: delegationResponse.GetDelegationResponse().GetDelegation().DelegatorAddress,
+			Shares:           delegationResponse.GetDelegationResponse().GetDelegation().Shares.BigInt(),
+			Decimals:         big.NewInt(sdk.Precision),
+			ValidatorAddress: delegationResponse.GetDelegationResponse().GetDelegation().ValidatorAddress,
+		},
+	}
+
+	return method.Outputs.Pack(delegation)
+}