[MEX-569] Rc/redist fees collector

common/modules/sc_whitelist_module/src/sc_whitelist_module.rs

@@ -40,6 +40,7 @@ pub trait SCWhitelistModule {
         match opt_orig_caller {
             OptionalValue::Some(opt_caller) => {
                 self.require_sc_address_whitelisted(caller);
+
                 opt_caller
             }
             OptionalValue::None => caller.clone(),

dex/pair/tests/pair_rs_test.rs

@@ -1600,6 +1600,7 @@ fn fees_collector_pair_test() {
                 sc.init(
                     managed_token_id!(LOCKED_TOKEN_ID),
                     managed_address!(&energy_factory_mock_addr),
+                    MultiValueEncoded::new(),
                 );
                 let _ = sc.known_contracts().insert(managed_address!(&pair_addr));
 

energy-integration/common-modules/weekly-rewards-splitting/src/base_impl.rs

@@ -64,13 +64,15 @@ pub trait WeeklyRewardsSplittingTraitsModule {
         let total_rewards = self.collect_and_get_rewards_for_week(sc, week);
         for weekly_reward in &total_rewards {
             let reward_amount = weekly_reward.amount * energy_amount / total_energy;
-            if reward_amount > 0 {
-                user_rewards.push(EsdtTokenPayment::new(
-                    weekly_reward.token_identifier,
-                    0,
-                    reward_amount,
-                ));
+            if reward_amount == 0 {
+                continue;
             }
+
+            user_rewards.push(EsdtTokenPayment::new(
+                weekly_reward.token_identifier,
+                0,
+                reward_amount,
+            ));
         }
 
         user_rewards

energy-integration/common-modules/weekly-rewards-splitting/src/global_info.rs

@@ -1,6 +1,6 @@
 multiversx_sc::imports!();
 
-use common_types::Week;
+use common_types::{PaymentsVec, Week};
 use energy_query::Energy;
 use week_timekeeping::EPOCHS_IN_WEEK;
 
@@ -147,10 +147,7 @@ pub trait WeeklyRewardsGlobalInfo:
 
     #[view(getTotalRewardsForWeek)]
     #[storage_mapper("totalRewardsForWeek")]
-    fn total_rewards_for_week(
-        &self,
-        week: Week,
-    ) -> SingleValueMapper<ManagedVec<EsdtTokenPayment<Self::Api>>>;
+    fn total_rewards_for_week(&self, week: Week) -> SingleValueMapper<PaymentsVec<Self::Api>>;
 
     #[view(getTotalEnergyForWeek)]
     #[storage_mapper("totalEnergyForWeek")]

energy-integration/fees-collector/src/additional_locked_tokens.rs

@@ -20,19 +20,18 @@ pub trait AdditionalLockedTokensModule:
 
     fn accumulate_additional_locked_tokens(&self) {
         let last_update_week_mapper = self.last_locked_token_add_week();
-        let mut last_update_week = last_update_week_mapper.get();
+        let last_update_week = last_update_week_mapper.get();
         let current_week = self.get_current_week();
         if last_update_week == current_week {
             return;
         }
 
-        last_update_week = current_week - 1;
         let blocks_in_week = BLOCKS_IN_WEEK;
         let amount_per_block = self.locked_tokens_per_block().get();
         let new_tokens_amount = amount_per_block * blocks_in_week;
 
         let locked_token_id = self.locked_token_id().get();
-        self.accumulated_fees(last_update_week, &locked_token_id)
+        self.accumulated_fees(current_week - 1, &locked_token_id)
             .update(|fees| *fees += new_tokens_amount);
 
         last_update_week_mapper.set(current_week);

energy-integration/fees-collector/src/claim.rs

@@ -0,0 +1,204 @@
+use core::marker::PhantomData;
+
+use common_types::{PaymentsVec, Week};
+use weekly_rewards_splitting::base_impl::WeeklyRewardsSplittingTraitsModule;
+
+multiversx_sc::imports!();
+
+#[multiversx_sc::module]
+pub trait ClaimModule:
+    crate::config::ConfigModule
+    + crate::events::FeesCollectorEventsModule
+    + weekly_rewards_splitting::WeeklyRewardsSplittingModule
+    + weekly_rewards_splitting::events::WeeklyRewardsSplittingEventsModule
+    + weekly_rewards_splitting::global_info::WeeklyRewardsGlobalInfo
+    + weekly_rewards_splitting::locked_token_buckets::WeeklyRewardsLockedTokenBucketsModule
+    + weekly_rewards_splitting::update_claim_progress_energy::UpdateClaimProgressEnergyModule
+    + crate::fees_accumulation::FeesAccumulationModule
+    + crate::additional_locked_tokens::AdditionalLockedTokensModule
+    + locking_module::lock_with_energy_module::LockWithEnergyModule
+    + energy_query::EnergyQueryModule
+    + week_timekeeping::WeekTimekeepingModule
+    + multiversx_sc_modules::pause::PauseModule
+    + utils::UtilsModule
+    + sc_whitelist_module::SCWhitelistModule
+    + multiversx_sc_modules::only_admin::OnlyAdminModule
+    + crate::redistribute_rewards::RedistributeRewardsModule
+{
+    #[endpoint(claimRewards)]
+    fn claim_rewards_endpoint(
+        &self,
+        opt_original_caller: OptionalValue<ManagedAddress>,
+    ) -> PaymentsVec<Self::Api> {
+        self.require_not_paused();
+
+        let caller = self.blockchain().get_caller();
+        let original_caller = self.get_orig_caller_from_opt(&caller, opt_original_caller);
+
+        self.claim_rewards(caller, original_caller)
+    }
+
+    #[endpoint(claimBoostedRewards)]
+    fn claim_boosted_rewards(
+        &self,
+        opt_original_caller: OptionalValue<ManagedAddress>,
+    ) -> PaymentsVec<Self::Api> {
+        self.require_not_paused();
+
+        let original_caller = match opt_original_caller {
+            OptionalValue::Some(user) => {
+                require!(
+                    self.allow_external_claim_rewards(&user).get(),
+                    "Cannot claim rewards for this address"
+                );
+
+                user
+            }
+            OptionalValue::None => self.blockchain().get_caller(),
+        };
+
+        self.claim_rewards(original_caller.clone(), original_caller)
+    }
+
+    fn claim_rewards(
+        &self,
+        caller: ManagedAddress,
+        original_caller: ManagedAddress,
+    ) -> PaymentsVec<Self::Api> {
+        self.accumulate_additional_locked_tokens();
+
+        let wrapper = FeesCollectorWrapper::new();
+        let mut rewards = self.claim_multi(&wrapper, &original_caller);
+        if rewards.is_empty() {
+            return rewards;
+        }
+
+        let locked_token_id = self.get_locked_token_id();
+        let mut i = 0;
+        let mut len = rewards.len();
+        let mut total_locked_token_rewards_amount = BigUint::zero();
+        while i < len {
+            let rew = rewards.get(i);
+            if rew.token_identifier != locked_token_id {
+                i += 1;
+                continue;
+            }
+
+            total_locked_token_rewards_amount += rew.amount;
+            len -= 1;
+            rewards.remove(i);
+        }
+
+        if !rewards.is_empty() {
+            self.send().direct_multi(&caller, &rewards);
+        }
+
+        if total_locked_token_rewards_amount > 0 {
+            let locked_rewards = self.lock_virtual(
+                self.get_base_token_id(),
+                total_locked_token_rewards_amount,
+                caller,
+                original_caller,
+            );
+
+            rewards.push(locked_rewards);
+        }
+
+        rewards
+    }
+}
+
+pub struct FeesCollectorWrapper<T: ClaimModule> {
+    phantom: PhantomData<T>,
+}
+
+impl<T: ClaimModule> Default for FeesCollectorWrapper<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T: ClaimModule> FeesCollectorWrapper<T> {
+    pub fn new() -> FeesCollectorWrapper<T> {
+        FeesCollectorWrapper {
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<T> WeeklyRewardsSplittingTraitsModule for FeesCollectorWrapper<T>
+where
+    T: ClaimModule,
+{
+    type WeeklyRewardsSplittingMod = T;
+
+    fn get_user_rewards_for_week(
+        &self,
+        sc: &Self::WeeklyRewardsSplittingMod,
+        week: Week,
+        energy_amount: &BigUint<<Self::WeeklyRewardsSplittingMod as ContractBase>::Api>,
+        total_energy: &BigUint<<Self::WeeklyRewardsSplittingMod as ContractBase>::Api>,
+    ) -> PaymentsVec<<Self::WeeklyRewardsSplittingMod as ContractBase>::Api> {
+        let mut user_rewards = ManagedVec::new();
+        if energy_amount == &0 || total_energy == &0 {
+            return user_rewards;
+        }
+
+        let total_rewards = self.collect_and_get_rewards_for_week(sc, week);
+        let remaining_rewards_mapper = sc.remaining_rewards(week);
+        let mut remaining_rewards = remaining_rewards_mapper.get();
+        for (i, weekly_reward) in total_rewards.iter().enumerate() {
+            let reward_amount = weekly_reward.amount * energy_amount / total_energy;
+            if reward_amount == 0 {
+                continue;
+            }
+
+            let mut rem_rew_entry = remaining_rewards.get_mut(i);
+            rem_rew_entry.amount -= &reward_amount;
+
+            user_rewards.push(EsdtTokenPayment::new(
+                weekly_reward.token_identifier,
+                0,
+                reward_amount,
+            ));
+        }
+
+        remaining_rewards_mapper.set(remaining_rewards);
+
+        user_rewards
+    }
+
+    fn collect_and_get_rewards_for_week(
+        &self,
+        sc: &Self::WeeklyRewardsSplittingMod,
+        week: Week,
+    ) -> PaymentsVec<<Self::WeeklyRewardsSplittingMod as ContractBase>::Api> {
+        let total_rewards_mapper = sc.total_rewards_for_week(week);
+        if total_rewards_mapper.is_empty() {
+            let total_rewards = self.collect_rewards_for_week(sc, week);
+            total_rewards_mapper.set(&total_rewards);
+            sc.remaining_rewards(week).set(&total_rewards);
+
+            total_rewards
+        } else {
+            total_rewards_mapper.get()
+        }
+    }
+
+    fn collect_rewards_for_week(
+        &self,
+        sc: &Self::WeeklyRewardsSplittingMod,
+        week: Week,
+    ) -> PaymentsVec<<Self::WeeklyRewardsSplittingMod as ContractBase>::Api> {
+        let mut results = ManagedVec::new();
+        let all_tokens = sc.all_tokens().get();
+        for token in &all_tokens {
+            let opt_accumulated_fees = sc.get_and_clear_accumulated_fees(week, &token);
+            if let Some(accumulated_fees) = opt_accumulated_fees {
+                results.push(EsdtTokenPayment::new(token, 0, accumulated_fees));
+            }
+        }
+
+        results
+    }
+}

energy-integration/fees-collector/src/config.rs

@@ -11,6 +11,7 @@ pub trait ConfigModule {
                 self.blockchain().is_smart_contract(&sc),
                 "Invalid SC address"
             );
+
             let _ = mapper.insert(sc);
         }
     }
@@ -32,10 +33,12 @@ pub trait ConfigModule {
         for token in tokens {
             require!(token.is_valid_esdt_identifier(), "Invalid token ID");
 
-            if !known_tokens_mapper.contains(&token) {
-                known_tokens_mapper.add(&token);
-                all_tokens_vec.push(token);
+            if known_tokens_mapper.contains(&token) {
+                continue;
             }
+
+            known_tokens_mapper.add(&token);
+            all_tokens_vec.push(token);
         }
 
         self.all_tokens().set(&all_tokens_vec);
@@ -47,28 +50,29 @@ pub trait ConfigModule {
         let mut all_tokens_vec = self.all_tokens().get();
         let known_tokens_mapper = self.known_tokens();
         for token in tokens {
-            if known_tokens_mapper.contains(&token) {
-                known_tokens_mapper.remove(&token);
+            if !known_tokens_mapper.contains(&token) {
+                continue;
+            }
 
-                unsafe {
-                    let index = all_tokens_vec.find(&token).unwrap_unchecked();
-                    all_tokens_vec.remove(index);
-                }
+            known_tokens_mapper.remove(&token);
+            unsafe {
+                let index = all_tokens_vec.find(&token).unwrap_unchecked();
+                all_tokens_vec.remove(index);
             }
         }
 
         self.all_tokens().set(&all_tokens_vec);
     }
 
-    #[view(getLockedTokenId)]
-    #[storage_mapper("lockedTokenId")]
-    fn locked_token_id(&self) -> SingleValueMapper<TokenIdentifier>;
-
     #[view(getAllTokens)]
     fn get_all_tokens(&self) -> MultiValueEncoded<TokenIdentifier> {
         self.all_tokens().get().into()
     }
 
+    #[view(getLockedTokenId)]
+    #[storage_mapper("lockedTokenId")]
+    fn locked_token_id(&self) -> SingleValueMapper<TokenIdentifier>;
+
     #[view(getAllKnownContracts)]
     #[storage_mapper("knownContracts")]
     fn known_contracts(&self) -> UnorderedSetMapper<ManagedAddress>;

energy-integration/fees-collector/src/events.rs

@@ -6,18 +6,18 @@ use common_types::Week;
 pub trait FeesCollectorEventsModule {
     fn emit_deposit_swap_fees_event(
         self,
-        caller: ManagedAddress,
+        caller: &ManagedAddress,
         current_week: Week,
-        payment: EsdtTokenPayment<Self::Api>,
+        payment: &EsdtTokenPayment,
     ) {
         self.deposit_swap_fees_event(caller, current_week, payment);
     }
 
     #[event("deposit_swap_fees_event")]
     fn deposit_swap_fees_event(
         &self,
-        #[indexed] caller: ManagedAddress,
+        #[indexed] caller: &ManagedAddress,
         #[indexed] current_week: Week,
-        #[indexed] payment: EsdtTokenPayment<Self::Api>,
+        #[indexed] payment: &EsdtTokenPayment,
     );
 }

energy-integration/fees-collector/src/fees_accumulation.rs

@@ -25,23 +25,25 @@ pub trait FeesAccumulationModule:
             self.known_tokens().contains(&payment.token_identifier),
             "Invalid payment token"
         );
-        let current_week = self.get_current_week();
 
         if payment.token_nonce > 0 {
             require!(
                 payment.token_identifier == self.locked_token_id().get(),
                 "Invalid locked token"
             );
+
             self.send().esdt_local_burn(
                 &payment.token_identifier,
                 payment.token_nonce,
                 &payment.amount,
             );
         }
+
+        let current_week = self.get_current_week();
         self.accumulated_fees(current_week, &payment.token_identifier)
             .update(|amt| *amt += &payment.amount);
 
-        self.emit_deposit_swap_fees_event(caller, current_week, payment);
+        self.emit_deposit_swap_fees_event(&caller, current_week, &payment);
     }
 
     fn get_and_clear_accumulated_fees(