Fix Geometric::new for small p>0

CHANGELOG.md

@@ -16,6 +16,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `Dirichlet` no longer uses `const` generics, which means that its size is not required at compile time. Essentially a revert of rand#1292. (#15)
 - Add `Dirichlet::new_with_size` constructor (#15)
 
+### Fixes
+- Fix `Geometric::new` for small `p > 0` where `1 - p` rounds to 1 (#36)
+
 ## [0.5.2]
 
 ### API Changes

Cargo.toml

@@ -39,9 +39,9 @@ serde = { version = "1.0.103", features = ["derive"], optional = true }
 serde_with = { version = "3", optional = true }
 
 [dev-dependencies]
-rand_pcg = { version = "0.9.0" }
+rand_pcg = { version = "0.10.0-rc.1" }
 # For inline examples
-rand = { version = "0.10.0-rc.0", features = ["small_rng"] }
+rand = { version = "0.10.0-rc.5", features = ["small_rng"] }
 # Histogram implementation for testing uniformity
 average = { version = "0.16", features = [ "std" ] }
 # Special functions for testing distributions

benches/Cargo.toml

@@ -7,8 +7,8 @@ publish = false
 [dependencies]
 
 [dev-dependencies]
-rand = { version = "0.10.0-rc.0", features = ["small_rng", "nightly"] }
-rand_pcg = "0.9.0"
+rand = { version = "0.10.0-rc.5", features = ["small_rng", "nightly"] }
+rand_pcg = "0.10.0-rc.1"
 rand_distr = { path = ".." }
 criterion = "0.5"
 criterion-cycles-per-byte = "0.6"

src/geometric.rs

@@ -67,18 +67,25 @@ impl fmt::Display for Error {
 impl std::error::Error for Error {}
 
 impl Geometric {
-    /// Construct a new `Geometric` with the given shape parameter `p`
-    /// (probability of success on each trial).
+    /// Construct a new `Geometric` distribution
+    ///
+    /// The shape parameter `p` is the probability of success on each trial.
+    ///
+    /// ### Edge cases
+    ///
+    /// If `p == 0.0` or `1.0 - p` rounds to `1.0` then sampling returns
+    /// `u64::MAX`.
     pub fn new(p: f64) -> Result<Self, Error> {
+        let mut pi = 1.0 - p;
         if !p.is_finite() || !(0.0..=1.0).contains(&p) {
             Err(Error::InvalidProbability)
-        } else if p == 0.0 || p >= 2.0 / 3.0 {
-            Ok(Geometric { p, pi: p, k: 0 })
+        } else if pi == 1.0 || p >= 2.0 / 3.0 {
+            Ok(Geometric { p, pi, k: 0 })
         } else {
             let (pi, k) = {
                 // choose smallest k such that pi = (1 - p)^(2^k) <= 0.5
                 let mut k = 1;
-                let mut pi = (1.0 - p).powi(2);
+                pi = pi * pi;
                 while pi > 0.5 {
                     k += 1;
                     pi = pi * pi;
@@ -106,7 +113,7 @@ impl Distribution<u64> for Geometric {
             return failures;
         }
 
-        if self.p == 0.0 {
+        if self.pi == 1.0 {
             return u64::MAX;
         }
 
@@ -264,4 +271,18 @@ mod test {
     fn geometric_distributions_can_be_compared() {
         assert_eq!(Geometric::new(1.0), Geometric::new(1.0));
     }
+
+    #[test]
+    fn small_p() {
+        let a = f64::EPSILON / 2.0;
+        assert!(1.0 - a < 1.0); // largest repr. value < 1
+        assert!(Geometric::new(a).is_ok());
+
+        let b = f64::EPSILON / 4.0;
+        assert!(b > 0.0);
+        assert!(1.0 - b == 1.0); // rounds to 1
+        let d = Geometric::new(b).unwrap();
+        let mut rng = crate::test::VoidRng;
+        assert_eq!(d.sample(&mut rng), u64::MAX);
+    }
 }

src/lib.rs

@@ -155,6 +155,20 @@ mod test {
     // NOTE: Some distributions have tests checking only that samples can be
     // generated. This is redundant with vector and correctness tests.
 
+    /// An RNG which panics on first use
+    pub struct VoidRng;
+    impl rand::RngCore for VoidRng {
+        fn next_u32(&mut self) -> u32 {
+            panic!("usage of VoidRng")
+        }
+        fn next_u64(&mut self) -> u64 {
+            panic!("usage of VoidRng")
+        }
+        fn fill_bytes(&mut self, _: &mut [u8]) {
+            panic!("usage of VoidRng")
+        }
+    }
+
     /// Construct a deterministic RNG with the given seed
     pub fn rng(seed: u64) -> impl rand::RngCore {
         // For tests, we want a statistically good, fast, reproducible RNG.