tycho_util/num/
safe_avg.rs

1#[derive(Default)]
2pub struct SafeUnsignedAvg {
3    sum: u128,
4    counter: u128,
5}
6
7impl SafeUnsignedAvg {
8    pub fn with_initial<V: Into<u128>>(v: V) -> Self {
9        Self {
10            sum: v.into(),
11            counter: 1,
12        }
13    }
14
15    pub fn accum<V: Into<u128>>(&mut self, v: V) {
16        let v = v.into();
17        self.sum = self.sum.checked_add(v).unwrap_or_else(|| {
18            let partial = self.sum / self.counter;
19            self.counter = 0;
20            partial.saturating_add(v)
21        });
22        self.counter += 1;
23    }
24
25    pub fn get_avg(&self) -> u128 {
26        self.get_avg_checked()
27            .expect("should check first with `get_avg_checked`")
28    }
29
30    pub fn get_avg_checked(&self) -> Option<u128> {
31        if self.counter == 0 {
32            None
33        } else if self.counter == 1 {
34            Some(self.sum)
35        } else {
36            Some(self.sum.saturating_div(self.counter))
37        }
38    }
39}
40
41#[derive(Default)]
42pub struct SafeSignedAvg {
43    sum: i128,
44    counter: i128,
45}
46
47impl SafeSignedAvg {
48    pub fn with_initial<V: Into<i128>>(v: V) -> Self {
49        Self {
50            sum: v.into(),
51            counter: 1,
52        }
53    }
54
55    pub fn accum<V: Into<i128>>(&mut self, v: V) {
56        let v = v.into();
57        self.sum = self.sum.checked_add(v).unwrap_or_else(|| {
58            let partial = self.sum / self.counter;
59            self.counter = 0;
60            partial.saturating_add(v)
61        });
62        self.counter += 1;
63    }
64
65    pub fn get_avg(&self) -> i128 {
66        self.get_avg_checked()
67            .expect("should check first with `get_avg_checked`")
68    }
69
70    pub fn get_avg_checked(&self) -> Option<i128> {
71        if self.counter == 0 {
72            None
73        } else if self.counter == 1 {
74            Some(self.sum)
75        } else {
76            Some(self.sum.saturating_div(self.counter))
77        }
78    }
79}
80
81pub struct SafeUnsignedVecAvg {
82    idx: usize,
83    counters: Vec<u128>,
84    sums: Vec<u128>,
85}
86
87impl SafeUnsignedVecAvg {
88    pub fn new(capacity: usize) -> Self {
89        Self {
90            idx: 0,
91            counters: vec![0; capacity],
92            sums: vec![0; capacity],
93        }
94    }
95
96    pub fn accum<V: Into<u128>>(&mut self, idx: usize, v: V) {
97        self.idx = idx;
98        self.accum_curr(v);
99    }
100
101    pub fn accum_next<V: Into<u128>>(&mut self, v: V) {
102        self.idx += 1;
103        self.accum_curr(v);
104    }
105
106    fn accum_curr<V: Into<u128>>(&mut self, v: V) {
107        let v = v.into();
108        self.sums[self.idx] = self.sums[self.idx].checked_add(v).unwrap_or_else(|| {
109            let partial = self.sums[self.idx] / self.counters[self.idx];
110            self.counters[self.idx] = 0;
111            partial.saturating_add(v)
112        });
113        self.counters[self.idx] += 1;
114    }
115
116    pub fn get_avg(&mut self, idx: usize) -> u128 {
117        self.get_avg_checked(idx)
118            .expect("should check first with `get_avg_checked`")
119    }
120
121    pub fn get_avg_checked(&mut self, idx: usize) -> Option<u128> {
122        self.idx = idx;
123        self.get_avg_curr()
124    }
125
126    pub fn get_avg_next(&mut self) -> u128 {
127        self.get_avg_next_checked().unwrap()
128    }
129
130    pub fn get_avg_next_checked(&mut self) -> Option<u128> {
131        self.idx += 1;
132        self.get_avg_curr()
133    }
134
135    fn get_avg_curr(&self) -> Option<u128> {
136        if self.counters[self.idx] == 0 {
137            None
138        } else {
139            Some(self.sums[self.idx].saturating_div(self.counters[self.idx]))
140        }
141    }
142}
tycho_util/num/safe_avg.rs

tycho_util/num/
safe_avg.rs
