diff --git a/rand_distr/src/binomial.rs b/rand_distr/src/binomial.rs index 0f49806aba..4cd21fd817 100644 --- a/rand_distr/src/binomial.rs +++ b/rand_distr/src/binomial.rs @@ -110,20 +110,34 @@ impl Distribution for Binomial { // Threshold for preferring the BINV algorithm. The paper suggests 10, // Ranlib uses 30, and GSL uses 14. const BINV_THRESHOLD: f64 = 10.; + + // Same value as in GSL. + // It is possible for BINV to get stuck, so we break if x > BINV_MAX_X and try again. + // It would be safer to set BINV_MAX_X to self.n, but it is extremely unlikely to be relevant. + // When n*p < 10, so is n*p*q which is the variance, so a result > 110 would be 100 / sqrt(10) = 31 standard deviations away. + const BINV_MAX_X : u64 = 110; if (self.n as f64) * p < BINV_THRESHOLD && self.n <= (core::i32::MAX as u64) { // Use the BINV algorithm. let s = p / q; let a = ((self.n + 1) as f64) * s; - let mut r = q.powi(self.n as i32); - let mut u: f64 = rng.gen(); - let mut x = 0; - while u > r as f64 { - u -= r; - x += 1; - r *= a / (x as f64) - s; + + result = 'outer: loop { + let mut r = q.powi(self.n as i32); + let mut u: f64 = rng.gen(); + let mut x = 0; + + while u > r as f64 { + u -= r; + x += 1; + if x > BINV_MAX_X { + continue 'outer; + } + r *= a / (x as f64) - s; + } + break x; + } - result = x; } else { // Use the BTPE algorithm. @@ -352,4 +366,15 @@ mod test { fn binomial_distributions_can_be_compared() { assert_eq!(Binomial::new(1, 1.0), Binomial::new(1, 1.0)); } + + #[test] + fn binomial_avoid_infinite_loop() { + let dist = Binomial::new(16000000, 3.1444753148558566e-10).unwrap(); + let mut sum: u64 = 0; + let mut rng = crate::test::rng(742); + for _ in 0..100_000 { + sum = sum.wrapping_add(dist.sample(&mut rng)); + } + assert_ne!(sum, 0); + } }