diff --git a/src/distributions/uniform.rs b/src/distributions/uniform.rs index 9a0acd018ac..87c4322c52a 100644 --- a/src/distributions/uniform.rs +++ b/src/distributions/uniform.rs @@ -693,7 +693,6 @@ uniform_float_impl! { f64x8, u64x8, f64, u64, 64 - 52 } #[cfg(feature = "std")] #[derive(Clone, Copy, Debug)] pub struct UniformDuration { - offset: Duration, mode: UniformDurationMode, } @@ -701,10 +700,17 @@ pub struct UniformDuration { #[derive(Debug, Copy, Clone)] enum UniformDurationMode { Small { + secs: u64, + nanos: Uniform, + }, + Medium { nanos: Uniform, }, Large { - size: Duration, + min_secs: u64, + min_nanos: u32, + max_secs: u64, + max_nanos: u32, secs: Uniform, } } @@ -737,52 +743,65 @@ impl UniformSampler for UniformDuration { let low = *low_b.borrow(); let high = *high_b.borrow(); assert!(low <= high, "Uniform::new_inclusive called with `low > high`"); - let size = high - low; - let nanos = size - .as_secs() - .checked_mul(1_000_000_000) - .and_then(|n| n.checked_add(size.subsec_nanos() as u64)); - - let mode = match nanos { - Some(nanos) => { - UniformDurationMode::Small { - nanos: Uniform::new_inclusive(0, nanos), - } + + let low_s = low.as_secs(); + let low_n = low.subsec_nanos(); + let high_s = high.as_secs(); + let high_n = high.subsec_nanos(); + + let mode = if low_s == high_s { + UniformDurationMode::Small { + secs: low_s, + nanos: Uniform::new_inclusive(low_n, high_n), } - None => { + } else { + let max = high_s + .checked_mul(1_000_000_000) + .and_then(|n| n.checked_add(high_n as u64)); + + if let Some(higher_bound) = max { + let lower_bound = low_s * 1_000_000_000 + low_n as u64; + UniformDurationMode::Medium { + nanos: Uniform::new_inclusive(lower_bound, higher_bound), + } + } else { UniformDurationMode::Large { - size: size, - secs: Uniform::new_inclusive(0, size.as_secs()), + min_secs: low_s, + min_nanos: low_n, + max_secs: high_s, + max_nanos: high_n, + secs: Uniform::new_inclusive(low_s, high_s), } } }; - UniformDuration { - mode, - offset: low, + mode } } #[inline] fn sample(&self, rng: &mut R) -> Duration { - let d = match self.mode { - UniformDurationMode::Small { nanos } => { + match self.mode { + UniformDurationMode::Small { secs, nanos } => { + let n = nanos.sample(rng); + Duration::new(secs, n) + } + UniformDurationMode::Medium { nanos } => { let nanos = nanos.sample(rng); Duration::new(nanos / 1_000_000_000, (nanos % 1_000_000_000) as u32) } - UniformDurationMode::Large { size, secs } => { + UniformDurationMode::Large { min_secs, min_nanos, max_secs, max_nanos, secs } => { // constant folding means this is at least as fast as `gen_range` let nano_range = Uniform::new(0, 1_000_000_000); loop { - let d = Duration::new(secs.sample(rng), nano_range.sample(rng)); - if d <= size { - break d; + let s = secs.sample(rng); + let n = nano_range.sample(rng); + if !((s == max_secs && n > max_nanos) || (s == min_secs && n < min_nanos)) { + break Duration::new(s, n); } } } - }; - - self.offset + d + } } }