feat(contracts): add rewards staking example

contracts/increment/src/contract.rs

@@ -1,25 +1,26 @@
-use archway_bindings::types::gov::VoteResponse;
-use archway_bindings::types::rewards::{
-    ContractMetadataResponse, FlatFeeResponse, RewardsRecordsResponse, WithdrawRewardsResponse,
-};
-use archway_bindings::{ArchwayMsg, ArchwayQuery, ArchwayResult, PageRequest};
 use cosmwasm_std::{
     coin, entry_point, to_binary, Addr, Binary, Deps, DepsMut, Env, MessageInfo, Reply, Response,
-    StdError, StdResult, SubMsg, Uint128,
+    StakingMsg, StdError, StdResult, SubMsg, Uint128,
 };
 use cw2::set_contract_version;
 use cw_utils::NativeBalance;
 
+use archway_bindings::types::gov::VoteResponse;
+use archway_bindings::types::rewards::{
+    ContractMetadataResponse, FlatFeeResponse, RewardsRecordsResponse, WithdrawRewardsResponse,
+};
+use archway_bindings::{ArchwayMsg, ArchwayQuery, ArchwayResult, PageRequest};
+
 use crate::error::ContractError;
 use crate::helpers;
 use crate::msg::{CountResponse, ExecuteMsg, InstantiateMsg, OutstandingRewardsResponse, QueryMsg};
-use crate::state::{State, STATE};
+use crate::state::{State, STAKE_REWARDS_VALIDATOR, STATE};
 
 // version info for migration info
 const CONTRACT_NAME: &str = "crates.io:increment";
 const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
 
-const REWARDS_WITHDRAW_REPLY: u64 = 1001;
+const STAKE_WITHDRAWN_REWARDS: u64 = 1001;
 
 #[entry_point]
 pub fn instantiate(
@@ -57,7 +58,9 @@ pub fn execute(
         ExecuteMsg::SetFlatFee { amount } => {
             set_flat_fee(deps.as_ref(), env.contract.address, amount)
         }
-        ExecuteMsg::WithdrawRewards {} => withdraw_rewards(),
+        ExecuteMsg::WithdrawRewards { stake_to_validator } => {
+            withdraw_rewards(deps, env, stake_to_validator)
+        }
     }
 }
 
@@ -115,43 +118,73 @@ fn set_flat_fee(
     Ok(res)
 }
 
