Skip to main content

des_sim/modeling/
sampler.rs

1//! The `sampler` module provides traits and implementations for generating random
2//! or pseudo-random durations within the simulation.
3//!
4//! It includes `DurationSampler` for basic sampling, `PendingDuration` for handling
5//! floating-point durations with rounding, and `CombinatorExt` for building complex
6//! sampling logic through method chaining.
7
8pub mod combinator;
9pub mod instance;
10
11use std::ops::{Add, Sub};
12
13use crate::primitive::time::{Duration, SimTime, TimeTick};
14use combinator::*;
15use rand::Rng;
16use rand_distr::num_traits::ToPrimitive;
17
18/// A wrapper for floating-point representations of [Duration] to mitigate rounding errors.
19#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
20pub struct PendingDuration(f64);
21
22impl From<Duration> for PendingDuration {
23    fn from(value: Duration) -> Self {
24        PendingDuration::from_duration(value)
25    }
26}
27
28impl Add<Duration> for PendingDuration {
29    type Output = PendingDuration;
30
31    fn add(self, rhs: Duration) -> Self::Output {
32        let result = self.0 + rhs.as_time_tick() as f64;
33        if result.is_finite() {
34            PendingDuration::new(result)
35        } else {
36            panic!("attempt to add with overflow")
37        }
38    }
39}
40
41impl Add<PendingDuration> for PendingDuration {
42    type Output = PendingDuration;
43
44    fn add(self, rhs: PendingDuration) -> Self::Output {
45        let result = self.0 + rhs.0;
46        if result.is_finite() {
47            PendingDuration::new(result)
48        } else {
49            panic!("attempt to add with overflow")
50        }
51    }
52}
53
54impl Sub<Duration> for PendingDuration {
55    type Output = PendingDuration;
56
57    fn sub(self, rhs: Duration) -> Self::Output {
58        let result = self.0 - rhs.as_time_tick() as f64;
59        if result.is_finite() {
60            PendingDuration::new(result)
61        } else {
62            panic!("attempt to subtract with overflow")
63        }
64    }
65}
66
67impl Sub<PendingDuration> for PendingDuration {
68    type Output = PendingDuration;
69
70    fn sub(self, rhs: PendingDuration) -> Self::Output {
71        let result = self.0 - rhs.0;
72        if result.is_finite() {
73            PendingDuration::new(result)
74        } else {
75            panic!("attempt to subtract with overflow")
76        }
77    }
78}
79
80impl PendingDuration {
81    /// Creates a new `PendingDuration` from a raw `f64` value.
82    pub fn new(value: f64) -> Self {
83        Self(value)
84    }
85
86    /// Creates a new `PendingDuration` from a [Duration].
87    pub fn from_duration(duration: Duration) -> Self {
88        Self::new(duration.as_time_tick() as f64)
89    }
90
91    /// Returns the raw floating-point value.
92    pub fn raw_value(&self) -> f64 {
93        self.0
94    }
95
96    /// Attempts to convert this instance into a [Duration].
97    /// Returns `None` if the value is negative.
98    pub fn try_to_duration(&self) -> Option<Duration> {
99        // Rounding and integer conversion occur here.
100        Some(Duration::ticks(self.raw_value().round().to_usize()?))
101    }
102
103    /// Converts this instance into a [Duration].
104    ///
105    /// # Panics
106    ///
107    /// Panics if the value is negative.
108    pub fn to_duration(&self) -> Duration {
109        self.try_to_duration().unwrap_or_else(|| {
110            panic!(
111                "Duration does not support negative values: {}",
112                self.raw_value()
113            )
114        })
115    }
116
117    /// Converts to a [Duration], clamping the result between 0 and `max` ticks.
118    pub fn to_duration_with_clamp(&self, max: TimeTick) -> Duration {
119        let raw_value = self.raw_value();
120        if raw_value >= max as f64 {
121            Duration::ticks(max)
122        } else {
123            Duration::ticks(
124                raw_value
125                    .round()
126                    .clamp(0.0, max as f64)
127                    .to_usize()
128                    .unwrap_or_else(|| {
129                        panic!("Unexpected clamping error for value: {}", raw_value)
130                    }),
131            )
132        }
133    }
134
135    /// Converts to a [Duration], invoking the provided closure if the conversion fails.
136    pub fn to_duration_or_else<F>(&self, f: F) -> Duration
137    where
138        F: FnOnce() -> Duration,
139    {
140        self.try_to_duration().unwrap_or_else(f)
141    }
142
143    /// Applies a function to the internal value and returns a new `PendingDuration`.
144    pub fn apply<F>(&mut self, f: F) -> PendingDuration
145    where
146        F: Fn(f64) -> f64,
147    {
148        PendingDuration::new(f(self.0))
149    }
150}
151
152/// A trait for generating [Duration] samples.
153pub trait DurationSampler {
154    /// Generates a sample given a random number generator and the current simulation time.
155    fn sample(&mut self, rng: &mut dyn Rng, current_tick: SimTime) -> PendingDuration;
156}
157
158/// A trait for cloning trait objects of [DurationSampler].
159pub trait ClonableDurationSampler: DurationSampler + Send + Sync {
160    fn box_clone(&self) -> Box<dyn ClonableDurationSampler>;
161}
162
163// Automatically implement `ClonableDurationSampler` for all compatible types.
164impl<S> ClonableDurationSampler for S
165where
166    S: DurationSampler + Clone + Send + Sync + 'static,
167{
168    fn box_clone(&self) -> Box<dyn ClonableDurationSampler> {
169        Box::new(self.clone())
170    }
171}
172
173// Enable cloning for boxed trait objects.
174impl Clone for Box<dyn ClonableDurationSampler> {
175    fn clone(&self) -> Self {
176        self.box_clone()
177    }
178}
179
180/// An extension trait for [DurationSampler] that provides fluent combinator methods
181/// for creating and composing complex duration sampling logic.
182pub trait CombinatorExt: DurationSampler + Sized + 'static {
183    /// Boxes the sampler into a `Box<dyn DurationSampler>` to erase its concrete type.
184    fn boxed(self) -> Box<dyn DurationSampler> {
185        Box::new(self)
186    }
187
188    /// Boxes the sampler as a cloneable trait object.
189    /// Requires the underlying sampler to be `Clone`, `Send`, and `Sync`.
190    fn boxed_clonable(self) -> Box<dyn ClonableDurationSampler>
191    where
192        Self: Clone + Send + Sync,
193    {
194        Box::new(self)
195    }
196
197    /// Transforms the output of this sampler using the provided closure `f`.
198    fn map<F>(self, f: F) -> MapSampler<Self, F>
199    where
200        F: FnMut(&mut dyn Rng, SimTime, f64) -> f64,
201    {
202        MapSampler::new(self, f)
203    }
204
205    /// Adds a delay to the result of this sampler using another sampler for the duration.
206    fn delay(self, delay: Box<dyn DurationSampler>) -> DelaySampler<Self> {
207        DelaySampler::new(self, delay)
208    }
209
210    /// Adds a jitter (noise) to the result of this sampler using another sampler for the variation.
211    fn jitter(self, jitter: Box<dyn DurationSampler>) -> JitterSampler<Self> {
212        JitterSampler::new(self, jitter)
213    }
214
215    /// Chains this sampler with another, combining their outputs using the provided closure `f`.
216    fn chain<F>(self, sampler: Box<dyn DurationSampler>, f: F) -> ChainSampler<Self, F>
217    where
218        F: FnMut(&mut dyn Rng, SimTime, f64, f64) -> f64,
219    {
220        ChainSampler::new(self, sampler, f)
221    }
222
223    /// Aggregates the results of this sampler and a collection of others into a single value
224    /// using the provided closure `f`.
225    fn aggregate<F>(
226        self,
227        others: impl IntoIterator<Item = Box<dyn DurationSampler>>,
228        f: F,
229    ) -> AggregateSampler<F>
230    where
231        F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
232    {
233        let mut samplers = vec![self.boxed()];
234        samplers.extend(others);
235
236        AggregateSampler::new(samplers, f)
237    }
238
239    /// Returns an `AggregateBuilder` initialized with this sampler, allowing for
240    /// a more flexible construction of aggregate samplers.
241    fn aggregate_builder(self) -> AggregateBuilder {
242        AggregateBuilder::from_sampler(self)
243    }
244
245    /// Clamps the output of this sampler within the range `[min, max]`.
246    fn clamp(self, min: f64, max: f64) -> ClampSampler<Self> {
247        ClampSampler::new(self, min, max)
248    }
249
250    /// Ensures the sampled duration is non-negative.
251    ///
252    /// If the value remains negative after `limit_try_count` attempts, the `fallback` closure is invoked.
253    fn ensure_non_negative<F>(
254        self,
255        limit_try_count: u8,
256        fallback: F,
257    ) -> EnsureNonNegativeSampler<Self, F>
258    where
259        F: FnMut(&mut dyn Rng, SimTime) -> Duration,
260    {
261        EnsureNonNegativeSampler::new(self, limit_try_count, fallback)
262    }
263}
264
265impl<T: DurationSampler + Sized + 'static> CombinatorExt for T {}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270    use rand::SeedableRng;
271    use rand::prelude::SmallRng;
272
273    #[test]
274    fn test_pending_duration_new() {
275        let pd = PendingDuration::new(10.5);
276        assert_eq!(pd.raw_value(), 10.5);
277    }
278
279    #[test]
280    fn test_pending_duration_from_duration() {
281        let d = Duration::ticks(100);
282        let pd = PendingDuration::from_duration(d);
283        assert_eq!(pd.raw_value(), 100.0);
284    }
285
286    #[test]
287    fn test_pending_duration_from_trait() {
288        let d = Duration::ticks(200);
289        let pd: PendingDuration = d.into();
290        assert_eq!(pd.raw_value(), 200.0);
291    }
292
293    #[test]
294    fn test_pending_duration_to_duration() {
295        let pd = PendingDuration::new(10.5);
296        let d = pd.to_duration();
297        assert_eq!(d, Duration::ticks(11)); // Rounds up
298    }
299
300    #[test]
301    fn test_pending_duration_to_duration_round_down() {
302        let pd = PendingDuration::new(10.4);
303        let d = pd.to_duration();
304        assert_eq!(d, Duration::ticks(10)); // Rounds down
305    }
306
307    #[test]
308    fn test_pending_duration_to_duration_negative_panics() {
309        let pd = PendingDuration::new(-5.0);
310        let result = std::panic::catch_unwind(|| pd.to_duration());
311        assert!(result.is_err());
312    }
313
314    #[test]
315    fn test_pending_duration_to_duration_with_clamp() {
316        let pd = PendingDuration::new(10.5);
317        let d = pd.to_duration_with_clamp(100);
318        assert_eq!(d, Duration::ticks(11));
319
320        let pd_large = PendingDuration::new(150.0);
321        let d_clamped = pd_large.to_duration_with_clamp(100);
322        assert_eq!(d_clamped, Duration::ticks(100));
323
324        let pd_negative = PendingDuration::new(-5.0);
325        let d_clamped_negative = pd_negative.to_duration_with_clamp(100);
326        assert_eq!(d_clamped_negative, Duration::ticks(0));
327    }
328
329    #[test]
330    fn test_pending_duration_to_duration_or_else() {
331        let pd_positive = PendingDuration::new(5.5);
332        let d = pd_positive.to_duration_or_else(|| Duration::ticks(0));
333        assert_eq!(d, Duration::ticks(6));
334
335        let pd_negative = PendingDuration::new(-5.5);
336        let d_else = pd_negative.to_duration_or_else(|| Duration::ticks(10));
337        assert_eq!(d_else, Duration::ticks(10));
338    }
339
340    #[test]
341    fn test_pending_duration_try_duration() {
342        let pd_positive = PendingDuration::new(5.5);
343        assert_eq!(pd_positive.try_to_duration(), Some(Duration::ticks(6)));
344
345        let pd_negative = PendingDuration::new(-5.5);
346        assert_eq!(pd_negative.try_to_duration(), None);
347    }
348
349    #[test]
350    fn test_pending_duration_add_duration() {
351        let pd = PendingDuration::new(10.0);
352        let d = Duration::ticks(5);
353        let result = pd + d;
354        assert_eq!(result.raw_value(), 15.0);
355    }
356
357    #[test]
358    fn test_pending_duration_add_pending_duration() {
359        let pd1 = PendingDuration::new(10.0);
360        let pd2 = PendingDuration::new(7.5);
361        let result = pd1 + pd2;
362        assert_eq!(result.raw_value(), 17.5);
363    }
364
365    #[test]
366    fn test_pending_duration_sub_duration() {
367        let pd = PendingDuration::new(10.0);
368        let d = Duration::ticks(3);
369        let result = pd - d;
370        assert_eq!(result.raw_value(), 7.0);
371    }
372
373    #[test]
374    fn test_pending_duration_sub_pending_duration() {
375        let pd1 = PendingDuration::new(10.0);
376        let pd2 = PendingDuration::new(2.5);
377        let result = pd1 - pd2;
378        assert_eq!(result.raw_value(), 7.5);
379    }
380
381    #[test]
382    fn test_pending_duration_apply() {
383        let mut pd = PendingDuration::new(10.0);
384        let result = pd.apply(|v| v * 2.0);
385        assert_eq!(result.raw_value(), 20.0);
386        assert_eq!(pd.raw_value(), 10.0); // Original value remains unchanged
387    }
388
389    /// Mock sampler for testing `CombinatorExt` functionality.
390    struct MockSampler {
391        value: f64,
392    }
393
394    impl DurationSampler for MockSampler {
395        fn sample(&mut self, _rng: &mut dyn Rng, _current_tick: SimTime) -> PendingDuration {
396            PendingDuration::new(self.value)
397        }
398    }
399
400    #[test]
401    fn test_combinator_ext_boxed() {
402        let mut sampler = MockSampler { value: 10.0 };
403        let mut rng = SmallRng::seed_from_u64(2);
404        let current_tick = SimTime::from_ticks(0);
405        assert_eq!(sampler.sample(&mut rng, current_tick).raw_value(), 10.0);
406    }
407
408    #[test]
409    fn test_combinator_ext_map() {
410        let sampler = MockSampler { value: 10.0 };
411        let mut map_sampler = sampler.map(|_rng, _tick, v| v * 2.0);
412        let mut rng = SmallRng::seed_from_u64(2);
413        let current_tick = SimTime::from_ticks(0);
414        assert_eq!(map_sampler.sample(&mut rng, current_tick).raw_value(), 20.0);
415    }
416
417    #[test]
418    fn test_combinator_ext_delay() {
419        let base_sampler = MockSampler { value: 10.0 };
420        let delay_sampler_impl = MockSampler { value: 5.0 };
421        let mut delay_sampler = base_sampler.delay(delay_sampler_impl.boxed());
422        let mut rng = SmallRng::seed_from_u64(2);
423        let current_tick = SimTime::from_ticks(0);
424        assert_eq!(
425            delay_sampler.sample(&mut rng, current_tick).raw_value(),
426            15.0
427        ); // base + delay
428    }
429
430    #[test]
431    fn test_combinator_ext_jitter() {
432        let base_sampler = MockSampler { value: 10.0 };
433        let jitter_sampler_impl = MockSampler { value: 2.0 };
434        let mut jitter_sampler = base_sampler.jitter(jitter_sampler_impl.boxed());
435        let mut rng = SmallRng::seed_from_u64(2);
436        let current_tick = SimTime::from_ticks(0);
437        // Jitter adds or subtracts; for a fixed mock value, we expect base + jitter.
438        assert_eq!(
439            jitter_sampler.sample(&mut rng, current_tick).raw_value(),
440            12.0
441        );
442    }
443
444    #[test]
445    fn test_combinator_ext_chain() {
446        let s1 = MockSampler { value: 10.0 };
447        let s2 = MockSampler { value: 5.0 };
448        let mut chain_sampler = s1.chain(s2.boxed(), |_rng, _tick, v1, v2| v1 + v2 * 2.0);
449        let mut rng = SmallRng::seed_from_u64(2);
450        let current_tick = SimTime::from_ticks(0);
451        assert_eq!(
452            chain_sampler.sample(&mut rng, current_tick).raw_value(),
453            20.0
454        ); // 10.0 + 5.0 * 2.0
455    }
456
457    #[test]
458    fn test_combinator_ext_aggregate() {
459        let s1 = MockSampler { value: 10.0 };
460        let s2 = MockSampler { value: 5.0 };
461        let s3 = MockSampler { value: 2.0 };
462        let others = vec![s2.boxed(), s3.boxed()];
463        let mut aggregate_sampler =
464            s1.aggregate(others, |_rng, _tick, values| values.iter().sum::<f64>());
465        let mut rng = SmallRng::seed_from_u64(2);
466        let current_tick = SimTime::from_ticks(0);
467        assert_eq!(
468            aggregate_sampler.sample(&mut rng, current_tick).raw_value(),
469            17.0
470        ); // 10.0 + 5.0 + 2.0
471    }
472
473    #[test]
474    fn test_combinator_ext_clamp() {
475        let sampler = MockSampler { value: 10.0 };
476        let mut clamp_sampler = sampler.clamp(5.0, 8.0);
477        let mut rng = SmallRng::seed_from_u64(2);
478        let current_tick = SimTime::from_ticks(0);
479        // Clamped to max
480        assert_eq!(
481            clamp_sampler.sample(&mut rng, current_tick).raw_value(),
482            8.0
483        );
484
485        let sampler_low = MockSampler { value: 3.0 };
486        let mut clamp_sampler_low = sampler_low.clamp(5.0, 8.0);
487        // Clamped to min
488        assert_eq!(
489            clamp_sampler_low.sample(&mut rng, current_tick).raw_value(),
490            5.0
491        );
492
493        let sampler_in_range = MockSampler { value: 6.0 };
494        let mut clamp_sampler_in_range = sampler_in_range.clamp(5.0, 8.0);
495        // In range
496        assert_eq!(
497            clamp_sampler_in_range
498                .sample(&mut rng, current_tick)
499                .raw_value(),
500            6.0
501        );
502    }
503
504    #[test]
505    fn test_combinator_ext_ensure_non_negative() {
506        struct NegativeSampler;
507        impl DurationSampler for NegativeSampler {
508            fn sample(&mut self, _rng: &mut dyn Rng, _current_tick: SimTime) -> PendingDuration {
509                PendingDuration::new(-5.0)
510            }
511        }
512
513        let sampler = NegativeSampler;
514        let mut ensure_sampler = sampler.ensure_non_negative(3, |_rng, _tick| Duration::ticks(10));
515        let mut rng = SmallRng::seed_from_u64(2);
516        let current_tick = SimTime::from_ticks(0);
517        // Should fall back to the provided duration
518        assert_eq!(
519            ensure_sampler.sample(&mut rng, current_tick).raw_value(),
520            10.0
521        );
522
523        struct PositiveSampler;
524        impl DurationSampler for PositiveSampler {
525            fn sample(&mut self, _rng: &mut dyn Rng, _current_tick: SimTime) -> PendingDuration {
526                PendingDuration::new(5.0)
527            }
528        }
529        let sampler_pos = PositiveSampler;
530        let mut ensure_sampler_pos =
531            sampler_pos.ensure_non_negative(3, |_rng, _tick| Duration::ticks(10));
532        assert_eq!(
533            ensure_sampler_pos
534                .sample(&mut rng, current_tick)
535                .raw_value(),
536            5.0
537        );
538    }
539}