From aa5de2684659f930cfcba95c2d0de9ecebbfc7d4 Mon Sep 17 00:00:00 2001
From: gangov <gangov1@gmail.com>
Date: Thu, 30 Nov 2023 20:33:23 +0200
Subject: [PATCH 01/10] Stake: first implementation at checking the rewards
 before unbond; Multihop: formatting

---
 contracts/stake/src/contract.rs     | 68 +++++------------------------
 contracts/stake/src/distribution.rs | 66 +++++++++++++++++++++++++++-
 2 files changed, 74 insertions(+), 60 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index c6f357dbb..010cb30e7 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -2,9 +2,9 @@ use soroban_sdk::{contract, contractimpl, contractmeta, log, vec, Address, Env,
 
 use crate::{
     distribution::{
-        calculate_annualized_payout, get_distribution, get_reward_curve, get_withdraw_adjustment,
-        save_distribution, save_reward_curve, save_withdraw_adjustment, update_rewards,
-        withdrawable_rewards, Distribution, SHARES_SHIFT,
+        calculate_annualized_payout, do_withdraw_rewards, get_distribution, get_reward_curve,
+        get_withdraw_adjustment, get_withdrawable_rewards, save_distribution, save_reward_curve,
+        save_withdraw_adjustment, update_rewards, withdrawable_rewards, Distribution, SHARES_SHIFT,
     },
     msg::{
         AnnualizedReward, AnnualizedRewardsResponse, ConfigResponse, StakedResponse,
@@ -181,6 +181,11 @@ impl StakingTrait for Staking {
 
         let config = get_config(&env);
 
+        let withdrawable_rewards = get_withdrawable_rewards(&env, &sender);
+        if !withdrawable_rewards.rewards.is_empty() {
+            do_withdraw_rewards(&env, &sender);
+        }
+
         let mut stakes = get_stakes(&env, &sender);
         remove_stake(&mut stakes.stakes, stake_amount, stake_timestamp);
         stakes.total_stake -= stake_amount as u128;
@@ -280,43 +285,7 @@ impl StakingTrait for Staking {
     }
 
     fn withdraw_rewards(env: Env, sender: Address) {
-        env.events().publish(("withdraw_rewards", "user"), &sender);
-
-        for distribution_address in get_distributions(&env) {
-            // get distribution data for the given reward
-            let mut distribution = get_distribution(&env, &distribution_address);
-            // get withdraw adjustment for the given distribution
-            let mut withdraw_adjustment =
-                get_withdraw_adjustment(&env, &sender, &distribution_address);
-            // calculate current reward amount given the distribution and subtracting withdraw
-            // adjustments
-            let reward_amount =
-                withdrawable_rewards(&env, &sender, &distribution, &withdraw_adjustment);
-
-            if reward_amount == 0 {
-                continue;
-            }
-
-            withdraw_adjustment.withdrawn_rewards += reward_amount;
-            distribution.withdrawable_total -= reward_amount;
-
-            save_distribution(&env, &distribution_address, &distribution);
-            save_withdraw_adjustment(&env, &sender, &distribution_address, &withdraw_adjustment);
-
-            let reward_token_client = token_contract::Client::new(&env, &distribution_address);
-            reward_token_client.transfer(
-                &env.current_contract_address(),
-                &sender,
-                &(reward_amount as i128),
-            );
-
-            env.events().publish(
-                ("withdraw_rewards", "reward_token"),
-                &reward_token_client.address,
-            );
-            env.events()
-                .publish(("withdraw_rewards", "reward_amount"), reward_amount as i128);
-        }
+        do_withdraw_rewards(&env, &sender)
     }
 
     fn fund_distribution(
@@ -442,24 +411,7 @@ impl StakingTrait for Staking {
     }
 
     fn query_withdrawable_rewards(env: Env, user: Address) -> WithdrawableRewardsResponse {
-        // iterate over all distributions and calculate withdrawable rewards
-        let mut rewards = vec![&env];
-        for distribution_address in get_distributions(&env) {
-            // get distribution data for the given reward
-            let distribution = get_distribution(&env, &distribution_address);
-            // get withdraw adjustment for the given distribution
-            let withdraw_adjustment = get_withdraw_adjustment(&env, &user, &distribution_address);
-            // calculate current reward amount given the distribution and subtracting withdraw
-            // adjustments
-            let reward_amount =
-                withdrawable_rewards(&env, &user, &distribution, &withdraw_adjustment);
-            rewards.push_back(WithdrawableReward {
-                reward_address: distribution_address,
-                reward_amount,
-            });
-        }
-
-        WithdrawableRewardsResponse { rewards }
+        get_withdrawable_rewards(&env, &user)
     }
 
     fn query_distributed_rewards(env: Env, asset: Address) -> u128 {
diff --git a/contracts/stake/src/distribution.rs b/contracts/stake/src/distribution.rs
index 9cf87b047..0a43a983b 100644
--- a/contracts/stake/src/distribution.rs
+++ b/contracts/stake/src/distribution.rs
@@ -1,9 +1,12 @@
-use soroban_sdk::{contracttype, Address, Env};
+use soroban_sdk::{contracttype, vec, Address, Env};
 
 use curve::Curve;
 use decimal::Decimal;
 
-use crate::storage::get_stakes;
+use crate::{
+    msg::{WithdrawableReward, WithdrawableRewardsResponse},
+    storage::{get_stakes, utils::get_distributions},
+};
 
 /// How much points is the worth of single token in rewards distribution.
 /// The scaling is performed to have better precision of fixed point division.
@@ -216,3 +219,62 @@ pub fn calculate_annualized_payout(reward_curve: Option<Curve>, now: u64) -> Dec
         None => Decimal::zero(),
     }
 }
+
+pub fn get_withdrawable_rewards(env: &Env, user: &Address) -> WithdrawableRewardsResponse {
+    // iterate over all distributions and calculate withdrawable rewards
+    let mut rewards = vec![&env];
+    for distribution_address in get_distributions(&env) {
+        // get distribution data for the given reward
+        let distribution = get_distribution(&env, &distribution_address);
+        // get withdraw adjustment for the given distribution
+        let withdraw_adjustment = get_withdraw_adjustment(&env, &user, &distribution_address);
+        // calculate current reward amount given the distribution and subtracting withdraw
+        // adjustments
+        let reward_amount = withdrawable_rewards(&env, &user, &distribution, &withdraw_adjustment);
+        rewards.push_back(WithdrawableReward {
+            reward_address: distribution_address,
+            reward_amount,
+        });
+    }
+
+    WithdrawableRewardsResponse { rewards }
+}
+
+pub fn do_withdraw_rewards(env: &Env, sender: &Address) {
+    env.events().publish(("withdraw_rewards", "user"), sender);
+
+    for distribution_address in get_distributions(&env) {
+        // get distribution data for the given reward
+        let mut distribution = get_distribution(&env, &distribution_address);
+        // get withdraw adjustment for the given distribution
+        let mut withdraw_adjustment = get_withdraw_adjustment(&env, &sender, &distribution_address);
+        // calculate current reward amount given the distribution and subtracting withdraw
+        // adjustments
+        let reward_amount =
+            withdrawable_rewards(&env, &sender, &distribution, &withdraw_adjustment);
+
+        if reward_amount == 0 {
+            continue;
+        }
+
+        withdraw_adjustment.withdrawn_rewards += reward_amount;
+        distribution.withdrawable_total -= reward_amount;
+
+        save_distribution(&env, &distribution_address, &distribution);
+        save_withdraw_adjustment(&env, &sender, &distribution_address, &withdraw_adjustment);
+
+        let reward_token_client = token_contract::Client::new(&env, &distribution_address);
+        reward_token_client.transfer(
+            &env.current_contract_address(),
+            &sender,
+            &(reward_amount as i128),
+        );
+
+        env.events().publish(
+            ("withdraw_rewards", "reward_token"),
+            &reward_token_client.address,
+        );
+        env.events()
+            .publish(("withdraw_rewards", "reward_amount"), reward_amount as i128);
+    }
+}

From 6c762feac57244f017e61de7298090cb8885951a Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Tue, 5 Mar 2024 17:14:34 +0200
Subject: [PATCH 02/10] stake: lints

---
 contracts/stake/src/contract.rs     |  6 +++---
 contracts/stake/src/distribution.rs | 25 ++++++++++++-------------
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index 010cb30e7..a3580df01 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -3,12 +3,12 @@ use soroban_sdk::{contract, contractimpl, contractmeta, log, vec, Address, Env,
 use crate::{
     distribution::{
         calculate_annualized_payout, do_withdraw_rewards, get_distribution, get_reward_curve,
-        get_withdraw_adjustment, get_withdrawable_rewards, save_distribution, save_reward_curve,
-        save_withdraw_adjustment, update_rewards, withdrawable_rewards, Distribution, SHARES_SHIFT,
+        get_withdrawable_rewards, save_distribution, save_reward_curve, update_rewards,
+        Distribution, SHARES_SHIFT,
     },
     msg::{
         AnnualizedReward, AnnualizedRewardsResponse, ConfigResponse, StakedResponse,
-        WithdrawableReward, WithdrawableRewardsResponse,
+        WithdrawableRewardsResponse,
     },
     storage::{
         get_config, get_stakes, save_config, save_stakes,
diff --git a/contracts/stake/src/distribution.rs b/contracts/stake/src/distribution.rs
index 0a43a983b..66112a4ce 100644
--- a/contracts/stake/src/distribution.rs
+++ b/contracts/stake/src/distribution.rs
@@ -223,14 +223,14 @@ pub fn calculate_annualized_payout(reward_curve: Option<Curve>, now: u64) -> Dec
 pub fn get_withdrawable_rewards(env: &Env, user: &Address) -> WithdrawableRewardsResponse {
     // iterate over all distributions and calculate withdrawable rewards
     let mut rewards = vec![&env];
-    for distribution_address in get_distributions(&env) {
+    for distribution_address in get_distributions(env) {
         // get distribution data for the given reward
-        let distribution = get_distribution(&env, &distribution_address);
+        let distribution = get_distribution(env, &distribution_address);
         // get withdraw adjustment for the given distribution
-        let withdraw_adjustment = get_withdraw_adjustment(&env, &user, &distribution_address);
+        let withdraw_adjustment = get_withdraw_adjustment(env, user, &distribution_address);
         // calculate current reward amount given the distribution and subtracting withdraw
         // adjustments
-        let reward_amount = withdrawable_rewards(&env, &user, &distribution, &withdraw_adjustment);
+        let reward_amount = withdrawable_rewards(env, user, &distribution, &withdraw_adjustment);
         rewards.push_back(WithdrawableReward {
             reward_address: distribution_address,
             reward_amount,
@@ -243,15 +243,14 @@ pub fn get_withdrawable_rewards(env: &Env, user: &Address) -> WithdrawableReward
 pub fn do_withdraw_rewards(env: &Env, sender: &Address) {
     env.events().publish(("withdraw_rewards", "user"), sender);
 
-    for distribution_address in get_distributions(&env) {
+    for distribution_address in get_distributions(env) {
         // get distribution data for the given reward
-        let mut distribution = get_distribution(&env, &distribution_address);
+        let mut distribution = get_distribution(env, &distribution_address);
         // get withdraw adjustment for the given distribution
-        let mut withdraw_adjustment = get_withdraw_adjustment(&env, &sender, &distribution_address);
+        let mut withdraw_adjustment = get_withdraw_adjustment(env, sender, &distribution_address);
         // calculate current reward amount given the distribution and subtracting withdraw
         // adjustments
-        let reward_amount =
-            withdrawable_rewards(&env, &sender, &distribution, &withdraw_adjustment);
+        let reward_amount = withdrawable_rewards(env, sender, &distribution, &withdraw_adjustment);
 
         if reward_amount == 0 {
             continue;
@@ -260,13 +259,13 @@ pub fn do_withdraw_rewards(env: &Env, sender: &Address) {
         withdraw_adjustment.withdrawn_rewards += reward_amount;
         distribution.withdrawable_total -= reward_amount;
 
-        save_distribution(&env, &distribution_address, &distribution);
-        save_withdraw_adjustment(&env, &sender, &distribution_address, &withdraw_adjustment);
+        save_distribution(env, &distribution_address, &distribution);
+        save_withdraw_adjustment(env, sender, &distribution_address, &withdraw_adjustment);
 
-        let reward_token_client = token_contract::Client::new(&env, &distribution_address);
+        let reward_token_client = crate::token_contract::Client::new(env, &distribution_address);
         reward_token_client.transfer(
             &env.current_contract_address(),
-            &sender,
+            sender,
             &(reward_amount as i128),
         );
 

From 341f00c44920d70dec2430e50a139bc2abdec556 Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Wed, 6 Mar 2024 19:05:10 +0200
Subject: [PATCH 03/10] stake: refactors unbond and adds test

---
 contracts/stake/src/contract.rs     | 46 +++++++++++++++++++++---
 contracts/stake/src/distribution.rs |  2 +-
 contracts/stake/src/tests/bond.rs   | 54 +++++++++++++++++++++++++++++
 3 files changed, 96 insertions(+), 6 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index a3580df01..b79123aaa 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -2,9 +2,9 @@ use soroban_sdk::{contract, contractimpl, contractmeta, log, vec, Address, Env,
 
 use crate::{
     distribution::{
-        calculate_annualized_payout, do_withdraw_rewards, get_distribution, get_reward_curve,
-        get_withdrawable_rewards, save_distribution, save_reward_curve, update_rewards,
-        Distribution, SHARES_SHIFT,
+        calculate_annualized_payout, get_distribution, get_reward_curve, get_withdraw_adjustment,
+        get_withdrawable_rewards, save_distribution, save_reward_curve, save_withdraw_adjustment,
+        update_rewards, withdraw_rewards, withdrawable_rewards, Distribution, SHARES_SHIFT,
     },
     msg::{
         AnnualizedReward, AnnualizedRewardsResponse, ConfigResponse, StakedResponse,
@@ -183,7 +183,7 @@ impl StakingTrait for Staking {
 
         let withdrawable_rewards = get_withdrawable_rewards(&env, &sender);
         if !withdrawable_rewards.rewards.is_empty() {
-            do_withdraw_rewards(&env, &sender);
+            withdraw_rewards(&env, &sender);
         }
 
         let mut stakes = get_stakes(&env, &sender);
@@ -285,7 +285,43 @@ impl StakingTrait for Staking {
     }
 
     fn withdraw_rewards(env: Env, sender: Address) {
-        do_withdraw_rewards(&env, &sender)
+        env.events().publish(("withdraw_rewards", "user"), &sender);
+
+        for distribution_address in get_distributions(&env) {
+            // get distribution data for the given reward
+            let mut distribution = get_distribution(&env, &distribution_address);
+            // get withdraw adjustment for the given distribution
+            let mut withdraw_adjustment =
+                get_withdraw_adjustment(&env, &sender, &distribution_address);
+            // calculate current reward amount given the distribution and subtracting withdraw
+            // adjustments
+            let reward_amount =
+                withdrawable_rewards(&env, &sender, &distribution, &withdraw_adjustment);
+
+            if reward_amount == 0 {
+                continue;
+            }
+
+            withdraw_adjustment.withdrawn_rewards += reward_amount;
+            distribution.withdrawable_total -= reward_amount;
+
+            save_distribution(&env, &distribution_address, &distribution);
+            save_withdraw_adjustment(&env, &sender, &distribution_address, &withdraw_adjustment);
+
+            let reward_token_client = token_contract::Client::new(&env, &distribution_address);
+            reward_token_client.transfer(
+                &env.current_contract_address(),
+                &sender,
+                &(reward_amount as i128),
+            );
+
+            env.events().publish(
+                ("withdraw_rewards", "reward_token"),
+                &reward_token_client.address,
+            );
+            env.events()
+                .publish(("withdraw_rewards", "reward_amount"), reward_amount as i128);
+        }
     }
 
     fn fund_distribution(
diff --git a/contracts/stake/src/distribution.rs b/contracts/stake/src/distribution.rs
index 66112a4ce..ba0c54935 100644
--- a/contracts/stake/src/distribution.rs
+++ b/contracts/stake/src/distribution.rs
@@ -240,7 +240,7 @@ pub fn get_withdrawable_rewards(env: &Env, user: &Address) -> WithdrawableReward
     WithdrawableRewardsResponse { rewards }
 }
 
-pub fn do_withdraw_rewards(env: &Env, sender: &Address) {
+pub fn withdraw_rewards(env: &Env, sender: &Address) {
     env.events().publish(("withdraw_rewards", "user"), sender);
 
     for distribution_address in get_distributions(env) {
diff --git a/contracts/stake/src/tests/bond.rs b/contracts/stake/src/tests/bond.rs
index 0ad020fe3..2ac6440cc 100644
--- a/contracts/stake/src/tests/bond.rs
+++ b/contracts/stake/src/tests/bond.rs
@@ -226,3 +226,57 @@ fn unbond_wrong_user_stake_not_found() {
 
     staking.unbond(&user2, &10_000, &2_000);
 }
+
+#[test]
+fn pay_rewards_during_unbond() {
+    const AMOUNT: i128 = 100_000;
+    const WITHDRAW_TIMESTAMP: u64 = 2_000;
+
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let admin = Address::generate(&env);
+    let user = Address::generate(&env);
+    let manager = Address::generate(&env);
+    let owner = Address::generate(&env);
+
+    let lp_token = deploy_token_contract(&env, &admin);
+    let staking = deploy_staking_contract(&env, admin.clone(), &lp_token.address, &manager, &owner);
+
+    lp_token.mint(&user, &AMOUNT);
+    lp_token.mint(&owner, &AMOUNT);
+
+    staking.create_distribution_flow(&owner, &lp_token.address);
+    staking.fund_distribution(&owner, &0u64, &10_000u64, &lp_token.address, &AMOUNT);
+
+    env.ledger().with_mut(|li| {
+        li.timestamp = WITHDRAW_TIMESTAMP;
+    });
+    staking.bond(&user, &10_000);
+
+    env.ledger().with_mut(|li| {
+        li.timestamp = 4_000;
+    });
+
+    // user hasn't unbonded yet, no rewards to withdraw
+    assert_eq!(
+        staking
+            .query_withdrawable_rewards(&user)
+            .rewards
+            .iter()
+            .map(|reward| reward.reward_amount)
+            .sum::<u128>(),
+        0
+    );
+    assert_eq!(
+        staking.query_undistributed_rewards(&lp_token.address),
+        110_000
+    );
+    staking.unbond(&user, &10_000, &WITHDRAW_TIMESTAMP);
+
+    // user unbonded we automatically distribute rewards
+    assert_eq!(
+        staking.query_undistributed_rewards(&lp_token.address),
+        100_000
+    );
+}

From c61d5c4400cf8d1cf7827c5718ee49dc22ccb76f Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Fri, 8 Mar 2024 10:50:43 +0200
Subject: [PATCH 04/10] stake: removes get_withdrawable_rewards and uses the
 old code instead

---
 contracts/stake/src/contract.rs     | 45 +++++++++++++++++++++++++----
 contracts/stake/src/distribution.rs | 27 ++---------------
 2 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index b79123aaa..b75e82b47 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -3,12 +3,12 @@ use soroban_sdk::{contract, contractimpl, contractmeta, log, vec, Address, Env,
 use crate::{
     distribution::{
         calculate_annualized_payout, get_distribution, get_reward_curve, get_withdraw_adjustment,
-        get_withdrawable_rewards, save_distribution, save_reward_curve, save_withdraw_adjustment,
-        update_rewards, withdraw_rewards, withdrawable_rewards, Distribution, SHARES_SHIFT,
+        save_distribution, save_reward_curve, save_withdraw_adjustment, update_rewards,
+        withdraw_rewards, withdrawable_rewards, Distribution, SHARES_SHIFT,
     },
     msg::{
         AnnualizedReward, AnnualizedRewardsResponse, ConfigResponse, StakedResponse,
-        WithdrawableRewardsResponse,
+        WithdrawableReward, WithdrawableRewardsResponse,
     },
     storage::{
         get_config, get_stakes, save_config, save_stakes,
@@ -181,8 +181,24 @@ impl StakingTrait for Staking {
 
         let config = get_config(&env);
 
-        let withdrawable_rewards = get_withdrawable_rewards(&env, &sender);
-        if !withdrawable_rewards.rewards.is_empty() {
+        // check for rewards and withdraw them
+        let mut rewards = vec![&env];
+        for distribution_address in get_distributions(&env) {
+            // get distribution data for the given reward
+            let distribution = get_distribution(&env, &distribution_address);
+            // get withdraw adjustment for the given distribution
+            let withdraw_adjustment = get_withdraw_adjustment(&env, &sender, &distribution_address);
+            // calculate current reward amount given the distribution and subtracting withdraw
+            // adjustments
+            let reward_amount =
+                withdrawable_rewards(&env, &sender, &distribution, &withdraw_adjustment);
+            rewards.push_back(WithdrawableReward {
+                reward_address: distribution_address,
+                reward_amount,
+            });
+        }
+
+        if !rewards.is_empty() {
             withdraw_rewards(&env, &sender);
         }
 
@@ -447,7 +463,24 @@ impl StakingTrait for Staking {
     }
 
     fn query_withdrawable_rewards(env: Env, user: Address) -> WithdrawableRewardsResponse {
-        get_withdrawable_rewards(&env, &user)
+        // iterate over all distributions and calculate withdrawable rewards
+        let mut rewards = vec![&env];
+        for distribution_address in get_distributions(&env) {
+            // get distribution data for the given reward
+            let distribution = get_distribution(&env, &distribution_address);
+            // get withdraw adjustment for the given distribution
+            let withdraw_adjustment = get_withdraw_adjustment(&env, &user, &distribution_address);
+            // calculate current reward amount given the distribution and subtracting withdraw
+            // adjustments
+            let reward_amount =
+                withdrawable_rewards(&env, &user, &distribution, &withdraw_adjustment);
+            rewards.push_back(WithdrawableReward {
+                reward_address: distribution_address,
+                reward_amount,
+            });
+        }
+
+        WithdrawableRewardsResponse { rewards }
     }
 
     fn query_distributed_rewards(env: Env, asset: Address) -> u128 {
diff --git a/contracts/stake/src/distribution.rs b/contracts/stake/src/distribution.rs
index ba0c54935..943412d7c 100644
--- a/contracts/stake/src/distribution.rs
+++ b/contracts/stake/src/distribution.rs
@@ -1,12 +1,9 @@
-use soroban_sdk::{contracttype, vec, Address, Env};
+use soroban_sdk::{contracttype, Address, Env};
 
 use curve::Curve;
 use decimal::Decimal;
 
-use crate::{
-    msg::{WithdrawableReward, WithdrawableRewardsResponse},
-    storage::{get_stakes, utils::get_distributions},
-};
+use crate::storage::{get_stakes, utils::get_distributions};
 
 /// How much points is the worth of single token in rewards distribution.
 /// The scaling is performed to have better precision of fixed point division.
@@ -220,26 +217,6 @@ pub fn calculate_annualized_payout(reward_curve: Option<Curve>, now: u64) -> Dec
     }
 }
 
-pub fn get_withdrawable_rewards(env: &Env, user: &Address) -> WithdrawableRewardsResponse {
-    // iterate over all distributions and calculate withdrawable rewards
-    let mut rewards = vec![&env];
-    for distribution_address in get_distributions(env) {
-        // get distribution data for the given reward
-        let distribution = get_distribution(env, &distribution_address);
-        // get withdraw adjustment for the given distribution
-        let withdraw_adjustment = get_withdraw_adjustment(env, user, &distribution_address);
-        // calculate current reward amount given the distribution and subtracting withdraw
-        // adjustments
-        let reward_amount = withdrawable_rewards(env, user, &distribution, &withdraw_adjustment);
-        rewards.push_back(WithdrawableReward {
-            reward_address: distribution_address,
-            reward_amount,
-        });
-    }
-
-    WithdrawableRewardsResponse { rewards }
-}
-
 pub fn withdraw_rewards(env: &Env, sender: &Address) {
     env.events().publish(("withdraw_rewards", "user"), sender);
 

From 5723dc32706d82461e4825c73b1c5c9d24346d69 Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Fri, 8 Mar 2024 12:26:15 +0200
Subject: [PATCH 05/10] stake: removes unnecessary code duplication

---
 contracts/stake/src/contract.rs | 19 +++----------------
 1 file changed, 3 insertions(+), 16 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index b75e82b47..f4a7927c8 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -182,23 +182,10 @@ impl StakingTrait for Staking {
         let config = get_config(&env);
 
         // check for rewards and withdraw them
-        let mut rewards = vec![&env];
-        for distribution_address in get_distributions(&env) {
-            // get distribution data for the given reward
-            let distribution = get_distribution(&env, &distribution_address);
-            // get withdraw adjustment for the given distribution
-            let withdraw_adjustment = get_withdraw_adjustment(&env, &sender, &distribution_address);
-            // calculate current reward amount given the distribution and subtracting withdraw
-            // adjustments
-            let reward_amount =
-                withdrawable_rewards(&env, &sender, &distribution, &withdraw_adjustment);
-            rewards.push_back(WithdrawableReward {
-                reward_address: distribution_address,
-                reward_amount,
-            });
-        }
+        let found_rewards: WithdrawableRewardsResponse =
+            Self::query_withdrawable_rewards(env.clone(), sender.clone());
 
-        if !rewards.is_empty() {
+        if !found_rewards.rewards.is_empty() {
             withdraw_rewards(&env, &sender);
         }
 

From 9a5de08b13c37da02c235c35db74dbe7eb1e8b4e Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Fri, 8 Mar 2024 15:56:05 +0200
Subject: [PATCH 06/10] stake: additional refactoring

---
 contracts/stake/src/contract.rs     |  4 +--
 contracts/stake/src/distribution.rs | 40 +----------------------------
 2 files changed, 3 insertions(+), 41 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index f4a7927c8..1af9401eb 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -4,7 +4,7 @@ use crate::{
     distribution::{
         calculate_annualized_payout, get_distribution, get_reward_curve, get_withdraw_adjustment,
         save_distribution, save_reward_curve, save_withdraw_adjustment, update_rewards,
-        withdraw_rewards, withdrawable_rewards, Distribution, SHARES_SHIFT,
+        withdrawable_rewards, Distribution, SHARES_SHIFT,
     },
     msg::{
         AnnualizedReward, AnnualizedRewardsResponse, ConfigResponse, StakedResponse,
@@ -186,7 +186,7 @@ impl StakingTrait for Staking {
             Self::query_withdrawable_rewards(env.clone(), sender.clone());
 
         if !found_rewards.rewards.is_empty() {
-            withdraw_rewards(&env, &sender);
+            Self::withdraw_rewards(env.clone(), sender.clone());
         }
 
         let mut stakes = get_stakes(&env, &sender);
diff --git a/contracts/stake/src/distribution.rs b/contracts/stake/src/distribution.rs
index 943412d7c..9cf87b047 100644
--- a/contracts/stake/src/distribution.rs
+++ b/contracts/stake/src/distribution.rs
@@ -3,7 +3,7 @@ use soroban_sdk::{contracttype, Address, Env};
 use curve::Curve;
 use decimal::Decimal;
 
-use crate::storage::{get_stakes, utils::get_distributions};
+use crate::storage::get_stakes;
 
 /// How much points is the worth of single token in rewards distribution.
 /// The scaling is performed to have better precision of fixed point division.
@@ -216,41 +216,3 @@ pub fn calculate_annualized_payout(reward_curve: Option<Curve>, now: u64) -> Dec
         None => Decimal::zero(),
     }
 }
-
-pub fn withdraw_rewards(env: &Env, sender: &Address) {
-    env.events().publish(("withdraw_rewards", "user"), sender);
-
-    for distribution_address in get_distributions(env) {
-        // get distribution data for the given reward
-        let mut distribution = get_distribution(env, &distribution_address);
-        // get withdraw adjustment for the given distribution
-        let mut withdraw_adjustment = get_withdraw_adjustment(env, sender, &distribution_address);
-        // calculate current reward amount given the distribution and subtracting withdraw
-        // adjustments
-        let reward_amount = withdrawable_rewards(env, sender, &distribution, &withdraw_adjustment);
-
-        if reward_amount == 0 {
-            continue;
-        }
-
-        withdraw_adjustment.withdrawn_rewards += reward_amount;
-        distribution.withdrawable_total -= reward_amount;
-
-        save_distribution(env, &distribution_address, &distribution);
-        save_withdraw_adjustment(env, sender, &distribution_address, &withdraw_adjustment);
-
-        let reward_token_client = crate::token_contract::Client::new(env, &distribution_address);
-        reward_token_client.transfer(
-            &env.current_contract_address(),
-            sender,
-            &(reward_amount as i128),
-        );
-
-        env.events().publish(
-            ("withdraw_rewards", "reward_token"),
-            &reward_token_client.address,
-        );
-        env.events()
-            .publish(("withdraw_rewards", "reward_amount"), reward_amount as i128);
-    }
-}

From 8b7760d254c5ee5e45915b0b6f425440a0d7e6ac Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Fri, 8 Mar 2024 22:01:32 +0200
Subject: [PATCH 07/10] stake: debugging why unbond does not pay rewards

---
 contracts/stake/src/contract.rs   |  2 ++
 contracts/stake/src/tests/bond.rs | 33 +++++++++++++++----------------
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index 1af9401eb..eb972d607 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -1,3 +1,4 @@
+use soroban_sdk::testutils::arbitrary::std::dbg;
 use soroban_sdk::{contract, contractimpl, contractmeta, log, vec, Address, Env, String, Vec};
 
 use crate::{
@@ -185,6 +186,7 @@ impl StakingTrait for Staking {
         let found_rewards: WithdrawableRewardsResponse =
             Self::query_withdrawable_rewards(env.clone(), sender.clone());
 
+        dbg!(found_rewards.clone());
         if !found_rewards.rewards.is_empty() {
             Self::withdraw_rewards(env.clone(), sender.clone());
         }
diff --git a/contracts/stake/src/tests/bond.rs b/contracts/stake/src/tests/bond.rs
index 2ac6440cc..168140c0d 100644
--- a/contracts/stake/src/tests/bond.rs
+++ b/contracts/stake/src/tests/bond.rs
@@ -1,4 +1,5 @@
 use pretty_assertions::assert_eq;
+use soroban_sdk::testutils::arbitrary::std::dbg;
 use soroban_sdk::{
     testutils::{Address as _, Ledger},
     vec, Address, Env,
@@ -229,9 +230,7 @@ fn unbond_wrong_user_stake_not_found() {
 
 #[test]
 fn pay_rewards_during_unbond() {
-    const AMOUNT: i128 = 100_000;
-    const WITHDRAW_TIMESTAMP: u64 = 2_000;
-
+    const STAKED_AMOUNT: i128 = 1_000;
     let env = Env::default();
     env.mock_all_auths();
 
@@ -241,21 +240,19 @@ fn pay_rewards_during_unbond() {
     let owner = Address::generate(&env);
 
     let lp_token = deploy_token_contract(&env, &admin);
+    let reward_token = deploy_token_contract(&env, &admin);
     let staking = deploy_staking_contract(&env, admin.clone(), &lp_token.address, &manager, &owner);
 
-    lp_token.mint(&user, &AMOUNT);
-    lp_token.mint(&owner, &AMOUNT);
+    lp_token.mint(&user, &10_000);
+    reward_token.mint(&admin, &10_000);
 
-    staking.create_distribution_flow(&owner, &lp_token.address);
-    staking.fund_distribution(&owner, &0u64, &10_000u64, &lp_token.address, &AMOUNT);
+    staking.create_distribution_flow(&manager, &reward_token.address);
+    staking.fund_distribution(&admin, &0u64, &10_000u64, &reward_token.address, &10_000);
 
-    env.ledger().with_mut(|li| {
-        li.timestamp = WITHDRAW_TIMESTAMP;
-    });
-    staking.bond(&user, &10_000);
+    staking.bond(&user, &STAKED_AMOUNT);
 
     env.ledger().with_mut(|li| {
-        li.timestamp = 4_000;
+        li.timestamp = 10_000;
     });
 
     // user hasn't unbonded yet, no rewards to withdraw
@@ -269,14 +266,16 @@ fn pay_rewards_during_unbond() {
         0
     );
     assert_eq!(
-        staking.query_undistributed_rewards(&lp_token.address),
-        110_000
+        staking.query_undistributed_rewards(&reward_token.address),
+        10_000
     );
-    staking.unbond(&user, &10_000, &WITHDRAW_TIMESTAMP);
+    dbg!(reward_token.balance(&user));
+    staking.unbond(&user, &STAKED_AMOUNT, &0);
+    dbg!(reward_token.balance(&user));
 
     // user unbonded we automatically distribute rewards
     assert_eq!(
-        staking.query_undistributed_rewards(&lp_token.address),
-        100_000
+        staking.query_undistributed_rewards(&reward_token.address),
+        10_000
     );
 }

From 6d276495d66d2c159b6052afa72f32f44d644ef4 Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Mon, 11 Mar 2024 10:09:36 +0200
Subject: [PATCH 08/10] stake: removes unnecessary code

---
 contracts/stake/src/tests/bond.rs | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/contracts/stake/src/tests/bond.rs b/contracts/stake/src/tests/bond.rs
index 168140c0d..e3e574509 100644
--- a/contracts/stake/src/tests/bond.rs
+++ b/contracts/stake/src/tests/bond.rs
@@ -255,16 +255,6 @@ fn pay_rewards_during_unbond() {
         li.timestamp = 10_000;
     });
 
-    // user hasn't unbonded yet, no rewards to withdraw
-    assert_eq!(
-        staking
-            .query_withdrawable_rewards(&user)
-            .rewards
-            .iter()
-            .map(|reward| reward.reward_amount)
-            .sum::<u128>(),
-        0
-    );
     assert_eq!(
         staking.query_undistributed_rewards(&reward_token.address),
         10_000

From 480e4901f467bf7fa93e1617d4b948879d1afac9 Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Mon, 11 Mar 2024 10:37:10 +0200
Subject: [PATCH 09/10] stake: fixes the test

---
 contracts/stake/src/contract.rs   |  2 --
 contracts/stake/src/tests/bond.rs | 27 +++++++++++++++++----------
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index eb972d607..1af9401eb 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -1,4 +1,3 @@
-use soroban_sdk::testutils::arbitrary::std::dbg;
 use soroban_sdk::{contract, contractimpl, contractmeta, log, vec, Address, Env, String, Vec};
 
 use crate::{
@@ -186,7 +185,6 @@ impl StakingTrait for Staking {
         let found_rewards: WithdrawableRewardsResponse =
             Self::query_withdrawable_rewards(env.clone(), sender.clone());
 
-        dbg!(found_rewards.clone());
         if !found_rewards.rewards.is_empty() {
             Self::withdraw_rewards(env.clone(), sender.clone());
         }
diff --git a/contracts/stake/src/tests/bond.rs b/contracts/stake/src/tests/bond.rs
index e3e574509..052b56e27 100644
--- a/contracts/stake/src/tests/bond.rs
+++ b/contracts/stake/src/tests/bond.rs
@@ -1,5 +1,4 @@
 use pretty_assertions::assert_eq;
-use soroban_sdk::testutils::arbitrary::std::dbg;
 use soroban_sdk::{
     testutils::{Address as _, Ledger},
     vec, Address, Env,
@@ -252,20 +251,28 @@ fn pay_rewards_during_unbond() {
     staking.bond(&user, &STAKED_AMOUNT);
 
     env.ledger().with_mut(|li| {
-        li.timestamp = 10_000;
+        li.timestamp = 5_000;
     });
+    staking.distribute_rewards();
 
+    // user has bonded for 5_000 time, initial rewards are 10_000
+    // so user should have 5_000 rewards
+    // 5_000 rewards are still undistributed
     assert_eq!(
         staking.query_undistributed_rewards(&reward_token.address),
-        10_000
+        5_000
     );
-    dbg!(reward_token.balance(&user));
-    staking.unbond(&user, &STAKED_AMOUNT, &0);
-    dbg!(reward_token.balance(&user));
-
-    // user unbonded we automatically distribute rewards
     assert_eq!(
-        staking.query_undistributed_rewards(&reward_token.address),
-        10_000
+        staking
+            .query_withdrawable_rewards(&user)
+            .rewards
+            .iter()
+            .map(|reward| reward.reward_amount)
+            .sum::<u128>(),
+        5_000
     );
+
+    assert_eq!(reward_token.balance(&user), 0);
+    staking.unbond(&user, &STAKED_AMOUNT, &0);
+    assert_eq!(reward_token.balance(&user), 5_000);
 }

From 3d29152bffdc0cada4fac4f5de78e6777b64b964 Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <gangov1@gmail.com>
Date: Mon, 11 Mar 2024 13:27:58 +0200
Subject: [PATCH 10/10] CHANGELOG.md: Add entry for #80 (stake pay rewards
 during unbond)

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index aae3bafe2..53d3c743c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ and this project adheres to
 
 ## Changed
 
+- Stake: unbonding now automatically pays the rewards ([#170])
 - Update soroban-sdk version from v20.0.3 to v20.1.0 ([#193])
 - Fixes documentation and naming ([#200])
 - Multihop: adds a new field in the Swap struct, that hold max_belief_price ([234])
@@ -19,6 +20,7 @@ and this project adheres to
 - Factory, Multihop, Pool, Pool_stable, Phoenix: adds lp_token decimal's as a const instead a user input ([#241])
 - Factory, Multihop, Pool, Pool_stable: adds a new parameter for creating liquidity pool ([#243])
 
+[#170]: https://github.com/Phoenix-Protocol-Group/phoenix-contracts/pull/170
 [#200]: https://github.com/Phoenix-Protocol-Group/phoenix-contracts/pull/200
 [#234]: https://github.com/Phoenix-Protocol-Group/phoenix-contracts/pull/234
 [#235]: https://github.com/Phoenix-Protocol-Group/phoenix-contracts/pull/235