des_sim/modeling/sampler/combinator/
chain.rs

1//! The `chain` module provides the `ChainSampler`, a combinator that links two
2//! `DurationSampler` instances.
3//!
4//! It takes the output of both samplers and combines them using a user-defined
5//! closure, allowing for sequential or interdependent sampling logic.
6
7use crate::modeling::sampler::{DurationSampler, PendingDuration};
8use crate::primitive::time::SimTime;
9use rand::Rng;
10
11/// A sampler that chains two sub-samplers together, combining their results
12/// using a closure.
13pub struct ChainSampler<S1, F>
14where
15    S1: DurationSampler,
16    F: FnMut(&mut dyn Rng, SimTime, f64, f64) -> f64,
17{
18    sampler_1: S1,
19    sampler_2: Box<dyn DurationSampler>,
20    f: F,
21}
22
23impl<S1, F> DurationSampler for ChainSampler<S1, F>
24where
25    S1: DurationSampler,
26    F: FnMut(&mut dyn Rng, SimTime, f64, f64) -> f64,
27{
28    fn sample(&mut self, rng: &mut dyn Rng, current_tick: SimTime) -> PendingDuration {
29        let sampled_1 = self.sampler_1.sample(rng, current_tick);
30        let sampled_2 = self.sampler_2.sample(rng, current_tick);
31
32        // Pass the raw f64 values to the closure
33        PendingDuration::new((self.f)(
34            rng,
35            current_tick,
36            sampled_1.raw_value(),
37            sampled_2.raw_value(),
38        ))
39    }
40}
41
42impl<S1, F> ChainSampler<S1, F>
43where
44    S1: DurationSampler,
45    F: FnMut(&mut dyn Rng, SimTime, f64, f64) -> f64,
46{
47    /// Creates a new `ChainSampler`.
48    pub fn new(sampler_1: S1, sampler_2: Box<dyn DurationSampler>, f: F) -> Self {
49        ChainSampler {
50            sampler_1,
51            sampler_2,
52            f,
53        }
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use crate::modeling::sampler::CombinatorExt;
61    use crate::modeling::sampler::instance::ConstantSampler;
62    use rand::SeedableRng;
63    use rand::rngs::SmallRng;
64
65    #[test]
66    fn test_chain_sampler_add() {
67        let mut rng = SmallRng::seed_from_u64(2);
68        let s1 = ConstantSampler::new(10.0);
69        let s2 = ConstantSampler::new(20.0);
70
71        let mut sampler = ChainSampler::new(s1, s2.boxed(), |_, _, val1, val2| val1 + val2);
72
73        let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
74        assert_eq!(sample.raw_value(), 30.0);
75    }
76
77    #[test]
78    fn test_chain_sampler_multiply() {
79        let mut rng = SmallRng::seed_from_u64(2);
80        let s1 = ConstantSampler::new(5.0);
81        let s2 = ConstantSampler::new(4.0);
82
83        let mut sampler = ChainSampler::new(s1, s2.boxed(), |_, _, val1, val2| val1 * val2);
84
85        let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
86        assert_eq!(sample.raw_value(), 20.0);
87    }
88}
des_sim/modeling/sampler/combinator/chain.rs

des_sim/modeling/sampler/combinator/
chain.rs
