Merge pull request #41 from mstoeckl/zipf-edge-case

Avoid hangs and debug asserts on invalid parameters for Zipf

Author: dhardy

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Use `direct-minimal-versions` (#38)
 - Fix panic in `FisherF::new` on almost zero parameters (#39)
 - Fix panic in `NormalInverseGaussian::new` with very large `alpha`; this is a Value-breaking change (#40)
+- Fix hang and debug assertion in `Zipf::new` on invalid parameters (#41)
 - Fix panic in `Binomial::sample` with `n ≥ 2^63`; this is a Value-breaking change (#43)
 - Error instead of producing `-inf` output for `Exp` when `lambda` is `-0.0` (#44)
 

src/zipf.rs

@@ -66,17 +66,20 @@ where
 /// Error type returned from [`Zipf::new`].
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Error {
-    /// `s < 0` or `nan`.
+    /// `s < 0` or `s` is `nan`
     STooSmall,
-    /// `n < 1`.
+    /// `n < 1` or `n` is `nan`
     NTooSmall,
+    /// `n = inf` and `s <= 1`
+    IllDefined,
 }
 
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(match self {
             Error::STooSmall => "s < 0 or is NaN in Zipf distribution",
-            Error::NTooSmall => "n < 1 in Zipf distribution",
+            Error::NTooSmall => "n < 1 or is NaN in Zipf distribution",
+            Error::IllDefined => "n = inf and s <= 1 in Zipf distribution",
         })
     }
 }
@@ -100,17 +103,22 @@ where
         if !(s >= F::zero()) {
             return Err(Error::STooSmall);
         }
-        if n < F::one() {
+        if !(n >= F::one()) {
             return Err(Error::NTooSmall);
         }
+        if n.is_infinite() && s <= F::one() {
+            return Err(Error::IllDefined);
+        }
         let q = if s != F::one() {
             // Make sure to calculate the division only once.
             F::one() / (F::one() - s)
         } else {
             // This value is never used.
             F::zero()
         };
-        let t = if s != F::one() {
+        let t = if s == F::infinity() {
+            F::one()
+        } else if s != F::one() {
             (n.powf(F::one() - s) - s) * q
         } else {
             F::one() + n.ln()
@@ -220,6 +228,16 @@ mod tests {
         // TODO: verify that this is a uniform distribution
     }
 
+    #[test]
+    fn zipf_sample_s_inf() {
+        let d = Zipf::new(10., f64::infinity()).unwrap();
+        let mut rng = crate::test::rng(2);
+        for _ in 0..1000 {
+            let r = d.sample(&mut rng);
+            assert!(r == 1.);
+        }
+    }
+
     #[test]
     fn zipf_sample_large_n() {
         let d = Zipf::new(f64::MAX, 1.5).unwrap();