1use super::{Bandwidth, Delay, DuplicatePattern, Duration, LossPattern, RwndDecision};
47use super::{BwTrace, DelayPerPacketTrace, DelayTrace, DuplicateTrace, LossTrace, RwndTrace};
48
49pub struct Peekable<Trace, Item> {
58 inner: Trace,
59 next: fn(&mut Trace) -> Option<Item>,
60 peeked: Option<Item>,
61}
62
63impl<Trace, Item> Peekable<Trace, Item> {
64 pub fn new(inner: Trace, next: fn(&mut Trace) -> Option<Item>) -> Self {
66 Self {
67 inner,
68 next,
69 peeked: None,
70 }
71 }
72
73 #[inline(always)]
74 fn update(&mut self) {
75 if self.peeked.is_none() {
76 self.peeked = (self.next)(&mut self.inner);
77 }
78 }
79
80 pub fn peek(&mut self) -> Option<&Item> {
83 self.update();
84 self.peeked.as_ref()
85 }
86
87 pub fn peek_mut(&mut self) -> Option<&mut Item> {
90 self.update();
91 self.peeked.as_mut()
92 }
93
94 #[inline(always)]
95 fn take_next(&mut self) -> Option<Item> {
96 self.update();
97 self.peeked.take()
98 }
99}
100
101macro_rules! impl_peekable_for {
102 ($ext:ident, $trace:ident, $next:ident, $item:ty) => {
103 #[doc = concat!("Adds `peekable()` to every type that implements [`", stringify!($trace), "`].")]
104 #[doc = concat!("For `Box<dyn ", stringify!($trace), ">`, the same method is provided on the trait object itself.")]
106 pub trait $ext: $trace + Sized {
107 fn peekable(self) -> Peekable<Self, $item> {
109 Peekable::new(self, <Self as $trace>::$next)
110 }
111 }
112
113 impl<T: $trace> $ext for T {}
114
115 impl dyn $trace {
119 pub fn peekable(self: Box<Self>) -> Peekable<Box<Self>, $item> {
121 Peekable::new(self, |inner| inner.$next())
122 }
123 }
124
125 impl<T: $trace> $trace for Peekable<T, $item> {
127 fn $next(&mut self) -> Option<$item> {
128 self.take_next()
129 }
130 }
131
132 impl $trace for Peekable<Box<dyn $trace>, $item> {
133 fn $next(&mut self) -> Option<$item> {
134 self.take_next()
135 }
136 }
137 };
138}
139
140impl_peekable_for!(BwTracePeekExt, BwTrace, next_bw, (Bandwidth, Duration));
141impl_peekable_for!(DelayTracePeekExt, DelayTrace, next_delay, (Delay, Duration));
142impl_peekable_for!(
143 DelayPerPacketTracePeekExt,
144 DelayPerPacketTrace,
145 next_delay,
146 Delay
147);
148impl_peekable_for!(
149 LossTracePeekExt,
150 LossTrace,
151 next_loss,
152 (LossPattern, Duration)
153);
154impl_peekable_for!(
155 DuplicateTracePeekExt,
156 DuplicateTrace,
157 next_duplicate,
158 (DuplicatePattern, Duration)
159);
160impl_peekable_for!(
161 RwndTracePeekExt,
162 RwndTrace,
163 next_rwnd,
164 (RwndDecision, Duration)
165);
166
167#[cfg(test)]
168mod test {
169 use super::*;
170
171 struct VecDelay(std::vec::IntoIter<Delay>);
173
174 impl VecDelay {
175 fn new(millis: &[u64]) -> Self {
176 let delays: Vec<_> = millis.iter().map(|&ms| Delay::from_millis(ms)).collect();
177 Self(delays.into_iter())
178 }
179 }
180
181 impl DelayTrace for VecDelay {
182 fn next_delay(&mut self) -> Option<(Delay, Duration)> {
183 self.0.next().map(|delay| (delay, Duration::from_secs(1)))
184 }
185 }
186
187 fn item(ms: u64) -> (Delay, Duration) {
188 (Delay::from_millis(ms), Duration::from_secs(1))
189 }
190
191 fn drain(trace: &mut impl DelayTrace) -> Vec<(Delay, Duration)> {
193 std::iter::from_fn(|| trace.next_delay()).collect()
194 }
195
196 #[test]
197 fn test_peek_does_not_consume() {
198 let mut trace = VecDelay::new(&[10, 20]).peekable();
199 assert_eq!(trace.peek(), Some(&item(10)));
200 assert_eq!(trace.peek(), Some(&item(10)));
201 assert_eq!(trace.next_delay(), Some(item(10)));
202 assert_eq!(trace.peek(), Some(&item(20)));
203 assert_eq!(trace.next_delay(), Some(item(20)));
204 assert_eq!(trace.peek(), None);
205 assert_eq!(trace.next_delay(), None);
206 }
207
208 #[test]
209 fn test_peek_mut_modifies_next() {
210 let mut trace = VecDelay::new(&[10, 20]).peekable();
211 if let Some((delay, _)) = trace.peek_mut() {
212 *delay = Delay::from_millis(15);
213 }
214 assert_eq!(trace.next_delay(), Some(item(15)));
215 assert_eq!(trace.next_delay(), Some(item(20)));
216 }
217
218 #[test]
219 fn test_peekable_is_drop_in_trace() {
220 let mut trace = VecDelay::new(&[10, 20, 30]).peekable();
221 assert_eq!(trace.peek(), Some(&item(10)));
222 assert_eq!(drain(&mut trace), vec![item(10), item(20), item(30)]);
223 }
224
225 #[test]
226 fn test_boxed_dyn_trace() {
227 let boxed: Box<dyn DelayTrace> = Box::new(VecDelay::new(&[10, 20]));
228 let mut trace = boxed.peekable();
229 assert_eq!(trace.peek(), Some(&item(10)));
230 assert_eq!(drain(&mut trace), vec![item(10), item(20)]);
231 }
232
233 #[test]
234 fn test_boxed_peekable_as_dyn_trace() {
235 let boxed: Box<dyn DelayTrace> = Box::new(VecDelay::new(&[10, 20]));
236 let mut trace: Box<dyn DelayTrace> = Box::new(boxed.peekable());
237 assert_eq!(trace.next_delay(), Some(item(10)));
238 assert_eq!(trace.next_delay(), Some(item(20)));
239 assert_eq!(trace.next_delay(), None);
240 }
241
242 #[cfg(feature = "delay-model")]
243 mod delay_model {
244 use super::*;
245 use crate::model::{DelayTraceConfig, StaticDelayConfig};
246
247 #[test]
248 fn test_static_delay() {
249 let mut trace = StaticDelayConfig::new()
250 .delay(Delay::from_millis(10))
251 .duration(Duration::from_secs(1))
252 .build()
253 .peekable();
254 assert_eq!(trace.peek(), Some(&item(10)));
255 assert_eq!(trace.next_delay(), Some(item(10)));
256 assert_eq!(trace.peek(), None);
257 assert_eq!(trace.next_delay(), None);
258 }
259
260 #[test]
261 fn test_model_from_config() {
262 let config = Box::new(
263 StaticDelayConfig::new()
264 .delay(Delay::from_millis(10))
265 .duration(Duration::from_secs(1)),
266 ) as Box<dyn DelayTraceConfig>;
267 let mut trace = config.into_model().peekable();
268 assert_eq!(trace.peek(), Some(&item(10)));
269 assert_eq!(drain(&mut trace), vec![item(10)]);
270 }
271 }
272
273 #[cfg(feature = "delay-per-packet-model")]
277 mod delay_per_packet_model {
278 use super::*;
279 use crate::model::StaticDelayPerPacketConfig;
280
281 #[test]
282 fn test_boxed_static_delay_per_packet() {
283 let boxed: Box<dyn DelayPerPacketTrace> = StaticDelayPerPacketConfig::new()
284 .delay(Delay::from_millis(10))
285 .count(2)
286 .build()
287 .into();
288 let mut trace = boxed.peekable();
289 assert_eq!(trace.peek(), Some(&Delay::from_millis(10)));
290 assert_eq!(trace.next_delay(), Some(Delay::from_millis(10)));
291 assert_eq!(trace.next_delay(), Some(Delay::from_millis(10)));
292 assert_eq!(trace.peek(), None);
293 }
294 }
295}