Skip to main content

netem_trace/
peek.rs

1//! One-item lookahead for traces.
2//!
3//! [`Peekable`] wraps a trace and lets you inspect, or modify, the next item
4//! before it is consumed. Every trace trait in this crate has a matching
5//! extension trait (e.g. [`DelayTracePeekExt`]) that adds a `peekable()` method
6//! to all of its implementors, and boxed trait objects (e.g.
7//! `Box<dyn DelayTrace>`) get the same method without any import.
8//!
9//! The wrapper implements the trace trait of the trace it wraps, so items are
10//! consumed with the usual method (e.g. `next_delay()`) and a [`Peekable`] can
11//! be used anywhere the original trace was.
12//!
13//! ## Examples
14//!
15//! ```
16//! # use netem_trace::model::StaticDelayConfig;
17//! # use netem_trace::{Delay, DelayTrace, Duration};
18//! use netem_trace::peek::DelayTracePeekExt;
19//!
20//! let mut trace = StaticDelayConfig::new()
21//!     .delay(Delay::from_millis(10))
22//!     .duration(Duration::from_secs(1))
23//!     .build()
24//!     .peekable();
25//! let item = (Delay::from_millis(10), Duration::from_secs(1));
26//! assert_eq!(trace.peek(), Some(&item));
27//! assert_eq!(trace.next_delay(), Some(item));
28//! assert_eq!(trace.peek(), None);
29//! ```
30//!
31//! A boxed trace, such as one built from a configuration, works the same way:
32//!
33//! ```
34//! # use netem_trace::model::{DelayTraceConfig, StaticDelayConfig};
35//! # use netem_trace::{Delay, DelayTrace, Duration};
36//! let config = Box::new(
37//!     StaticDelayConfig::new()
38//!         .delay(Delay::from_millis(10))
39//!         .duration(Duration::from_secs(1)),
40//! ) as Box<dyn DelayTraceConfig>;
41//! let mut trace = config.into_model().peekable();
42//! let item = (Delay::from_millis(10), Duration::from_secs(1));
43//! assert_eq!(trace.peek(), Some(&item));
44//! assert_eq!(trace.next_delay(), Some(item));
45//! ```
46use super::{Bandwidth, Delay, DuplicatePattern, Duration, LossPattern, RwndDecision};
47use super::{BwTrace, DelayPerPacketTrace, DelayTrace, DuplicateTrace, LossTrace, RwndTrace};
48
49/// A trace wrapper that can look ahead one item.
50///
51/// `Trace` is the wrapped trace and `Item` is what it yields. It is usually
52/// created with a `peekable()` method rather than with [`Peekable::new`]; see
53/// the [module documentation](self) for examples.
54///
55/// An item is pulled from the wrapped trace only when it is peeked or consumed,
56/// and a peeked item is held until it is consumed.
57pub 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    /// Wrap `inner`, using `next` to pull its items.
65    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    /// Return a reference to the next item without consuming it, or `None` if
81    /// the trace has ended.
82    pub fn peek(&mut self) -> Option<&Item> {
83        self.update();
84        self.peeked.as_ref()
85    }
86
87    /// Return a mutable reference to the next item without consuming it, or
88    /// `None` if the trace has ended. Changes are seen by the next consumer.
89    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        ///
105        #[doc = concat!("For `Box<dyn ", stringify!($trace), ">`, the same method is provided on the trait object itself.")]
106        pub trait $ext: $trace + Sized {
107            /// Wrap `self` in a [`Peekable`] that can look ahead one item.
108            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        // `Box<dyn $trace>` does not implement `$trace` itself, so the
116        // blanket impl above does not cover it. An inherent method on the
117        // trait object gives boxed traces the same `peekable()` call.
118        impl dyn $trace {
119            /// Wrap the boxed trace in a [`Peekable`] that can look ahead one item.
120            pub fn peekable(self: Box<Self>) -> Peekable<Box<Self>, $item> {
121                Peekable::new(self, |inner| inner.$next())
122            }
123        }
124
125        // Make the wrapper a drop-in replacement for the wrapped trace.
126        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    /// A trace that yields the given delays, each lasting one second.
172    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    // Accepts any trace, to check the wrapper is usable as a drop-in replacement.
192    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    // `DelayPerPacketTrace` shares the `next_delay` name with `DelayTrace` but
274    // yields a bare `Delay`, and `Box<dyn DelayPerPacketTrace>` has a `From`
275    // impl for every model, so check that its boxed form resolves on its own.
276    #[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}