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