templar_common/
accumulator.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
use near_sdk::{json_types::U128, near, require};

use crate::asset::{AssetClass, FungibleAssetAmount};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[near(serializers = [borsh, json])]
pub struct Accumulator<T: AssetClass> {
    total: FungibleAssetAmount<T>,
    fraction_as_u128_dividend: U128,
    next_snapshot_index: u32,
    #[borsh(skip)]
    #[serde(default, skip_serializing_if = "FungibleAssetAmount::is_zero")]
    pub pending_estimate: FungibleAssetAmount<T>,
}

impl<T: AssetClass> Accumulator<T> {
    pub fn new(next_snapshot_index: u32) -> Self {
        Self {
            total: 0.into(),
            fraction_as_u128_dividend: U128(0),
            next_snapshot_index,
            pending_estimate: 0.into(),
        }
    }

    pub fn get_next_snapshot_index(&self) -> u32 {
        self.next_snapshot_index
    }

    pub fn get_total(&self) -> FungibleAssetAmount<T> {
        self.total
    }

    pub fn clear(&mut self, next_snapshot_index: u32) {
        self.total = 0.into();
        self.next_snapshot_index = next_snapshot_index;
    }

    pub fn remove(&mut self, amount: FungibleAssetAmount<T>) -> Option<FungibleAssetAmount<T>> {
        self.total.split(amount)
    }

    pub fn add_once(&mut self, amount: FungibleAssetAmount<T>) -> Option<()> {
        self.total.join(amount)?;
        Some(())
    }

    pub fn accumulate(
        &mut self,
        AccumulationRecord {
            mut amount,
            fraction_as_u128_dividend: fraction,
            next_snapshot_index,
        }: AccumulationRecord<T>,
    ) -> Option<()> {
        require!(
            next_snapshot_index >= self.next_snapshot_index,
            "Invariant violation: Asset accumulations cannot occur retroactively.",
        );
        let (fraction, carry) = self.fraction_as_u128_dividend.0.overflowing_add(fraction);
        if carry {
            amount.join(1u128)?;
        }
        self.add_once(amount)?;
        self.fraction_as_u128_dividend.0 = fraction;
        self.next_snapshot_index = next_snapshot_index;
        Some(())
    }
}

#[must_use]
#[derive(Debug, Clone)]
pub struct AccumulationRecord<T: AssetClass> {
    pub(crate) amount: FungibleAssetAmount<T>,
    pub(crate) fraction_as_u128_dividend: u128,
    pub(crate) next_snapshot_index: u32,
}

impl<T: AssetClass> AccumulationRecord<T> {
    pub fn get_amount(&self) -> FungibleAssetAmount<T> {
        self.amount
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fraction() {
        let mut a = Accumulator::<crate::asset::BorrowAsset>::new(1);

        a.accumulate(AccumulationRecord {
            amount: 100.into(),
            fraction_as_u128_dividend: 1 << 127,
            next_snapshot_index: 2,
        });

        assert_eq!(a.get_total(), 100.into());

        a.accumulate(AccumulationRecord {
            amount: 100.into(),
            fraction_as_u128_dividend: 1 << 127,
            next_snapshot_index: 3,
        });

        assert_eq!(a.get_total(), 201.into());
    }
}