more refactoring

multiversx · Dec 16, 2024 · a8b4ca7 · a8b4ca7
1 parent 3da60f8
commit a8b4ca7
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 19 deletions.
diff --git a/contracts/examples/lottery-esdt/src/specific/award.rs b/contracts/examples/lottery-esdt/src/specific/award.rs
@@ -163,11 +163,13 @@ pub trait AwardingModule: views::ViewsModule + storage::StorageModule + utils::U
                         .get(total_winning_tickets - *index_last_winner),
                 ),
             );
-            if prize > 0 {
-                self.assign_prize_to_winner(info.token_identifier.clone(), &prize, &winner_address);
-
-                info.unawarded_amount -= prize;
+            if prize == 0 {
+                return;
             }
+
+            self.assign_prize_to_winner(info.token_identifier.clone(), &prize, &winner_address);
+
+            info.unawarded_amount -= prize;
         } else {
             // insert token in accumulated rewards first place
             let first_place_winner = ticket_holders_mapper.get(*index_last_winner);

diff --git a/contracts/examples/lottery-esdt/src/specific/claim.rs b/contracts/examples/lottery-esdt/src/specific/claim.rs
@@ -20,24 +20,11 @@ pub trait ClaimModule: storage::StorageModule {
 
         if tokens.is_empty() {
             // if wanted tokens were not specified claim all, and clear user_accumulated_token_rewards storage mapper
-
-            let mut all_tokens: ManagedVec<Self::Api, EgldOrEsdtTokenIdentifier> =
-                ManagedVec::new();
-
-            for token_id in self.user_accumulated_token_rewards(&caller_id).iter() {
-                require!(
-                    !self.accumulated_rewards(&token_id, &caller_id).is_empty(),
-                    "Token requested not available for claim"
-                );
-                all_tokens.push(token_id);
-            }
-
-            self.claim_rewards_user(
-                all_tokens,
+            self.handle_claim_with_unspecified_tokens(
                 &caller_id,
                 &mut accumulated_egld_rewards,
                 &mut accumulated_esdt_rewards,
-            )
+            );
         } else {
             // otherwise claim just what was requested and remove those tokens from the user_accumulated_token_rewards storage mapper
 
@@ -63,6 +50,30 @@ pub trait ClaimModule: storage::StorageModule {
         }
     }
 
+    fn handle_claim_with_unspecified_tokens(
+        &self,
+        caller_id: &u64,
+        accumulated_egld_rewards: &mut BigUint,
+        accumulated_esdt_rewards: &mut ManagedVec<Self::Api, EsdtTokenPayment>,
+    ) {
+        let mut all_tokens: ManagedVec<Self::Api, EgldOrEsdtTokenIdentifier> = ManagedVec::new();
+
+        for token_id in self.user_accumulated_token_rewards(caller_id).iter() {
+            require!(
+                !self.accumulated_rewards(&token_id, caller_id).is_empty(),
+                "Token requested not available for claim"
+            );
+            all_tokens.push(token_id);
+        }
+
+        self.claim_rewards_user(
+            all_tokens,
+            caller_id,
+            accumulated_egld_rewards,
+            accumulated_esdt_rewards,
+        )
+    }
+
     fn claim_rewards_user(
         &self,
         tokens: ManagedVec<Self::Api, EgldOrEsdtTokenIdentifier>,