Skip to content

Commit

Permalink
collapse reward sends working
Browse files Browse the repository at this point in the history
  • Loading branch information
@DrPresident
DrPresident committed Jun 24, 2024
1 parent 015939e commit 2d74070
Show file tree
Hide file tree
Showing 2 changed files with 132 additions and 112 deletions.
3 changes: 3 additions & 0 deletions contracts/basic_staking/src/contract.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,8 @@ pub fn migrate(deps: DepsMut, env: Env, _msg: MigrateMsg) -> StdResult<Response>

let mut msgs = vec![];
let mut stake_is_reward = false;

// Register receive on all reward tokens
for token in reward_tokens {
if token.address == stake_token.address {
stake_is_reward = true;
Expand All @@ -231,6 +233,7 @@ pub fn migrate(deps: DepsMut, env: Env, _msg: MigrateMsg) -> StdResult<Response>
)?);
}

// Register if not done with rewards
if !stake_is_reward {
msgs.push(register_receive(
env.contract.code_hash.clone(),
Expand Down
241 changes: 129 additions & 112 deletions contracts/basic_staking/src/execute.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,10 @@ use shade_protocol::{
};

use crate::storage::*;
use std::cmp::{max, min};
use std::{
cmp::{max, min},
collections::HashMap,
};

pub fn update_config(
deps: DepsMut,
Expand Down Expand Up @@ -368,6 +371,85 @@ pub fn reward_pool_claim(

Ok(user_reward)
}
pub fn _do_claim(
storage: &mut dyn Storage,
mut response: Response,
mut reward_pools: Vec<RewardPoolInternal>,
user_staked: Uint128,
user: Addr,
) -> StdResult<(Response, Vec<RewardPoolInternal>)> {
let mut rewards: HashMap<Contract, Uint128> = HashMap::new();

for reward_pool in reward_pools.iter_mut() {
let reward_claimed = reward_pool_claim(storage, user.clone(), user_staked, &reward_pool)?;

reward_pool.claimed += reward_claimed;
rewards
.entry(reward_pool.token.clone())
.and_modify(|reward| *reward += reward_claimed)
.or_insert(reward_claimed);
}

for (token, reward) in rewards.into_iter() {
if reward.is_zero() {
continue;
}

response = response
.add_message(send_msg(user.clone(), reward, None, None, None, &token)?)
.add_attribute(token.address.to_string(), reward);
}

Ok((response, reward_pools))
}

pub fn _do_compound(
storage: &mut dyn Storage,
mut response: Response,
mut reward_pools: Vec<RewardPoolInternal>,
mut user_staked: Uint128,
mut total_staked: Uint128,
stake_token: Contract,
user: Addr,
) -> StdResult<(Response, Vec<RewardPoolInternal>, Uint128, Uint128)> {
let mut claim_rewards: HashMap<Contract, Uint128> = HashMap::new();

let mut compound_amount = Uint128::zero();

for reward_pool in reward_pools.iter_mut() {
let reward_claimed = reward_pool_claim(storage, user.clone(), user_staked, &reward_pool)?;
reward_pool.claimed += reward_claimed;

if reward_pool.token == stake_token {
// Compound stake_token rewards
compound_amount += reward_claimed;
} else {
// Claim non-stake_token rewards
claim_rewards
.entry(reward_pool.token.clone())
.and_modify(|reward| *reward += reward_claimed)
.or_insert(reward_claimed);
}
}

if compound_amount > Uint128::zero() {
response = response.add_attribute("compounded", compound_amount);
user_staked += compound_amount;
total_staked += compound_amount;
}

for (token, reward) in claim_rewards.into_iter() {
if reward.is_zero() {
continue;
}

response = response
.add_message(send_msg(user.clone(), reward, None, None, None, &token)?)
.add_attribute(token.address.to_string(), reward);
}

Ok((response, reward_pools, user_staked, total_staked))
}

pub fn claim(deps: DepsMut, env: Env, info: MessageInfo) -> StdResult<Response> {
let user_staked = USER_STAKED.load(deps.storage, info.sender.clone())?;
Expand All @@ -387,56 +469,34 @@ pub fn claim(deps: DepsMut, env: Env, info: MessageInfo) -> StdResult<Response>

let mut response = Response::new();

for reward_pool in reward_pools.iter_mut() {
let reward_claimed =
reward_pool_claim(deps.storage, info.sender.clone(), user_staked, &reward_pool)?;

reward_pool.claimed += reward_claimed;

if !reward_claimed.is_zero() {
response = response
.add_message(send_msg(
info.sender.clone(),
reward_claimed,
None,
None,
None,
&reward_pool.token,
)?)
.add_attribute(reward_pool.token.address.to_string(), reward_claimed);
}
}
(response, reward_pools) = _do_claim(
deps.storage,
response,
reward_pools,
user_staked,
info.sender,
)?;

REWARD_POOLS.save(deps.storage, &reward_pools)?;

Ok(response.set_data(to_binary(&ExecuteAnswer::Claim {
//claimed:
status: ResponseStatus::Success,
})?))
}

pub fn unbond(
deps: DepsMut,
mut deps: DepsMut,
env: Env,
info: MessageInfo,
amount: Uint128,
compound: bool,
) -> StdResult<Response> {
let config = CONFIG.load(deps.storage)?;

if amount.is_zero() {
return Err(StdError::generic_err("Must unbond non-zero amount"));
}

if let Some(mut user_staked) = USER_STAKED.may_load(deps.storage, info.sender.clone())? {
// if not compounding, check staked >= unbond amount
if !compound && user_staked < amount {
return Err(StdError::generic_err(format!(
"Cannot unbond {}, only {} staked",
amount, user_staked
)));
}

let now = env.block.time.seconds();

let mut total_staked = TOTAL_STAKED.load(deps.storage)?;
Expand All @@ -445,48 +505,39 @@ pub fn unbond(
update_rewards(env.clone(), &REWARD_POOLS.load(deps.storage)?, total_staked);

let stake_token = STAKE_TOKEN.load(deps.storage)?;
let mut compound_amount = Uint128::zero();

let mut response = Response::new();

// Claim/Compound rewards
for reward_pool in reward_pools.iter_mut() {
let reward_claimed =
reward_pool_claim(deps.storage, info.sender.clone(), user_staked, &reward_pool)?;
reward_pool.claimed += reward_claimed;

if compound && reward_pool.token == stake_token {
// Compound stake_token rewards
compound_amount += reward_claimed;
} else {
// Claim if not compound or not stake token rewards
response = response
.add_message(send_msg(
info.sender.clone(),
reward_claimed,
None,
None,
None,
&reward_pool.token,
)?)
.add_attribute(reward_pool.token.address.to_string(), reward_claimed);
}
if compound {
(response, reward_pools, user_staked, total_staked) = _do_compound(
deps.storage,
response,
reward_pools,
user_staked,
total_staked,
stake_token,
info.sender.clone(),
)?;
} else {
(response, reward_pools) = _do_claim(
deps.storage,
response,
reward_pools,
user_staked,
info.sender.clone(),
)?;
}

// if compounding, check staked + compounded >= unbond amount
if user_staked + compound_amount < amount {
if user_staked < amount {
return Err(StdError::generic_err(format!(
"Cannot unbond {}, only {} staked after compounding",
amount,
user_staked + compound_amount,
"Cannot unbond {}, only {} staked",
amount, user_staked,
)));
}
if compound_amount > Uint128::zero() {
response = response.add_attribute("compounded", compound_amount);
}

user_staked = (user_staked + compound_amount) - amount;
total_staked = (total_staked + compound_amount) - amount;
// Reduce by unbonding
user_staked -= amount;

TOTAL_STAKED.save(deps.storage, &total_staked)?;
USER_STAKED.save(deps.storage, info.sender.clone(), &user_staked)?;
Expand Down Expand Up @@ -603,58 +654,37 @@ pub fn withdraw(
pub fn compound(deps: DepsMut, env: Env, info: MessageInfo) -> StdResult<Response> {
let mut response = Response::new();

let user_staked = USER_STAKED
let mut user_staked = USER_STAKED
.may_load(deps.storage, info.sender.clone())?
.unwrap_or(Uint128::zero());

if user_staked.is_zero() {
return Err(StdError::generic_err("User has no stake"));
}
let pre_compound_user_staked = user_staked;

let total_staked = TOTAL_STAKED.load(deps.storage)?;
let mut total_staked = TOTAL_STAKED.load(deps.storage)?;
let mut reward_pools =
update_rewards(env.clone(), &REWARD_POOLS.load(deps.storage)?, total_staked);
let stake_token = STAKE_TOKEN.load(deps.storage)?;

let mut compound_amount = Uint128::zero();

for reward_pool in reward_pools.iter_mut() {
let reward_claimed =
reward_pool_claim(deps.storage, info.sender.clone(), user_staked, &reward_pool)?;
reward_pool.claimed += reward_claimed;

if reward_pool.token == stake_token {
// Compound stake_token rewards
compound_amount += reward_claimed;
} else {
// Claim non-stake_token rewards
response = response
.add_message(send_msg(
info.sender.clone(),
reward_claimed,
None,
None,
None,
&reward_pool.token,
)?)
.add_attribute(reward_pool.token.address.to_string(), reward_claimed);
}
}
REWARD_POOLS.save(deps.storage, &reward_pools)?;

if compound_amount > Uint128::zero() {
response = response.add_attribute("compounded", compound_amount);
}

USER_STAKED.save(
(response, reward_pools, user_staked, total_staked) = _do_compound(
deps.storage,
response,
reward_pools,
user_staked,
total_staked,
stake_token,
info.sender.clone(),
&(user_staked + compound_amount),
)?;
TOTAL_STAKED.save(deps.storage, &(total_staked + compound_amount))?;

REWARD_POOLS.save(deps.storage, &reward_pools)?;

USER_STAKED.save(deps.storage, info.sender, &user_staked)?;
TOTAL_STAKED.save(deps.storage, &(total_staked))?;

Ok(response.set_data(to_binary(&ExecuteAnswer::Compound {
compounded: compound_amount,
compounded: user_staked - pre_compound_user_staked,
status: ResponseStatus::Success,
})?))
}
Expand Down Expand Up @@ -696,15 +726,13 @@ pub fn end_reward_pool(

// Delete reward pool if it hasn't started
let deleted = if reward_pool.start > now {
println!("DELETING BEFORE START");
extract_amount = reward_pool.amount;
true
}
// Reward pool hasn't ended, trim off un-emitted tokens & edit pool to end now
else if reward_pool.end > now {
// remove rewards from now -> end
extract_amount = reward_pool.rate * (reward_pool.end - now) / Uint128::new(10u128.pow(18));
println!("EXTRACTING {}", extract_amount);
reward_pool.end = now;
reward_pool.amount -= extract_amount;

Expand Down Expand Up @@ -849,14 +877,12 @@ pub fn transfer_stake(

// Claim/Compound rewards for Sender
for reward_pool in reward_pools.iter_mut() {
println!("reward pool claim sender");
let reward_claimed = reward_pool_claim(
deps.storage,
info.sender.clone(),
sender_staked,
&reward_pool,
)?;
println!("POST reward pool claim sender");
reward_pool.claimed += reward_claimed;

if compound && reward_pool.token == stake_token {
Expand Down Expand Up @@ -885,8 +911,6 @@ pub fn transfer_stake(
)));
}

println!("sender compound amount {}", sender_compound_amount);

if sender_compound_amount > Uint128::zero() {
response = response.add_attribute("compounded", sender_compound_amount);
}
Expand All @@ -905,7 +929,6 @@ pub fn transfer_stake(

// Claim rewards for Receiver (no compound)
for reward_pool in reward_pools.iter_mut() {
println!("reward pool claim recipient");
let reward_claimed = reward_pool_claim(
deps.storage,
recipient.clone(),
Expand All @@ -914,12 +937,6 @@ pub fn transfer_stake(
)?;
reward_pool.claimed += reward_claimed;

// Claim if not compound or not stake token rewards
println!(
"SENDING RECIPIETN REWARD {} {}",
reward_claimed,
recipient.clone()
);
response = response.add_message(send_msg(
recipient.clone(),
reward_claimed,
Expand Down

0 comments on commit 2d74070

Please sign in to comment.