-fn withdraw_rewards() -> ArchwayResult<ContractError> {
+fn withdraw_rewards(
+    deps: DepsMut<ArchwayQuery>,
+    env: Env,
+    stake_to_validator: Option<Addr>,
+) -> ArchwayResult<ContractError> {
     let msg = ArchwayMsg::withdraw_rewards_by_limit(0);
+    let withdraw_rewards_msg = if let Some(validator_addr) = stake_to_validator {
+        let tx_index = env.transaction.map(|tx| tx.index).unwrap_or(0);
+        STAKE_REWARDS_VALIDATOR.save(
+            deps.storage,
+            (env.block.height, tx_index),
+            &validator_addr,
+        )?;
+        SubMsg::reply_always(msg, STAKE_WITHDRAWN_REWARDS)
+    } else {
+        SubMsg::new(msg)
+    };
 
     let res = Response::new()
-        .add_submessage(SubMsg::reply_on_success(msg, REWARDS_WITHDRAW_REPLY))
+        .add_submessage(withdraw_rewards_msg)
         .add_attribute("method", "withdraw_rewards");
 
     Ok(res)
 }
 
 #[entry_point]
-pub fn reply(deps: DepsMut<ArchwayQuery>, _env: Env, msg: Reply) -> StdResult<Response> {
+pub fn reply(deps: DepsMut<ArchwayQuery>, env: Env, msg: Reply) -> StdResult<Response> {
     match msg.id {
-        REWARDS_WITHDRAW_REPLY => after_rewards_withdrawn(deps, msg),
+        STAKE_WITHDRAWN_REWARDS => stake_contract_rewards(deps, env, msg),
         id => Err(StdError::not_found(format!("Unknown reply id: {}", id))),
     }
 }
 
-fn after_rewards_withdrawn(_deps: DepsMut<ArchwayQuery>, msg: Reply) -> StdResult<Response> {
+fn stake_contract_rewards(
+    deps: DepsMut<ArchwayQuery>,
+    env: Env,
+    msg: Reply,
+) -> StdResult<Response> {
+    let tx_index = env.transaction.map(|tx| tx.index).unwrap_or(0);
+    let validator_addr =
+        STAKE_REWARDS_VALIDATOR.load(deps.storage, (env.block.height, tx_index))?;
+
+    STAKE_REWARDS_VALIDATOR.remove(deps.storage, (env.block.height, tx_index));
+
     let data = helpers::parse_reply_data(msg)?;
     let withdraw_response: WithdrawRewardsResponse =
         serde_json_wasm::from_slice::<WithdrawRewardsResponse>(&data.0)
             .map_err(|e| StdError::generic_err(e.to_string()))?;
 
-    let mut rewards_balance = NativeBalance(withdraw_response.total_rewards);
-    rewards_balance.normalize();
-
-    let total_rewards: Vec<String> = rewards_balance
-        .into_vec()
+    let staking_denom = deps.querier.query_bonded_denom()?;
+    let rewards_amount = withdraw_response
+        .total_rewards
         .iter()
-        .map(|coin| coin.to_string())
-        .collect();
+        .find(|coin| coin.denom == staking_denom)
+        .ok_or_else(|| StdError::generic_err("Could not find coin with staking denom"))?;
+
+    let staking_msg = StakingMsg::Delegate {
+        validator: validator_addr.to_string(),
+        amount: rewards_amount.clone(),
+    };
 
     let res = Response::new()
-        .add_attribute("method", "after_rewards_withdrawn")
+        .add_message(staking_msg)
+        .add_attribute("method", "stake_withdrawn_rewards")
         .add_attribute("records_num", withdraw_response.records_num.to_string())
-        .add_attribute("total_rewards", total_rewards.concat());
+        .add_attribute("staked_rewards", rewards_amount.to_string())
+        .add_attribute("validator", validator_addr);
 
     Ok(res)
 }
@@ -232,13 +265,14 @@ fn gov_vote(
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use cosmwasm_std::testing::{mock_env, mock_info};
+    use cosmwasm_std::{coins, from_binary, ContractResult, QueryResponse};
+
     use archway_bindings::testing::{mock_dependencies, mock_dependencies_with_balance};
     use archway_bindings::types::rewards::RewardsRecord;
     use archway_bindings::PageResponse;
 
-    use cosmwasm_std::testing::{mock_env, mock_info};
-    use cosmwasm_std::{coins, from_binary, ContractResult, QueryResponse};
+    use super::*;
 
     #[test]
     fn proper_initialization() {

contracts/increment/src/msg.rs

@@ -16,7 +16,7 @@ pub enum ExecuteMsg {
     Reset { count: i32 },
     UpdateRewardsAddress { rewards_address: Option<Addr> },
     SetFlatFee { amount: Uint128 },
-    WithdrawRewards {},
+    WithdrawRewards { stake_to_validator: Option<Addr> },
 }
 
 #[cw_serde]

contracts/increment/src/state.rs

@@ -2,7 +2,7 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 use cosmwasm_std::Addr;
-use cw_storage_plus::Item;
+use cw_storage_plus::{Item, Map};
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 pub struct State {
@@ -11,3 +11,5 @@ pub struct State {
 }
 
 pub const STATE: Item<State> = Item::new("state");
+// Map with key (block height, tx index) and value validator address
+pub const STAKE_REWARDS_VALIDATOR: Map<(u64, u32), Addr> = Map::new("stake_rewards_validator");