iota_common/
moving_window.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::VecDeque, fmt::Debug, time::Duration};
6
7/// A moving window that maintains the last N values of type `T` and calculates
8/// their arithmetic mean. All values in the window have equal weight and the
9/// oldest value is dropped when the window exceeds its configured capacity.
10#[derive(Debug, Clone)]
11pub struct MovingWindow<T: MovingWindowValue> {
12    values: VecDeque<T>,
13    max_size: usize,
14    sum: T,
15}
16
17impl<T: MovingWindowValue> MovingWindow<T> {
18    /// Creates a new `MovingWindow` with the specified maximum size and an
19    /// `init_value`. The provided `max_size` must be greater than 0.
20    pub fn new(init_value: T, max_size: usize) -> Self {
21        assert!(max_size > 0, "Window size must be greater than 0");
22        let mut window = Self {
23            values: VecDeque::with_capacity(max_size),
24            max_size,
25            sum: T::zero(),
26        };
27        window.add_value(init_value);
28        window
29    }
30
31    /// Adds a new value to the window. If the window is at capacity, the oldest
32    /// value is removed before adding the new value.
33    pub fn add_value(&mut self, value: T) {
34        if self.values.len() == self.max_size {
35            if let Some(old_value) = self.values.pop_front() {
36                T::sub_assign(&mut self.sum, old_value);
37            }
38        }
39
40        self.values.push_back(value);
41        T::add_assign(&mut self.sum, value);
42    }
43
44    /// Get the current average of all values in the window. Returns the value's
45    /// zero if the window is empty.
46    pub fn get(&self) -> T {
47        if self.values.is_empty() {
48            T::zero()
49        } else {
50            T::average(self.sum, self.values.len())
51        }
52    }
53
54    /// Get the number of values currently in the window.
55    pub fn len(&self) -> usize {
56        self.values.len()
57    }
58
59    pub fn is_empty(&self) -> bool {
60        self.values.is_empty()
61    }
62}
63
64pub trait MovingWindowValue: Copy + Debug {
65    fn zero() -> Self;
66    fn add_assign(target: &mut Self, value: Self);
67    fn sub_assign(target: &mut Self, value: Self);
68    fn average(total: Self, divisor: usize) -> Self;
69}
70
71impl MovingWindowValue for Duration {
72    fn zero() -> Self {
73        Duration::ZERO
74    }
75
76    fn add_assign(target: &mut Self, value: Self) {
77        *target += value;
78    }
79
80    fn sub_assign(target: &mut Self, value: Self) {
81        *target -= value;
82    }
83
84    fn average(total: Self, divisor: usize) -> Self {
85        let divisor = u32::try_from(divisor).expect("window size too large for Duration average");
86        total / divisor
87    }
88}
89
90impl MovingWindowValue for f64 {
91    fn zero() -> Self {
92        0.0
93    }
94
95    fn add_assign(target: &mut Self, value: Self) {
96        *target += value;
97    }
98
99    fn sub_assign(target: &mut Self, value: Self) {
100        *target -= value;
101    }
102
103    fn average(total: Self, divisor: usize) -> Self {
104        total / divisor as f64
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn test_with_initial_value() {
114        let window = MovingWindow::new(10.0, 3);
115        assert_eq!(window.len(), 1);
116        assert_eq!(window.get(), 10.0);
117    }
118
119    #[test]
120    fn test_duration_window() {
121        let mut window = MovingWindow::new(Duration::ZERO, 3);
122        assert_eq!(window.get(), Duration::ZERO);
123
124        window.add_value(Duration::from_millis(100));
125        assert_eq!(window.get(), Duration::from_millis(50));
126        assert_eq!(window.len(), 2);
127
128        window.add_value(Duration::from_millis(200));
129        assert_eq!(window.get(), Duration::from_millis(100));
130        assert_eq!(window.len(), 3);
131
132        window.add_value(Duration::from_millis(300));
133        assert_eq!(window.get(), Duration::from_millis(200));
134        assert_eq!(window.len(), 3);
135
136        window.add_value(Duration::from_millis(400));
137        assert_eq!(window.get(), Duration::from_millis(300));
138        assert_eq!(window.len(), 3);
139    }
140
141    #[test]
142    fn test_float_window() {
143        let mut window = MovingWindow::new(0.0_f64, 3);
144        assert_eq!(window.get(), 0.0);
145
146        window.add_value(1.0);
147        assert_eq!(window.get(), 0.5);
148        assert_eq!(window.len(), 2);
149
150        window.add_value(2.0);
151        assert_eq!(window.get(), 1.0);
152        assert_eq!(window.len(), 3);
153
154        window.add_value(3.0);
155        assert_eq!(window.get(), 2.0);
156        assert_eq!(window.len(), 3);
157
158        window.add_value(4.0);
159        assert_eq!(window.get(), 3.0);
160        assert_eq!(window.len(), 3);
161    }
162
163    #[test]
164    #[should_panic(expected = "Window size must be greater than 0")]
165    fn test_zero_size_panics() {
166        let _window = MovingWindow::new(0.0_f64, 0);
167    }
168}
iota_common/moving_window.rs

iota_common/
moving_window.rs
