diff --git a/contracts/alliance-hub/src/query.rs b/contracts/alliance-hub/src/query.rs index 2b455c5..004a84b 100644 --- a/contracts/alliance-hub/src/query.rs +++ b/contracts/alliance-hub/src/query.rs @@ -5,11 +5,12 @@ use alliance_protocol::alliance_protocol::{ }; #[cfg(not(feature = "library"))] use cosmwasm_std::entry_point; -use cosmwasm_std::{to_json_binary, Binary, Deps, Env, Order, StdResult, Uint128}; +use cosmwasm_std::{to_json_binary, Binary, Deps, Env, Order, StdError, StdResult, Uint128}; use cw_asset::AssetInfo; use std::cmp::min; use std::collections::HashMap; use ve3_shared::constants::SECONDS_PER_YEAR; +use ve3_shared::msgs_asset_staking::AssetConfigRuntime; use crate::state::{ ASSET_CONFIG, ASSET_REWARD_DISTRIBUTION, ASSET_REWARD_RATE, CONFIG, SHARES, @@ -23,11 +24,11 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> { QueryMsg::Validators {} => get_validators(deps)?, QueryMsg::WhitelistedAssets {} => get_whitelisted_assets(deps)?, QueryMsg::RewardDistribution {} => get_rewards_distribution(deps)?, - QueryMsg::StakedBalance(asset_query) => get_staked_balance(deps, asset_query)?, + QueryMsg::StakedBalance(asset_query) => get_staked_balance(deps, env, asset_query)?, QueryMsg::PendingRewards(asset_query) => get_pending_rewards(deps, asset_query)?, QueryMsg::AllStakedBalances(query) => get_all_staked_balances(deps, env, query)?, QueryMsg::AllPendingRewards(query) => get_all_pending_rewards(deps, query)?, - QueryMsg::TotalStakedBalances {} => get_total_staked_balances(deps)?, + QueryMsg::TotalStakedBalances {} => get_total_staked_balances(deps, env)?, }) } @@ -62,15 +63,17 @@ fn get_rewards_distribution(deps: Deps) -> StdResult<Binary> { to_json_binary(&asset_rewards_distr) } -fn get_staked_balance(deps: Deps, asset_query: AssetQuery) -> StdResult<Binary> { +fn get_staked_balance(deps: Deps, env: Env, asset_query: AssetQuery) -> StdResult<Binary> { let addr = deps.api.addr_validate(&asset_query.address)?; let key = (addr, &asset_query.asset); let user_shares = SHARES.load(deps.storage, key)?; let (balance, shares) = TOTAL_BALANCES_SHARES.load(deps.storage, &asset_query.asset)?; - let config = ASSET_CONFIG + let mut config = ASSET_CONFIG .may_load(deps.storage, &asset_query.asset)? .unwrap_or_default(); + compute_take_amount(&env, balance, &mut config)?; + to_json_binary(&StakedBalanceRes { asset: asset_query.asset, balance: compute_balance_amount(shares, user_shares, balance.checked_sub(config.taken)?), @@ -79,6 +82,25 @@ fn get_staked_balance(deps: Deps, asset_query: AssetQuery) -> StdResult<Binary> }) } +fn compute_take_amount( + env: &Env, + balance: Uint128, + config: &mut AssetConfigRuntime, +) -> Result<(), StdError> { + if config.last_taken_s != 0 { + let take_diff_s = Uint128::new((env.block.time.seconds() - config.last_taken_s).into()); + let relevant_balance = balance.saturating_sub(config.taken); + let take_amount = config.yearly_take_rate + * relevant_balance.multiply_ratio( + min(take_diff_s, Uint128::new(SECONDS_PER_YEAR.into())), + SECONDS_PER_YEAR, + ); + config.last_taken_s = env.block.time.seconds(); + config.taken = config.taken.checked_add(take_amount)? + }; + Ok(()) +} + fn get_pending_rewards(deps: Deps, asset_query: AssetQuery) -> StdResult<Binary> { let config = CONFIG.load(deps.storage)?; let addr = deps.api.addr_validate(&asset_query.address)?; @@ -117,17 +139,7 @@ fn get_all_staked_balances( let (balance, shares) = TOTAL_BALANCES_SHARES.load(deps.storage, &asset_info)?; let mut config = ASSET_CONFIG.load(deps.storage, &asset_info)?; - if config.last_taken_s != 0 { - let take_diff_s = Uint128::new((env.block.time.seconds() - config.last_taken_s).into()); - let relevant_balance = balance.saturating_sub(config.taken); - let take_amount = config.yearly_take_rate - * relevant_balance.multiply_ratio( - min(take_diff_s, Uint128::new(SECONDS_PER_YEAR.into())), - SECONDS_PER_YEAR, - ); - config.last_taken_s = env.block.time.seconds(); - config.taken = config.taken.checked_add(take_amount)? - }; + compute_take_amount(&env, balance, &mut config)?; // Append the request res.push(StakedBalanceRes { @@ -170,16 +182,18 @@ fn get_all_pending_rewards(deps: Deps, query: AllPendingRewardsQuery) -> StdResu to_json_binary(&all_pending_rewards?) } -fn get_total_staked_balances(deps: Deps) -> StdResult<Binary> { +fn get_total_staked_balances(deps: Deps, env: Env) -> StdResult<Binary> { let total_staked_balances: StdResult<Vec<StakedBalanceRes>> = TOTAL_BALANCES_SHARES .range(deps.storage, None, None, Order::Ascending) .map(|total_balance| -> StdResult<StakedBalanceRes> { let (asset, (balance, shares)) = total_balance?; - let config = ASSET_CONFIG.load(deps.storage, &asset)?; + let mut config = ASSET_CONFIG.load(deps.storage, &asset)?; + + compute_take_amount(&env, balance, &mut config)?; Ok(StakedBalanceRes { asset, - balance, + balance: compute_balance_amount(shares, shares, balance.checked_sub(config.taken)?), shares, config, })