just some cleanup

locked-asset/token-unstake/src/events.rs

@@ -1,3 +1,5 @@
+use common_structs::Epoch;
+
 use crate::tokens_per_user::UnstakePair;
 
 multiversx_sc::imports!();
@@ -24,7 +26,7 @@ pub trait EventsModule {
         &self,
         #[indexed] caller: &ManagedAddress,
         #[indexed] block: u64,
-        #[indexed] epoch: u64,
+        #[indexed] epoch: Epoch,
         #[indexed] timestamp: u64,
         data: ManagedVec<UnstakePair<Self::Api>>,
     );

locked-asset/token-unstake/src/fees_handler.rs

@@ -1,6 +1,8 @@
 multiversx_sc::imports!();
 
-pub const MAX_PENALTY_PERCENTAGE: u64 = 10_000;
+pub const MAX_PENALTY_PERCENTAGE: Percent = 10_000;
+
+use common_structs::Percent;
 
 use crate::{events, tokens_per_user::UnstakePair};
 
@@ -103,7 +105,7 @@ pub trait FeesHandlerModule:
 
     #[view(getFeesBurnPercentage)]
     #[storage_mapper("feesBurnPercentage")]
-    fn fees_burn_percentage(&self) -> SingleValueMapper<u64>;
+    fn fees_burn_percentage(&self) -> SingleValueMapper<Percent>;
 
     #[view(getFeesCollectorAddress)]
     #[storage_mapper("feesCollectorAddress")]

locked-asset/token-unstake/src/lib.rs

@@ -8,6 +8,8 @@ pub mod fees_handler;
 pub mod tokens_per_user;
 pub mod unbond_tokens;
 
+use common_structs::{Epoch, Percent};
+
 use crate::fees_handler::MAX_PENALTY_PERCENTAGE;
 
 #[multiversx_sc::contract]
@@ -24,9 +26,9 @@ pub trait TokenUnstakeModule:
     #[init]
     fn init(
         &self,
-        unbond_epochs: u64,
+        unbond_epochs: Epoch,
         energy_factory_address: ManagedAddress,
-        fees_burn_percentage: u64,
+        fees_burn_percentage: Percent,
         fees_collector_address: ManagedAddress,
     ) {
         self.require_sc_address(&energy_factory_address);
@@ -37,8 +39,8 @@ pub trait TokenUnstakeModule:
         );
 
         self.unbond_epochs().set(unbond_epochs);
-        self.energy_factory_address().set(&energy_factory_address);
-        self.fees_collector_address().set(&fees_collector_address);
+        self.energy_factory_address().set(energy_factory_address);
+        self.fees_collector_address().set(fees_collector_address);
         self.fees_burn_percentage().set(fees_burn_percentage);
     }
 

locked-asset/token-unstake/src/tokens_per_user.rs

@@ -1,3 +1,5 @@
+use common_structs::Epoch;
+
 multiversx_sc::imports!();
 multiversx_sc::derive_imports!();
 
@@ -13,7 +15,7 @@ multiversx_sc::derive_imports!();
     Debug,
 )]
 pub struct UnstakePair<M: ManagedTypeApi> {
-    pub unlock_epoch: u64,
+    pub unlock_epoch: Epoch,
     pub locked_tokens: EsdtTokenPayment<M>,
     pub unlocked_tokens: EsdtTokenPayment<M>,
 }
@@ -22,7 +24,7 @@ pub struct UnstakePair<M: ManagedTypeApi> {
 pub trait TokensPerUserModule {
     #[view(getUnbondEpochs)]
     #[storage_mapper("unbondEpochs")]
-    fn unbond_epochs(&self) -> SingleValueMapper<u64>;
+    fn unbond_epochs(&self) -> SingleValueMapper<Epoch>;
 
     #[view(getUnlockedTokensForUser)]
     #[storage_mapper("unlockedTokensForUser")]

locked-asset/token-unstake/src/unbond_tokens.rs

@@ -1,4 +1,6 @@
-use crate::events;
+use common_structs::PaymentsVec;
+
+use crate::{events, tokens_per_user::UnstakePair};
 
 multiversx_sc::imports!();
 
@@ -16,39 +18,23 @@ pub trait UnbondTokensModule:
         let current_epoch = self.blockchain().get_block_epoch();
         let mut output_payments = ManagedVec::new();
         let mut penalty_tokens = ManagedVec::<Self::Api, _>::new();
-        self.unlocked_tokens_for_user(&caller)
+        let new_unlocked_tokens = self
+            .unlocked_tokens_for_user(&caller)
             .update(|user_entries| {
                 while !user_entries.is_empty() {
                     let entry = user_entries.get(0);
                     if current_epoch < entry.unlock_epoch {
                         break;
                     }
 
-                    let locked_tokens = entry.locked_tokens;
-                    let unlocked_tokens = entry.unlocked_tokens;
-
-                    // we only burn the tokens that are not unlocked
-                    // the rest are sent back as penalty
-                    let locked_tokens_burn_amount = unlocked_tokens.amount.clone();
-                    self.send().esdt_local_burn(
-                        &locked_tokens.token_identifier,
-                        locked_tokens.token_nonce,
-                        &locked_tokens_burn_amount,
-                    );
-
-                    let penalty_amount = &locked_tokens.amount - &unlocked_tokens.amount;
-                    if penalty_amount > 0 {
-                        let penalty = EsdtTokenPayment::new(
-                            locked_tokens.token_identifier,
-                            locked_tokens.token_nonce,
-                            penalty_amount,
-                        );
-                        penalty_tokens.push(penalty);
-                    }
+                    self.handle_single_unbond_entry(&entry, &mut penalty_tokens);
 
+                    let unlocked_tokens = entry.unlocked_tokens;
                     output_payments.push(unlocked_tokens);
                     user_entries.remove(0);
                 }
+
+                (*user_entries).clone()
             });
 
         require!(!output_payments.is_empty(), "Nothing to unbond");
@@ -58,10 +44,38 @@ pub trait UnbondTokensModule:
         }
 
         self.send().direct_multi(&caller, &output_payments);
-
-        let new_unlocked_tokens = self.unlocked_tokens_for_user(&caller).get();
         self.emit_unlocked_tokens_event(&caller, new_unlocked_tokens);
 
         output_payments.into()
     }
+
+    fn handle_single_unbond_entry(
+        &self,
+        entry: &UnstakePair<Self::Api>,
+        penalty_tokens: &mut PaymentsVec<Self::Api>,
+    ) {
+        let locked_tokens = &entry.locked_tokens;
+        let unlocked_tokens = &entry.unlocked_tokens;
+
+        // we only burn the tokens that are not unlocked
+        // the rest are sent back as penalty
+        let locked_tokens_burn_amount = &unlocked_tokens.amount;
+        self.send().esdt_local_burn(
+            &locked_tokens.token_identifier,
+            locked_tokens.token_nonce,
+            locked_tokens_burn_amount,
+        );
+
+        let penalty_amount = &locked_tokens.amount - &unlocked_tokens.amount;
+        if penalty_amount == 0 {
+            return;
+        }
+
+        let penalty = EsdtTokenPayment::new(
+            locked_tokens.token_identifier.clone(),
+            locked_tokens.token_nonce,
+            penalty_amount,
+        );
+        penalty_tokens.push(penalty);
+    }
 }