diff --git a/contracts/treasurechest-contract/src/contract.rs b/contracts/treasurechest-contract/src/contract.rs index 3dedd08..9f899ff 100644 --- a/contracts/treasurechest-contract/src/contract.rs +++ b/contracts/treasurechest-contract/src/contract.rs @@ -73,7 +73,7 @@ pub fn execute( ExecuteMsg::ChangeTokenFactory { token_factory_type, } => change_token_factory(deps, info.sender, &token_factory_type), - ExecuteMsg::ReturnDust {} => return_dust(deps, env, info.sender), + ExecuteMsg::ReturnDust {limit} => return_dust(deps, env, info.sender, limit), } } diff --git a/contracts/treasurechest-contract/src/executions.rs b/contracts/treasurechest-contract/src/executions.rs index 3ceac55..131d66d 100644 --- a/contracts/treasurechest-contract/src/executions.rs +++ b/contracts/treasurechest-contract/src/executions.rs @@ -1,9 +1,13 @@ use std::{ops::Mul, str::FromStr}; +use std::collections::HashMap; + +const DEFAULT_SIZE:u32 = 20; use cosmwasm_std::{ Addr, BankMsg, Coin, CosmosMsg, DepsMut, Env, Event, MessageInfo, Order, Response, StdResult, Uint128, }; + use treasurechest::{errors::ContractError, tf::tokenfactory::TokenFactoryType}; use crate::state::{CONFIG, TOTAL_REWARDS}; @@ -76,27 +80,32 @@ pub fn change_token_factory( .add_attribute("action", "treasurechest/change_token_factory")) } -pub fn return_dust(deps: DepsMut, env: Env, sender: Addr) -> Result { +pub fn return_dust(deps: DepsMut, env: Env, sender: Addr, limit: Option) -> Result { cw_ownable::assert_owner(deps.storage, &sender)?; let config = CONFIG.load(deps.storage)?; let balances = deps .querier .query_all_balances(env.contract.address)? .into_iter() - .filter(|x| x.denom != config.denom) - .collect::>(); + .filter(|x| x.denom != config.denom).map(|coin| (coin.denom,coin.amount)) + .collect::>(); let mut balances_out = vec![]; - for entry in balances { - if let Some(one_amt) = TOTAL_REWARDS.may_load(deps.storage, entry.denom.clone())? { - if one_amt.to_uint_floor() > entry.amount { - TOTAL_REWARDS.remove(deps.storage, entry.denom.clone()); - balances_out.push(entry) + let rewards = TOTAL_REWARDS.range(deps.storage, None,None,Order::Ascending).take(limit.unwrap_or(DEFAULT_SIZE).try_into()?).collect::>>()?; + for reward in rewards { + let reward_amt = reward.1.to_uint_floor(); + if let Some(token_balance) = balances.get( &reward.0) { + if &reward_amt > token_balance{ + TOTAL_REWARDS.remove(deps.storage, reward.0.clone()); + balances_out.push(Coin{denom: reward.0, amount: token_balance.clone()}) } } } // balances_out should only contain the dust now. // TOTAL rewards should no longer show that token + if balances_out.is_empty() { + return Ok(Response::new().add_attribute("action", "treasurechest/return_dust").add_attribute("dust","no-dust")) + } let msg = CosmosMsg::Bank(BankMsg::Send { to_address: sender.to_string(), diff --git a/integration/src/treasurechest/contract.rs b/integration/src/treasurechest/contract.rs index f0b23c1..b031f07 100644 --- a/integration/src/treasurechest/contract.rs +++ b/integration/src/treasurechest/contract.rs @@ -42,8 +42,8 @@ impl TreasureChestContract { }) } - pub fn return_dust(&self) -> StdResult { - self.call(ExecuteMsg::ReturnDust {}) + pub fn return_dust(&self, limit:Option) -> StdResult { + self.call(ExecuteMsg::ReturnDust {limit}) } pub fn config(&self, app: &App) -> StdResult { diff --git a/integration/src/treasurechest/tests.rs b/integration/src/treasurechest/tests.rs index 77d78b8..ffba0e7 100644 --- a/integration/src/treasurechest/tests.rs +++ b/integration/src/treasurechest/tests.rs @@ -335,7 +335,7 @@ pub fn test_withdraw() -> AnyResult<()> { assert_eq!(state.outstanding, Uint128::zero()); app.update_block(|f| f.height += 1); - let res_admin = app.execute(Addr::unchecked(ADMIN), chest.return_dust()?)?; + let res_admin = app.execute(Addr::unchecked(ADMIN), chest.return_dust(None)?)?; for e in get_events("transfer", &res_admin.events) { if let Some(amount) = get_attribute("amount", &e.attributes) { let coins = amount.split(',').map(|x| x.to_string()).collect::>(); diff --git a/packages/treasurechest/src/chest.rs b/packages/treasurechest/src/chest.rs index f176709..6507b48 100644 --- a/packages/treasurechest/src/chest.rs +++ b/packages/treasurechest/src/chest.rs @@ -24,7 +24,7 @@ pub enum ExecuteMsg { /// Withdraw pending rewards Withdraw {}, /// If balance is below >1< tickets worth (ADMIN only) - ReturnDust {}, + ReturnDust { limit: Option}, /// change token factory type (ADMIN only) ChangeTokenFactory { token_factory_type: String, diff --git a/packages/treasurechest/src/errors.rs b/packages/treasurechest/src/errors.rs index 8aa4f4a..8811e1f 100644 --- a/packages/treasurechest/src/errors.rs +++ b/packages/treasurechest/src/errors.rs @@ -1,3 +1,4 @@ +use std::num::TryFromIntError; use cosmwasm_std::{CheckedFromRatioError, DivideByZeroError, OverflowError, StdError}; use cw_ownable::OwnershipError; use thiserror::Error; @@ -10,6 +11,8 @@ pub enum ContractError { Ownership(#[from] OwnershipError), #[error("{0}")] OverflowError(#[from] OverflowError), + #[error("{0}")] + TryFromIntError(#[from] TryFromIntError), #[error("{0}")] DivideByZeroError(#[from] DivideByZeroError),