1 files changed, 12 insertions, 8 deletions
diff --git a/lib/jxl/base/random.h b/lib/jxl/base/random.h
index 663b88c..b27815b 100644
--- a/lib/jxl/base/random.h
+++ b/lib/jxl/base/random.h
@@ -14,6 +14,7 @@
 #include <string.h>
 
 #include <algorithm>
+#include <cmath>
 
 #include "lib/jxl/base/status.h"
 
@@ -69,15 +70,18 @@ struct Rng {
   bool Bernoulli(float p) { return UniformF(0, 1) < p; }
 
   // State for geometric distributions.
-  struct GeometricDistribution {
-    explicit GeometricDistribution(float p);
-
-   private:
-    float inv_log_1mp;
-    friend struct Rng;
-  };
+  // The stored value is inv_log_1mp
+  using GeometricDistribution = float;
+  static GeometricDistribution MakeGeometric(float p) {
+    return 1.0 / std::log(1 - p);
+  }
 
-  uint32_t Geometric(const GeometricDistribution& dist);
+  uint32_t Geometric(const GeometricDistribution& dist) {
+    float f = UniformF(0, 1);
+    float inv_log_1mp = dist;
+    float log = std::log(1 - f) * inv_log_1mp;
+    return static_cast<uint32_t>(log);
+  }
 
   template <typename T>
   void Shuffle(T* t, size_t n) {