des_sim/modeling/sampler/combinator/
aggregate.rs1use crate::modeling::sampler::{DurationSampler, PendingDuration};
8use crate::primitive::time::SimTime;
9use rand::Rng;
10
11pub struct AggregateBuilder {
13 samplers: Vec<Box<dyn DurationSampler>>,
14}
15
16impl AggregateBuilder {
17 pub fn from_sampler(sampler: impl DurationSampler + 'static) -> Self {
19 AggregateBuilder {
20 samplers: vec![Box::new(sampler)],
21 }
22 }
23
24 pub fn add_sampler(mut self, sampler: Box<dyn DurationSampler>) -> Self {
26 self.samplers.push(sampler);
27 self
28 }
29
30 pub fn build<F>(self, f: F) -> AggregateSampler<F>
32 where
33 F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
34 {
35 AggregateSampler {
36 samplers: self.samplers,
37 f,
38 }
39 }
40}
41
42pub struct AggregateSampler<F>
44where
45 F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
46{
47 samplers: Vec<Box<dyn DurationSampler>>,
48 f: F,
49}
50
51impl<F> DurationSampler for AggregateSampler<F>
52where
53 F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
54{
55 fn sample(&mut self, rng: &mut dyn Rng, current_tick: SimTime) -> PendingDuration {
56 let mut sampled_list = Vec::with_capacity(self.samplers.len());
57 for sampler in &mut self.samplers {
58 let sampled = sampler.sample(rng, current_tick);
59 sampled_list.push(sampled.raw_value());
60 }
61
62 PendingDuration::new((self.f)(rng, current_tick, sampled_list))
63 }
64}
65
66impl<F> AggregateSampler<F>
67where
68 F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
69{
70 pub fn new(samplers: impl IntoIterator<Item = Box<dyn DurationSampler>>, f: F) -> Self {
72 let samplers: Vec<_> = samplers.into_iter().collect();
73
74 assert!(
75 !samplers.is_empty(),
76 "AggregateSampler requires at least one sampler."
77 );
78
79 AggregateSampler { samplers, f }
80 }
81}
82
83#[cfg(test)]
84mod tests {
85 use super::*;
86 use crate::modeling::sampler::CombinatorExt;
87 use crate::modeling::sampler::instance::ConstantSampler;
88 use rand::SeedableRng;
89 use rand::rngs::SmallRng;
90
91 #[test]
92 fn test_aggregate_sampler_sum() {
93 let mut rng = SmallRng::seed_from_u64(2);
94 let s1 = ConstantSampler::new(10.0);
95 let s2 = ConstantSampler::new(20.0);
96 let s3 = ConstantSampler::new(30.0);
97
98 let mut sampler =
99 AggregateSampler::new(vec![s1.boxed(), s2.boxed(), s3.boxed()], |_, _, samples| {
100 samples.iter().sum()
101 });
102
103 let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
104 assert_eq!(sample.raw_value(), 60.0);
105 }
106
107 #[test]
108 fn test_aggregate_sampler_builder() {
109 let mut rng = SmallRng::seed_from_u64(2);
110 let s1 = ConstantSampler::new(5.0);
111 let s2 = ConstantSampler::new(15.0);
112
113 let mut sampler = AggregateBuilder::from_sampler(s1)
114 .add_sampler(s2.boxed())
115 .build(|_, _, samples| samples[0] * samples[1]);
116
117 let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
118 assert_eq!(sample.raw_value(), 75.0); }
120
121 #[test]
122 #[should_panic]
123 fn test_aggregate_sampler_empty_samplers() {
124 let _sampler = AggregateSampler::new(vec![], |_, _, _| 0.0);
125 }
126}