1use std::collections::BTreeMap;
28
29#[derive(Debug, Default)]
33pub struct StreamMaxTracker {
34 pub per_stream: BTreeMap<u64, Vec<(f64, u64)>>,
36
37 pub flat: BTreeMap<u64, u64>,
39
40 pub sum_series: Vec<(f64, u64)>,
42
43 running_sum: u64,
45}
46
47impl StreamMaxTracker {
48 pub fn new() -> Self {
49 Self::default()
50 }
51
52 pub fn update(
54 &mut self, stream_id: u64, new_value: u64, ev_time: f64, init_val: u64,
55 ) -> Option<(u64, u64)> {
56 let entry = self
57 .per_stream
58 .entry(stream_id)
59 .or_insert_with(|| vec![(0.0, init_val)]);
60 entry.push((ev_time, new_value));
61
62 let old_max = self.flat.get(&stream_id).copied().unwrap_or(0);
63 let new_max = old_max.max(new_value);
64
65 let result = if new_max > old_max {
66 self.flat.insert(stream_id, new_max);
67 self.running_sum += new_max - old_max;
68 Some((old_max, new_max))
69 } else {
70 None
71 };
72
73 self.sum_series.push((ev_time, self.running_sum));
74
75 result
76 }
77
78 pub fn get_stream_max(&self, stream_id: u64) -> Option<u64> {
79 self.flat.get(&stream_id).copied()
80 }
81
82 pub fn current_sum(&self) -> u64 {
83 self.running_sum
84 }
85}
86
87#[cfg(test)]
88mod tests {
89 use super::*;
90
91 #[test]
92 fn test_tracker_empty() {
93 let tracker = StreamMaxTracker::new();
94 assert_eq!(tracker.current_sum(), 0);
95 assert_eq!(tracker.get_stream_max(0), None);
96 assert!(tracker.sum_series.is_empty());
97 assert!(tracker.flat.is_empty());
98 }
99
100 #[test]
101 fn test_tracker_single_stream_in_order() {
102 let mut tracker = StreamMaxTracker::new();
103
104 let result = tracker.update(0, 1000, 10.0, 0);
105 assert_eq!(result, Some((0, 1000)));
106 assert_eq!(tracker.get_stream_max(0), Some(1000));
107 assert_eq!(tracker.current_sum(), 1000);
108 assert_eq!(tracker.sum_series, vec![(10.0, 1000)]);
109
110 let result = tracker.update(0, 2000, 20.0, 0);
111 assert_eq!(result, Some((1000, 2000)));
112 assert_eq!(tracker.get_stream_max(0), Some(2000));
113 assert_eq!(tracker.current_sum(), 2000);
114 assert_eq!(tracker.sum_series, vec![(10.0, 1000), (20.0, 2000)]);
115 }
116
117 #[test]
118 fn test_tracker_out_of_order_events() {
119 let mut tracker = StreamMaxTracker::new();
120
121 let result = tracker.update(0, 2000, 20.0, 0);
122 assert_eq!(result, Some((0, 2000)));
123 assert_eq!(tracker.get_stream_max(0), Some(2000));
124 assert_eq!(tracker.current_sum(), 2000);
125
126 let result = tracker.update(0, 1000, 10.0, 0);
127 assert_eq!(result, None);
128 assert_eq!(tracker.get_stream_max(0), Some(2000));
129 assert_eq!(tracker.current_sum(), 2000);
130
131 assert_eq!(tracker.sum_series, vec![(20.0, 2000), (10.0, 2000)]);
132
133 let result = tracker.update(0, 3000, 30.0, 0);
134 assert_eq!(result, Some((2000, 3000)));
135 assert_eq!(tracker.get_stream_max(0), Some(3000));
136 assert_eq!(tracker.current_sum(), 3000);
137 }
138
139 #[test]
140 fn test_tracker_multiple_streams() {
141 let mut tracker = StreamMaxTracker::new();
142
143 tracker.update(0, 1000, 10.0, 0);
144 assert_eq!(tracker.current_sum(), 1000);
145
146 tracker.update(1, 500, 15.0, 0);
147 assert_eq!(tracker.current_sum(), 1500);
148
149 tracker.update(0, 1500, 20.0, 0);
150 assert_eq!(tracker.current_sum(), 2000);
151
152 tracker.update(2, 300, 25.0, 0);
153 assert_eq!(tracker.current_sum(), 2300);
154
155 assert_eq!(tracker.get_stream_max(0), Some(1500));
156 assert_eq!(tracker.get_stream_max(1), Some(500));
157 assert_eq!(tracker.get_stream_max(2), Some(300));
158 assert_eq!(tracker.get_stream_max(999), None);
159 }
160
161 #[test]
162 fn test_tracker_duplicate_values() {
163 let mut tracker = StreamMaxTracker::new();
164
165 let result = tracker.update(0, 1000, 10.0, 0);
166 assert_eq!(result, Some((0, 1000)));
167
168 let result = tracker.update(0, 1000, 20.0, 0);
169 assert_eq!(result, None);
170
171 assert_eq!(tracker.get_stream_max(0), Some(1000));
172 assert_eq!(tracker.current_sum(), 1000);
173
174 assert_eq!(tracker.sum_series.len(), 2);
175 assert_eq!(tracker.sum_series, vec![(10.0, 1000), (20.0, 1000)]);
176 }
177
178 #[test]
179 fn test_tracker_running_sum_correctness() {
180 let mut tracker = StreamMaxTracker::new();
181
182 tracker.update(0, 1000, 1.0, 0);
183 tracker.update(1, 2000, 2.0, 0);
184 tracker.update(2, 1500, 3.0, 0);
185 tracker.update(0, 1200, 4.0, 0);
186
187 let manual_sum: u64 = tracker.flat.values().sum();
188
189 assert_eq!(tracker.current_sum(), manual_sum);
190 assert_eq!(tracker.current_sum(), 4700);
191 }
192
193 #[test]
194 fn test_tracker_complex_interleaving() {
195 let mut tracker = StreamMaxTracker::new();
196
197 let updates = vec![
198 (0, 1000, 10.0),
199 (1, 500, 12.0),
200 (0, 800, 8.0),
201 (2, 2000, 15.0),
202 (1, 1000, 18.0),
203 (0, 1500, 20.0),
204 ];
205
206 for (stream_id, value, time) in updates {
207 tracker.update(stream_id, value, time, 0);
208 }
209
210 assert_eq!(tracker.get_stream_max(0), Some(1500));
211 assert_eq!(tracker.get_stream_max(1), Some(1000));
212 assert_eq!(tracker.get_stream_max(2), Some(2000));
213 assert_eq!(tracker.current_sum(), 4500);
214
215 assert_eq!(tracker.sum_series.len(), 6);
216 }
217
218 #[test]
219 fn test_tracker_init_value() {
220 let mut tracker = StreamMaxTracker::new();
221
222 tracker.update(0, 1000, 10.0, 500);
223
224 let stream_data = tracker.per_stream.get(&0).unwrap();
225 assert_eq!(stream_data.len(), 2);
226 assert_eq!(stream_data[0], (0.0, 500));
227 assert_eq!(stream_data[1], (10.0, 1000));
228
229 tracker.update(0, 1500, 20.0, 500);
230 let stream_data = tracker.per_stream.get(&0).unwrap();
231 assert_eq!(stream_data.len(), 3);
232 assert_eq!(stream_data[0], (0.0, 500));
233 assert_eq!(stream_data[1], (10.0, 1000));
234 assert_eq!(stream_data[2], (20.0, 1500));
235 }
236}