update junoswap migration contracts so users can access unclaimed rewards

contracts/junoswap-staking/src/contract.rs

@@ -1,3 +1,4 @@
+use std::cmp::min;
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
@@ -11,6 +12,7 @@ use wyndex::asset::{Asset, AssetInfo};
 
 use crate::error::ContractError;
 use crate::msg::{ExecuteMsg, MigrateMsg, QueryMsg};
+use crate::msg::QueryMsg::TotalStakedAtHeight;
 use crate::state::{MigrateConfig, MigrateStakersConfig, DESTINATION, MIGRATION};
 
 // this is the contract we are migrating from
@@ -32,13 +34,26 @@ pub fn instantiate(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn query(_deps: Deps, _env: Env, msg: QueryMsg) -> Result<Binary, ContractError> {
+    // Block before the migration was started
+    let migration_height: u64 = 6634356;
     match msg {
         QueryMsg::MigrationFinished {} => {
             let no_stakers = stake_cw20::state::STAKED_BALANCES
                 .keys(_deps.storage, None, None, Order::Ascending)
                 .next()
                 .is_none();
             Ok(to_binary(&no_stakers)?)
+        },
+        QueryMsg::TotalStakedAtHeight { height } => {
+            let new_height = min(migration_height,height.unwrap_or(u64::MAX));
+            let staked_total = stake_cw20::state::STAKED_TOTAL.may_load_at_height(_deps.storage, new_height)?.unwrap_or_default();
+            Ok(to_binary(&stake_cw20::msg::TotalStakedAtHeightResponse{ total: staked_total, height: new_height })?)
+        },
+        QueryMsg::StakedBalanceAtHeight { address, height } => {
+            let address = _deps.api.addr_validate(&address)?;
+            let new_height = min(migration_height,height.unwrap_or(u64::MAX));
+            let balance = stake_cw20::state::STAKED_BALANCES.may_load_at_height(_deps.storage, &address, new_height)?.unwrap_or_default();
+            Ok(to_binary(&stake_cw20::msg::StakedBalanceAtHeightResponse{ balance, height: new_height })?)
         }
     }
 }

contracts/junoswap-staking/src/msg.rs

@@ -1,11 +1,20 @@
 use cosmwasm_schema::{cw_serde, QueryResponses};
+use cosmwasm_std::Uint128;
+use stake_cw20::msg::{StakedBalanceAtHeightResponse, TotalStakedAtHeightResponse};
 
 #[cw_serde]
 #[derive(QueryResponses)]
 pub enum QueryMsg {
     /// Checks whether all stakers have been migrated
     #[returns(bool)]
     MigrationFinished {},
+    #[returns(StakedBalanceAtHeightResponse)]
+    StakedBalanceAtHeight {
+        address: String,
+        height: Option<u64>,
+    },
+    #[returns(TotalStakedAtHeightResponse)]
+    TotalStakedAtHeight { height: Option<u64> }
 }
 
 /// For existing contract, we need to specify which pool it can be withdrawn into

contracts/raw-migration/src/contract.rs

@@ -1,3 +1,4 @@
+use std::cmp::min;
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
@@ -35,13 +36,26 @@ pub fn instantiate(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn query(_deps: Deps, _env: Env, msg: QueryMsg) -> Result<Binary, ContractError> {
+    // Block before the migration was started
+    let migration_height: u64 = 6634356;
     match msg {
         QueryMsg::MigrationFinished {} => {
             let no_stakers = stake_cw20::state::STAKED_BALANCES
                 .keys(_deps.storage, None, None, Order::Ascending)
                 .next()
                 .is_none();
             Ok(to_binary(&no_stakers)?)
+        },
+        QueryMsg::TotalStakedAtHeight { height } => {
+            let new_height = min(migration_height,height.unwrap_or(u64::MAX));
+            let staked_total = stake_cw20::state::STAKED_TOTAL.may_load_at_height(_deps.storage, new_height)?.unwrap_or_default();
+            Ok(to_binary(&stake_cw20::msg::TotalStakedAtHeightResponse{ total: staked_total, height: new_height })?)
+        },
+        QueryMsg::StakedBalanceAtHeight { address, height } => {
+            let address = _deps.api.addr_validate(&address)?;
+            let new_height = min(migration_height,height.unwrap_or(u64::MAX));
+            let balance = stake_cw20::state::STAKED_BALANCES.may_load_at_height(_deps.storage, &address, new_height)?.unwrap_or_default();
+            Ok(to_binary(&stake_cw20::msg::StakedBalanceAtHeightResponse{ balance, height: new_height })?)
         }
     }
 }

contracts/raw-migration/src/msg.rs

@@ -1,5 +1,6 @@
 use cosmwasm_schema::{cw_serde, QueryResponses};
 use cosmwasm_std::Decimal;
+use stake_cw20::msg::{StakedBalanceAtHeightResponse, TotalStakedAtHeightResponse};
 
 #[cw_serde]
 pub struct InstantiateMsg {}
@@ -10,6 +11,13 @@ pub enum QueryMsg {
     /// Checks whether all stakers have been migrated
     #[returns(bool)]
     MigrationFinished {},
+    #[returns(StakedBalanceAtHeightResponse)]
+    StakedBalanceAtHeight {
+        address: String,
+        height: Option<u64>,
+    },
+    #[returns(TotalStakedAtHeightResponse)]
+    TotalStakedAtHeight { height: Option<u64> }
 }
 
 /// For existing contract, we need to specify which pool it can be withdrawn into