fix: add token info to account record

contract/src/burn/api.rs

@@ -20,7 +20,7 @@ impl BurnApi for Contract {
         let mut keys_to_remove = vec![];
         let now = now_seconds();
 
-        for (datetime, (_, total)) in self.accruals.iter() {
+        for (datetime, (_, total)) in self.get_sweat_accruals_unsafe().iter() {
             if now - datetime >= self.burn_period {
                 keys_to_remove.push(*datetime);
                 total_to_burn += total;
@@ -50,8 +50,9 @@ impl Contract {
             return U128(0);
         }
 
+        let sweat_accruals = self.get_sweat_accruals_unsafe_mut();
         for datetime in keys_to_remove {
-            self.accruals.remove(&datetime);
+            sweat_accruals.remove(&datetime);
         }
 
         emit(EventKind::Burn(BurnData {

contract/src/claim/api.rs

@@ -1,7 +1,7 @@
 use model::{
     api::ClaimApi,
     event::{emit, ClaimData, EventKind},
-    ClaimAvailabilityView, ClaimResultView, TokensAmount, UnixTimestamp,
+    AccrualIndex, ClaimAvailabilityView, ClaimResultView, TokensAmount, UnixTimestamp,
 };
 use near_sdk::{env, json_types::U128, near_bindgen, require, store::Vector, AccountId, PromiseOrValue};
 
@@ -10,19 +10,20 @@ use crate::{common::now_seconds, Contract, ContractExt, StorageKey::AccrualsEntr
 #[near_bindgen]
 impl ClaimApi for Contract {
     fn get_claimable_balance_for_account(&self, account_id: AccountId) -> U128 {
-        let Some(account_data) = self.accounts.get(&account_id) else {
+        let Some(account_data) = self.get_sweat_account_data(&account_id) else {
             return U128(0);
         };
 
         let mut total_accrual = 0;
         let now = now_seconds();
+        let sweat_accruals = self.get_sweat_accruals_unsafe();
 
-        for (datetime, index) in &account_data.accruals {
+        for (datetime, index) in account_data {
             if now - datetime > self.burn_period {
                 continue;
             }
 
-            let Some((accruals, _)) = self.accruals.get(datetime) else {
+            let Some((accruals, _)) = sweat_accruals.get(datetime) else {
                 continue;
             };
 
@@ -55,41 +56,51 @@ impl ClaimApi for Contract {
             "Claim is not available at the moment"
         );
 
-        let account_data = self.accounts.get_mut(&account_id).expect("Account data is not found");
-        require!(!account_data.is_locked, "Another operation is running");
+        require!(
+            !self.get_account_data_unsafe(&account_id).is_locked,
+            "Another operation is running"
+        );
 
-        account_data.is_locked = true;
+        self.get_account_data_unsafe_mut(&account_id).is_locked = true;
 
         let now = now_seconds();
         let mut total_accrual = 0;
-        let mut details = vec![];
+        let burn_period = self.burn_period;
 
-        for (datetime, index) in &account_data.accruals {
-            if now - datetime > self.burn_period {
+        let mut details: Vec<(UnixTimestamp, AccrualIndex)> = vec![];
+        for (datetime, index) in self.get_sweat_account_data_unsafe_mut(&account_id).iter() {
+            if now - datetime > burn_period {
                 continue;
             }
 
-            let Some((accruals, total)) = self.accruals.get_mut(datetime) else {
-                continue;
-            };
+            details.push((*datetime, *index));
+        }
 
-            let Some(amount) = accruals.get_mut(*index) else {
+        let mut amount_details: Vec<(UnixTimestamp, TokensAmount)> = vec![];
+        for (datetime, index) in &details {
+            let Some((accruals, _)) = self.get_sweat_accruals_unsafe_mut().get(datetime) else {
                 continue;
             };
 
-            details.push((*datetime, *amount));
+            if let Some(amount) = accruals.get(*index).cloned() {
+                let accrual = self.get_sweat_accruals_unsafe_mut().get_mut(datetime).unwrap();
+                accrual.0.set(*index, 0);
+                accrual.1 -= amount;
+
+                total_accrual += amount;
 
-            total_accrual += *amount;
-            *total -= *amount;
-            *amount = 0;
+                amount_details.push((*datetime, amount));
+            }
         }
 
-        account_data.accruals.clear();
+        self.get_account_data_unsafe_mut(&account_id)
+            .get_sweat_accruals_unsafe_mut()
+            .clear();
 
         if total_accrual > 0 {
-            self.transfer_external(now, account_id, total_accrual, details)
+            self.transfer_external(now, account_id, total_accrual, amount_details)
         } else {
-            account_data.is_locked = false;
+            self.get_account_data_unsafe_mut(&account_id).is_locked = false;
             PromiseOrValue::Value(ClaimResultView::new(0))
         }
     }
@@ -104,11 +115,10 @@ impl Contract {
         details: Vec<(UnixTimestamp, TokensAmount)>,
         is_success: bool,
     ) -> ClaimResultView {
-        let account = self.accounts.get_mut(&account_id).expect("Account not found");
-        account.is_locked = false;
+        self.get_account_data_unsafe_mut(&account_id).is_locked = false;
 
         if is_success {
-            account.claim_period_refreshed_at = now;
+            self.get_account_data_unsafe_mut(&account_id).claim_period_refreshed_at = now;
 
             let event_data = ClaimData {
                 account_id,
@@ -124,15 +134,23 @@ impl Contract {
         }
 
         for (timestamp, amount) in details {
-            let daily_accruals = self
-                .accruals
-                .entry(timestamp)
-                .or_insert_with(|| (Vector::new(AccrualsEntry(timestamp)), 0));
+            if !self.get_sweat_accruals_unsafe().contains_key(&timestamp) {
+                self.get_sweat_accruals_unsafe_mut()
+                    .insert(timestamp, (Vector::new(AccrualsEntry(timestamp)), 0));
+            }
+
+            self.get_sweat_accruals_unsafe_mut()
+                .get_mut(&timestamp)
+                .unwrap()
+                .0
+                .push(amount);
+            self.get_sweat_accruals_unsafe_mut().get_mut(&timestamp).unwrap().1 += amount;
 
-            daily_accruals.0.push(amount);
-            daily_accruals.1 += amount;
+            let accrual_index = self.get_sweat_accruals_unsafe().get(&timestamp).unwrap().0.len() - 1;
 
-            account.accruals.push((timestamp, daily_accruals.0.len() - 1));
+            self.get_account_data_unsafe_mut(&account_id)
+                .get_sweat_accruals_unsafe_mut()
+                .push((timestamp, accrual_index));
         }
 
         ClaimResultView::new(0)

contract/src/common/mod.rs

@@ -1,5 +1,10 @@
-use model::UnixTimestamp;
-use near_sdk::env::{block_timestamp_ms, panic_str};
+use model::{account_record::AccountRecord, AccrualIndex, UnixTimestamp};
+use near_sdk::{
+    env::{block_timestamp_ms, panic_str},
+    AccountId,
+};
+
+use crate::{AccrualsMap, Contract};
 
 mod asserts;
 pub(crate) mod tests;
@@ -13,6 +18,55 @@ pub(crate) fn now_seconds() -> UnixTimestamp {
     ms_timestamp_to_seconds(block_timestamp_ms())
 }
 
+impl Contract {
+    pub(crate) fn get_sweat_accruals(&self) -> Option<&AccrualsMap> {
+        self.accruals.get(&"SWEAT".to_string())
+    }
+    pub(crate) fn get_sweat_accruals_unsafe(&self) -> &AccrualsMap {
+        self.get_sweat_accruals().unwrap()
+    }
+
+    pub(crate) fn get_sweat_accruals_unsafe_mut(&mut self) -> &mut AccrualsMap {
+        self.accruals.get_mut(&"SWEAT".to_string()).unwrap()
+    }
+
+    pub(crate) fn get_account_data(&self, account_id: &AccountId) -> Option<&AccountRecord> {
+        self.accounts.get(account_id)
+    }
+
+    pub(crate) fn get_account_data_unsafe(&self, account_id: &AccountId) -> &AccountRecord {
+        self.get_account_data(account_id).unwrap()
+    }
+
+    pub(crate) fn get_account_data_unsafe_mut(&mut self, account_id: &AccountId) -> &mut AccountRecord {
+        self.get_account_data_mut(account_id).unwrap()
+    }
+
+    pub(crate) fn get_account_data_mut(&mut self, account_id: &AccountId) -> Option<&mut AccountRecord> {
+        self.accounts.get_mut(account_id)
+    }
+
+    pub(crate) fn get_sweat_account_data_unsafe(&self, account_id: &AccountId) -> &Vec<(UnixTimestamp, AccrualIndex)> {
+        self.get_account_data_unsafe(account_id).get_sweat_accruals_unsafe()
+    }
+
+    pub(crate) fn get_sweat_account_data(&self, account_id: &AccountId) -> Option<&Vec<(UnixTimestamp, AccrualIndex)>> {
+        if let Some(record) = self.get_account_data(account_id) {
+            record.get_sweat_accruals()
+        } else {
+            None
+        }
+    }
+
+    pub(crate) fn get_sweat_account_data_unsafe_mut(
+        &mut self,
+        account_id: &AccountId,
+    ) -> &mut Vec<(UnixTimestamp, AccrualIndex)> {
+        self.get_account_data_unsafe_mut(account_id)
+            .get_sweat_accruals_unsafe_mut()
+    }
+}
+
 #[test]
 fn convert_milliseconds_to_unix_timestamp_successfully() {
     let millis: u64 = 1_699_038_575_819;

contract/src/lib.rs

@@ -1,11 +1,13 @@
-use model::{account_record::AccountRecord, api::InitApi, Duration, TokensAmount, UnixTimestamp};
+use model::{account_record::AccountRecord, api::InitApi, Asset, Duration, TokensAmount, UnixTimestamp};
 use near_sdk::{
     borsh::{self, BorshDeserialize, BorshSerialize},
     near_bindgen,
     store::{LookupMap, UnorderedMap, UnorderedSet, Vector},
     AccountId, BorshStorageKey, PanicOnDefault,
 };
 
+use crate::StorageKey::AssetAccruals;
+
 mod auth;
 mod burn;
 mod claim;
@@ -17,6 +19,8 @@ mod record;
 const INITIAL_CLAIM_PERIOD_MS: u32 = 24 * 60 * 60;
 const INITIAL_BURN_PERIOD_MS: u32 = 30 * 24 * 60 * 60;
 
+pub(crate) type AccrualsMap = UnorderedMap<UnixTimestamp, (Vector<TokensAmount>, TokensAmount)>;
+
 /// The main structure representing a smart contract for managing fungible tokens.
 #[near_bindgen]
 #[derive(BorshDeserialize, BorshSerialize, PanicOnDefault)]
@@ -65,7 +69,7 @@ pub struct Contract {
     ///  │     [(1705066501, 2)]        │
     ///  └────────────┘      └──────────┘
     /// ```
-    accruals: UnorderedMap<UnixTimestamp, (Vector<TokensAmount>, TokensAmount)>,
+    accruals: UnorderedMap<Asset, AccrualsMap>,
 
     /// A map containing accrual and service details for each user account.
     ///
@@ -88,6 +92,7 @@ enum StorageKey {
     Accruals,
     AccrualsEntry(u32),
     Oracles,
+    AssetAccruals(Asset),
 }
 
 #[near_bindgen]
@@ -96,11 +101,19 @@ impl InitApi for Contract {
     fn init(token_account_id: AccountId) -> Self {
         Self::assert_private();
 
+        let mut accruals = UnorderedMap::new(StorageKey::Accruals);
+        accruals.insert(
+            "SWEAT".to_string(),
+            UnorderedMap::<UnixTimestamp, (Vector<TokensAmount>, TokensAmount)>::new(AssetAccruals(
+                "SWEAT".to_string(),
+            )),
+        );
+
         Self {
             token_account_id,
+            accruals,
 
             accounts: LookupMap::new(StorageKey::Accounts),
-            accruals: UnorderedMap::new(StorageKey::Accruals),
             oracles: UnorderedSet::new(StorageKey::Oracles),
 
             claim_period: INITIAL_CLAIM_PERIOD_MS,

contract/src/record/api.rs

@@ -3,7 +3,7 @@ use model::{
     api::RecordApi,
     event::{emit, EventKind, RecordData},
 };
-use near_sdk::{json_types::U128, near_bindgen, require, store::Vector, AccountId};
+use near_sdk::{json_types::U128, near_bindgen, store::Vector, AccountId};
 
 use crate::{common::now_seconds, Contract, ContractExt, StorageKey::AccrualsEntry};
 
@@ -31,24 +31,24 @@ impl RecordApi for Contract {
             balances.push(amount);
 
             if let Some(record) = self.accounts.get_mut(&account_id) {
-                record.accruals.push((now_seconds, index));
+                record.get_sweat_accruals_unsafe_mut().push((now_seconds, index));
             } else {
-                let record = AccountRecord {
-                    accruals: vec![(now_seconds, index)],
-                    ..AccountRecord::new(now_seconds)
-                };
+                let mut record = AccountRecord::new(now_seconds, None);
+                record.get_sweat_accruals_unsafe_mut().push((now_seconds, index));
 
                 self.accounts.insert(account_id, record);
             }
         }
 
-        let existing = self.accruals.insert(now_seconds, (balances, total_balance));
+        let sweat_accruals = self.get_sweat_accruals_unsafe_mut();
+        if sweat_accruals.contains_key(&now_seconds) {
+            let current_accruals = sweat_accruals.get_mut(&now_seconds).unwrap();
+            current_accruals.0.extend(balances.iter().cloned());
+            current_accruals.1 += total_balance;
+        } else {
+            sweat_accruals.insert(now_seconds, (balances, total_balance));
+        }
 
         emit(EventKind::Record(event_data));
-
-        require!(
-            existing.is_none(),
-            format!("Record for this timestamp: {now_seconds} already existed. It was overwritten.")
-        );
     }
 }

contract/src/record/tests.rs

@@ -1,7 +1,7 @@
 #![cfg(test)]
 
 use model::api::{ClaimApi, RecordApi};
-use near_sdk::json_types::U128;
+use near_sdk::{json_types::U128, AccountId};
 
 use crate::common::tests::Context;
 
@@ -44,11 +44,20 @@ fn record_by_not_oracle() {
 }
 
 #[test]
-#[should_panic(expected = "Record for this timestamp: 0 already existed. It was overwritten.")]
-fn test_record() {
+fn test_multiple_records_in_the_same_block() {
     let (mut context, mut contract, accounts) = Context::init_with_oracle();
 
+    let target_accruals = [10, 20];
+    let batches: Vec<Vec<(AccountId, U128)>> = target_accruals
+        .iter()
+        .map(|&amount| vec![(accounts.alice.clone(), amount.into())])
+        .collect();
+
     context.switch_account(&accounts.oracle);
-    contract.record_batch_for_hold(vec![(accounts.alice.clone(), 10.into())]);
-    contract.record_batch_for_hold(vec![(accounts.alice, 10.into())]);
+    contract.record_batch_for_hold(batches.get(0).unwrap().clone());
+    contract.record_batch_for_hold(batches.get(1).unwrap().clone());
+
+    let accruals = contract.get_sweat_accruals_unsafe().get(&0).unwrap();
+    assert_eq!(accruals.0.len(), target_accruals.len() as u32);
+    assert_eq!(accruals.1, target_accruals.iter().sum::<u128>());
 }