Merge pull request #36 from hadronlabs-org/feat/merge_delegations_bal…

…ance_icq

feat: merge delegations and balance ICQs

contracts/core/src/contract.rs

@@ -122,8 +122,12 @@ fn query_exchange_rate(
     let extra_amount = match FSM.get_current_state(deps.storage)? {
         ContractState::Transfering => PENDING_TRANSFER.load(deps.storage),
         ContractState::Staking => {
-            let (ica_balance, _) =
-                get_ica_balance_by_denom(deps, &config.puppeteer_contract, &config.remote_denom)?;
+            let (ica_balance, _) = get_ica_balance_by_denom(
+                deps,
+                &config.puppeteer_contract,
+                &config.remote_denom,
+                false,
+            )?;
             Ok(ica_balance)
         }
         _ => Ok(Uint128::zero()),
@@ -397,6 +401,7 @@ fn execute_tick_staking(
             deps.as_ref(),
             &config.puppeteer_contract,
             &config.remote_denom,
+            true,
         )?;
         PRE_UNBONDING_BALANCE.save(deps.storage, &pre_unbonding_balance)?;
         FSM.go_to(deps.storage, ContractState::Unbonding)?;
@@ -809,8 +814,12 @@ fn get_stake_msg<T>(
     config: &Config,
     funds: Vec<cosmwasm_std::Coin>,
 ) -> ContractResult<CosmosMsg<T>> {
-    let (balance, balance_height) =
-        get_ica_balance_by_denom(deps, &config.puppeteer_contract, &config.remote_denom)?;
+    let (balance, balance_height) = get_ica_balance_by_denom(
+        deps,
+        &config.puppeteer_contract,
+        &config.remote_denom,
+        false,
+    )?;
     ensure_ne!(balance, Uint128::zero(), ContractError::ICABalanceZero {});
     let last_ica_balance_change_height = LAST_ICA_BALANCE_CHANGE_HEIGHT.load(deps.storage)?;
     ensure!(
@@ -872,6 +881,7 @@ fn get_ica_balance_by_denom<T: CustomQuery>(
     deps: Deps<T>,
     puppeteer_contract: &str,
     remote_denom: &str,
+    can_be_zero: bool,
 ) -> ContractResult<(Uint128, u64)> {
     let (ica_balances, remote_height): lido_staking_base::msg::puppeteer::BalancesResponse =
         deps.querier.query_wasm_smart(
@@ -880,18 +890,20 @@ fn get_ica_balance_by_denom<T: CustomQuery>(
                 msg: lido_staking_base::msg::puppeteer::QueryExtMsg::Balances {},
             },
         )?;
-    let balance = ica_balances
-        .coins
-        .iter()
-        .find_map(|c| {
-            if c.denom == remote_denom {
-                Some(c.amount)
-            } else {
-                None
-            }
-        })
-        .ok_or(ContractError::ICABalanceZero {})?;
-    Ok((balance, remote_height))
+    let balance = ica_balances.coins.iter().find_map(|c| {
+        if c.denom == remote_denom {
+            Some(c.amount)
+        } else {
+            None
+        }
+    });
+    Ok((
+        match can_be_zero {
+            true => balance.unwrap_or(Uint128::zero()),
+            false => balance.ok_or(ContractError::ICABalanceZero {})?,
+        },
+        remote_height,
+    ))
 }
 
 fn new_unbond(now: u64) -> lido_staking_base::state::core::UnbondBatch {

contracts/factory/src/contract.rs

@@ -152,7 +152,7 @@ fn execute_proxy_msg(
                 )?);
                 messages.push(get_proxied_message(
                     state.puppeteer_contract,
-                    lido_staking_base::msg::puppeteer::ExecuteMsg::RegisterDelegatorDelegationsQuery { validators: validators.iter().map(|v| {v.valoper_address.to_string()}).collect() },
+                    lido_staking_base::msg::puppeteer::ExecuteMsg::RegisterBalanceAndDelegatorDelegationsQuery { validators: validators.iter().map(|v| {v.valoper_address.to_string()}).collect() },
                     info.funds,
                 )?)
             }

contracts/puppeteer/src/contract.rs

@@ -20,7 +20,10 @@ use cosmwasm_std::{
 };
 use cosmwasm_std::{Binary, DepsMut, Env, MessageInfo, Response, StdResult};
 use cw2::set_contract_version;
-use lido_helpers::answer::response;
+use lido_helpers::{
+    answer::response,
+    icq::{new_delegations_and_balance_query_msg, update_balance_and_delegations_query_msg},
+};
 use lido_puppeteer_base::{
     error::{ContractError, ContractResult},
     msg::{
@@ -34,7 +37,7 @@ use lido_puppeteer_base::{
 };
 use lido_staking_base::{
     msg::puppeteer::{ExecuteMsg, InstantiateMsg, MigrateMsg, QueryExtMsg},
-    state::puppeteer::{Config, KVQueryType, BALANCES, DELEGATIONS},
+    state::puppeteer::{Config, KVQueryType, DELEGATIONS_AND_BALANCE},
 };
 use neutron_sdk::interchain_queries::v045::new_register_delegator_unbonding_delegations_query_msg;
 use neutron_sdk::{
@@ -43,10 +46,7 @@ use neutron_sdk::{
         query::NeutronQuery,
         types::ProtobufAny,
     },
-    interchain_queries::v045::{
-        new_register_balance_query_msg, new_register_delegator_delegations_query_msg,
-        types::{Balances, Delegations},
-    },
+    interchain_queries::v045::types::{Balances, Delegations},
     interchain_txs::helpers::decode_message_response,
     sudo::msg::{RequestPacket, RequestPacketTimeoutHeight, SudoMsg},
     NeutronError, NeutronResult,
@@ -84,16 +84,16 @@ pub fn instantiate(
         proxy_address: None,
         transfer_channel_id: msg.transfer_channel_id,
     };
-    DELEGATIONS.save(
+    DELEGATIONS_AND_BALANCE.save(
         deps.storage,
         &(
             Delegations {
                 delegations: vec![],
             },
+            Balances { coins: vec![] },
             0,
         ),
     )?;
-    BALANCES.save(deps.storage, &(Balances { coins: vec![] }, 0))?;
     Puppeteer::default().instantiate(deps, config)
 }
 
@@ -105,12 +105,8 @@ pub fn query(
 ) -> ContractResult<Binary> {
     match msg {
         QueryMsg::Extention { msg } => match msg {
-            QueryExtMsg::Delegations {} => {
-                to_json_binary(&DELEGATIONS.load(deps.storage)?).map_err(ContractError::Std)
-            }
-            QueryExtMsg::Balances {} => {
-                to_json_binary(&BALANCES.load(deps.storage)?).map_err(ContractError::Std)
-            }
+            QueryExtMsg::Delegations {} => query_delegations(deps),
+            QueryExtMsg::Balances {} => query_balances(deps),
             QueryExtMsg::UnbondingDelegations {} => to_json_binary(
                 &Puppeteer::default()
                     .unbonding_delegations
@@ -124,6 +120,16 @@ pub fn query(
     }
 }
 
+fn query_delegations(deps: Deps<NeutronQuery>) -> ContractResult<Binary> {
+    let data = DELEGATIONS_AND_BALANCE.load(deps.storage)?;
+    to_json_binary(&(data.0, data.2)).map_err(ContractError::Std)
+}
+
+fn query_balances(deps: Deps<NeutronQuery>) -> ContractResult<Binary> {
+    let data = DELEGATIONS_AND_BALANCE.load(deps.storage)?;
+    to_json_binary(&(data.1, data.2)).map_err(ContractError::Std)
+}
+
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn execute(
     deps: DepsMut<NeutronQuery>,
@@ -180,13 +186,12 @@ pub fn execute(
         } => execute_claim_rewards_and_optionaly_transfer(
             deps, info, validators, transfer, timeout, reply_to,
         ),
-        ExecuteMsg::RegisterDelegatorDelegationsQuery { validators } => {
-            register_delegations_query(deps, info, validators)
+        ExecuteMsg::RegisterBalanceAndDelegatorDelegationsQuery { validators } => {
+            register_balance_delegations_query(deps, info, validators)
         }
         ExecuteMsg::RegisterDelegatorUnbondingDelegationsQuery { validators } => {
             register_unbonding_delegations_query(deps, info, validators)
         }
-        ExecuteMsg::RegisterBalanceQuery { denom } => register_balance_query(deps, denom),
         ExecuteMsg::IBCTransfer { timeout, reply_to } => {
             execute_ibc_transfer(deps, env, info, timeout, reply_to)
         }
@@ -250,41 +255,46 @@ fn execute_ibc_transfer(
     Ok(Response::default().add_submessages(vec![submsg]))
 }
 
-fn register_delegations_query(
+fn register_balance_delegations_query(
     deps: DepsMut<NeutronQuery>,
     info: MessageInfo,
     validators: Vec<String>,
 ) -> ContractResult<Response<NeutronMsg>> {
     let puppeteer_base = Puppeteer::default();
     let config = puppeteer_base.config.load(deps.storage)?;
     ensure_eq!(config.owner, info.sender, ContractError::Unauthorized {});
-    // remove old delegation query if any
     let kv_queries = puppeteer_base
         .kv_queries
         .range(deps.storage, None, None, Order::Ascending)
         .collect::<Result<Vec<(u64, KVQueryType)>, _>>()?;
+    let ica = puppeteer_base.ica.get_address(deps.storage)?;
     let mut messages = vec![];
+    let mut submessages = vec![];
     for (query_id, query_type) in kv_queries {
-        if query_type == KVQueryType::Delegations {
-            messages.push(NeutronMsg::remove_interchain_query(query_id));
+        if query_type == KVQueryType::DelegationsAndBalance {
+            messages.push(update_balance_and_delegations_query_msg(
+                query_id,
+                ica.to_string(),
+                config.remote_denom.to_string(),
+                validators.clone(),
+            )?); //no need to handle reply as nothing to update in the query
         }
     }
-    //
-    let delegator = puppeteer_base.ica.get_address(deps.storage)?;
-    let msg = SubMsg::reply_on_success(
-        new_register_delegator_delegations_query_msg(
-            config.connection_id,
-            delegator,
-            validators,
-            config.update_period,
-        )?,
-        ReplyMsg::KvDelegations.to_reply_id(),
-    );
-    deps.api.debug(&format!(
-        "WASMDEBUG: register_delegations_query {msg:?}",
-        msg = msg
-    ));
-    Ok(Response::new().add_messages(messages).add_submessage(msg))
+    if messages.is_empty() {
+        submessages.push(SubMsg::reply_on_success(
+            new_delegations_and_balance_query_msg(
+                config.connection_id.clone(),
+                ica.clone(),
+                config.remote_denom.clone(),
+                validators.clone(),
+                config.update_period,
+            )?,
+            ReplyMsg::KvDelegationsAndBalance.to_reply_id(),
+        ));
+    }
+    Ok(Response::new()
+        .add_messages(messages)
+        .add_submessages(submessages))
 }
 
 fn register_unbonding_delegations_query(
@@ -339,24 +349,6 @@ fn register_unbonding_delegations_query(
     Ok(Response::new().add_submessages(msgs))
 }
 
-fn register_balance_query(
-    deps: DepsMut<NeutronQuery>,
-    denom: String,
-) -> ContractResult<Response<NeutronMsg>> {
-    let puppeteer_base = Puppeteer::default();
-    let config = puppeteer_base.config.load(deps.storage)?;
-    let ica = puppeteer_base.ica.get_address(deps.storage)?;
-    let msg = SubMsg::reply_on_success(
-        new_register_balance_query_msg(config.connection_id, ica, denom, config.update_period)?,
-        ReplyMsg::KvBalance.to_reply_id(),
-    );
-    deps.api.debug(&format!(
-        "WASMDEBUG: register_balance_query {msg:?}",
-        msg = msg
-    ));
-    Ok(Response::new().add_submessage(msg))
-}
-
 fn execute_delegate(
     mut deps: DepsMut<NeutronQuery>,
     info: MessageInfo,
@@ -706,13 +698,15 @@ pub fn sudo(deps: DepsMut<NeutronQuery>, env: Env, msg: SudoMsg) -> NeutronResul
         } => puppeteer_base.sudo_tx_query_result(deps, env, query_id, height, data),
         SudoMsg::KVQueryResult { query_id } => {
             let query_type = puppeteer_base.kv_queries.load(deps.storage, query_id)?;
+            deps.api
+                .debug(&format!("WASMDEBUG: KVQueryResult type {:?}", query_type));
             match query_type {
-                KVQueryType::Balance => {
-                    puppeteer_base.sudo_kv_query_result(deps, env, query_id, BALANCES)
-                }
-                KVQueryType::Delegations => {
-                    puppeteer_base.sudo_kv_query_result(deps, env, query_id, DELEGATIONS)
-                }
+                KVQueryType::DelegationsAndBalance => puppeteer_base.sudo_kv_query_result(
+                    deps,
+                    env,
+                    query_id,
+                    DELEGATIONS_AND_BALANCE,
+                ),
                 KVQueryType::UnbondingDelegations => {
                     puppeteer_base.sudo_unbonding_delegations_kv_query_result(deps, env, query_id)
                 }
@@ -981,15 +975,8 @@ pub fn reply(deps: DepsMut, _env: Env, msg: Reply) -> StdResult<Response> {
     match ReplyMsg::from_reply_id(msg.id) {
         ReplyMsg::SudoPayload => puppeteer_base.submit_tx_reply(deps, msg),
         ReplyMsg::IbcTransfer => puppeteer_base.submit_ibc_transfer_reply(deps, msg),
-        ReplyMsg::KvBalance => {
-            deps.api
-                .debug(&format!("WASMDEBUG: KV_BALANCE_REPLY_ID {:?}", msg));
-            puppeteer_base.register_kv_query_reply(deps, msg, KVQueryType::Balance)
-        }
-        ReplyMsg::KvDelegations => {
-            deps.api
-                .debug(&format!("WASMDEBUG: DELEGATIONS_REPLY_ID {:?}", msg));
-            puppeteer_base.register_kv_query_reply(deps, msg, KVQueryType::Delegations)
+        ReplyMsg::KvDelegationsAndBalance => {
+            puppeteer_base.register_kv_query_reply(deps, msg, KVQueryType::DelegationsAndBalance)
         }
         ReplyMsg::KvUnbondingDelegations { validator_index } => {
             deps.api.debug(&format!(

integration_tests/src/generated/contractLib/lidoPuppeteer.ts

@@ -115,9 +115,8 @@ export interface LidoPuppeteerSchema {
   responses: ConfigResponse | Binary | IcaState | ArrayOfTransaction;
   query: ExtentionArgs;
   execute:
-    | RegisterDelegatorDelegationsQueryArgs
+    | RegisterBalanceAndDelegatorDelegationsQueryArgs
     | RegisterDelegatorUnbondingDelegationsQueryArgs
-    | RegisterBalanceQueryArgs
     | SetFeesArgs
     | DelegateArgs
     | UndelegateArgs
@@ -141,15 +140,12 @@ export interface TransferReadyBatchMsg {
 export interface ExtentionArgs {
   msg: QueryExtMsg;
 }
-export interface RegisterDelegatorDelegationsQueryArgs {
+export interface RegisterBalanceAndDelegatorDelegationsQueryArgs {
   validators: string[];
 }
 export interface RegisterDelegatorUnbondingDelegationsQueryArgs {
   validators: string[];
 }
-export interface RegisterBalanceQueryArgs {
-  denom: string;
-}
 export interface SetFeesArgs {
   ack_fee: Uint128;
   recv_fee: Uint128;
@@ -249,18 +245,14 @@ export class Client {
           if (!isSigningCosmWasmClient(this.client)) { throw this.mustBeSigningClient(); }
     return this.client.execute(sender, this.contractAddress, { register_query: {} }, fee || "auto", memo, funds);
   }
-  registerDelegatorDelegationsQuery = async(sender:string, args: RegisterDelegatorDelegationsQueryArgs, fee?: number | StdFee | "auto", memo?: string, funds?: Coin[]): Promise<ExecuteResult> =>  {
+  registerBalanceAndDelegatorDelegationsQuery = async(sender:string, args: RegisterBalanceAndDelegatorDelegationsQueryArgs, fee?: number | StdFee | "auto", memo?: string, funds?: Coin[]): Promise<ExecuteResult> =>  {
           if (!isSigningCosmWasmClient(this.client)) { throw this.mustBeSigningClient(); }
-    return this.client.execute(sender, this.contractAddress, { register_delegator_delegations_query: args }, fee || "auto", memo, funds);
+    return this.client.execute(sender, this.contractAddress, { register_balance_and_delegator_delegations_query: args }, fee || "auto", memo, funds);
   }
   registerDelegatorUnbondingDelegationsQuery = async(sender:string, args: RegisterDelegatorUnbondingDelegationsQueryArgs, fee?: number | StdFee | "auto", memo?: string, funds?: Coin[]): Promise<ExecuteResult> =>  {
           if (!isSigningCosmWasmClient(this.client)) { throw this.mustBeSigningClient(); }
     return this.client.execute(sender, this.contractAddress, { register_delegator_unbonding_delegations_query: args }, fee || "auto", memo, funds);
   }
-  registerBalanceQuery = async(sender:string, args: RegisterBalanceQueryArgs, fee?: number | StdFee | "auto", memo?: string, funds?: Coin[]): Promise<ExecuteResult> =>  {
-          if (!isSigningCosmWasmClient(this.client)) { throw this.mustBeSigningClient(); }
-    return this.client.execute(sender, this.contractAddress, { register_balance_query: args }, fee || "auto", memo, funds);
-  }
   setFees = async(sender:string, args: SetFeesArgs, fee?: number | StdFee | "auto", memo?: string, funds?: Coin[]): Promise<ExecuteResult> =>  {
           if (!isSigningCosmWasmClient(this.client)) { throw this.mustBeSigningClient(); }
     return this.client.execute(sender, this.contractAddress, { set_fees: args }, fee || "auto", memo, funds);