templar_common/
accumulator.rs

1use near_sdk::{json_types::U128, near, require};
2
3use crate::asset::{AssetClass, FungibleAssetAmount};
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq)]
6#[near(serializers = [borsh, json])]
7pub struct Accumulator<T: AssetClass> {
8    total: FungibleAssetAmount<T>,
9    fraction_as_u128_dividend: U128,
10    next_snapshot_index: u32,
11}
12
13impl<T: AssetClass> Accumulator<T> {
14    pub fn new(next_snapshot_index: u32) -> Self {
15        Self {
16            total: 0.into(),
17            fraction_as_u128_dividend: U128(0),
18            next_snapshot_index,
19        }
20    }
21
22    pub fn get_next_snapshot_index(&self) -> u32 {
23        self.next_snapshot_index
24    }
25
26    pub fn get_total(&self) -> FungibleAssetAmount<T> {
27        self.total
28    }
29
30    pub fn clear(&mut self, next_snapshot_index: u32) {
31        self.total = 0.into();
32        self.next_snapshot_index = next_snapshot_index;
33    }
34
35    pub fn remove(&mut self, amount: FungibleAssetAmount<T>) {
36        self.total -= amount;
37    }
38
39    pub fn add_once(&mut self, amount: FungibleAssetAmount<T>) {
40        self.total += amount;
41    }
42
43    pub fn accumulate(
44        &mut self,
45        AccumulationRecord {
46            mut amount,
47            fraction_as_u128_dividend: fraction,
48            next_snapshot_index,
49        }: AccumulationRecord<T>,
50    ) {
51        require!(
52            next_snapshot_index >= self.next_snapshot_index,
53            "Invariant violation: Asset accumulations cannot occur retroactively.",
54        );
55        let (fraction, carry) = self.fraction_as_u128_dividend.0.overflowing_add(fraction);
56        if carry {
57            amount += 1;
58        }
59        self.add_once(amount);
60        self.fraction_as_u128_dividend.0 = fraction;
61        self.next_snapshot_index = next_snapshot_index;
62    }
63}
64
65#[must_use]
66#[derive(Debug, Clone)]
67pub struct AccumulationRecord<T: AssetClass> {
68    pub(crate) amount: FungibleAssetAmount<T>,
69    pub(crate) fraction_as_u128_dividend: u128,
70    pub(crate) next_snapshot_index: u32,
71}
72
73impl<T: AssetClass> AccumulationRecord<T> {
74    pub fn get_amount(&self) -> FungibleAssetAmount<T> {
75        self.amount
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn fraction() {
85        let mut a = Accumulator::<crate::asset::BorrowAsset>::new(1);
86
87        a.accumulate(AccumulationRecord {
88            amount: 100.into(),
89            fraction_as_u128_dividend: 1 << 127,
90            next_snapshot_index: 2,
91        });
92
93        assert_eq!(a.get_total(), 100.into());
94
95        a.accumulate(AccumulationRecord {
96            amount: 100.into(),
97            fraction_as_u128_dividend: 1 << 127,
98            next_snapshot_index: 3,
99        });
100
101        assert_eq!(a.get_total(), 201.into());
102    }
103}