From 3929d06d668845aeb7cade6609d64f3e83c7fca1 Mon Sep 17 00:00:00 2001
From: stefan-aws <stefanze@amazon.com>
Date: Fri, 3 Nov 2023 19:22:00 +0000
Subject: [PATCH 1/2] skeleton

---
 src/interop/java/DRandomUniform.java | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/interop/java/DRandomUniform.java b/src/interop/java/DRandomUniform.java
index 12d310f5..0a3c444e 100644
--- a/src/interop/java/DRandomUniform.java
+++ b/src/interop/java/DRandomUniform.java
@@ -5,9 +5,7 @@
 import java.lang.Thread;
 
 public final class DRandomUniform {
-
   private static final ThreadLocal<SecureRandom> RNG = ThreadLocal.withInitial(DRandomUniform::createSecureRandom);
-  
   private DRandomUniform() {} // Prevent instantiation
 
   private static final SecureRandom createSecureRandom() {
@@ -17,9 +15,23 @@ private static final SecureRandom createSecureRandom() {
     return rng;
   }
 
-  public static BigInteger Uniform(BigInteger n) {
-    // `n.intValueExact` will throw an `ArithmeticException` if `n` does not fit in an `int`.
-    // see https://docs.oracle.com/javase/8/docs/api/java/math/BigInteger.html#intValueExact--
-    return BigInteger.valueOf(RNG.get().nextInt(n.intValueExact()));
+  /**
+   * Sample a uniform value using rejection sampling between [0, n).
+   * 
+   * @param n an integer (must be >= 1)
+   * @return a uniform value between 0 and n-1
+   * @throws IllegalArgumentException if `n` is less than 1
+   */
+  public static BigInteger Uniform(final BigInteger n) {
+    if (n.compareTo(BigInteger.ONE) < 0) {
+      throw new IllegalArgumentException("n must be positive");
+    }
+
+    BigInteger sampleValue;
+    do {
+      sampleValue = new BigInteger(n.bitLength(), RNG);
+    } while (sampleValue.compareTo(n) >= 0);
+
+    return sampleValue;
   }
-}
+}
\ No newline at end of file

From f60fa57e16a2691378d5da7834c4086c05e392b2 Mon Sep 17 00:00:00 2001
From: stefan-aws <stefanze@amazon.com>
Date: Fri, 3 Nov 2023 19:30:28 +0000
Subject: [PATCH 2/2] fix

---
 src/interop/java/DRandomUniform.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/interop/java/DRandomUniform.java b/src/interop/java/DRandomUniform.java
index 0a3c444e..b962022b 100644
--- a/src/interop/java/DRandomUniform.java
+++ b/src/interop/java/DRandomUniform.java
@@ -29,7 +29,7 @@ public static BigInteger Uniform(final BigInteger n) {
 
     BigInteger sampleValue;
     do {
-      sampleValue = new BigInteger(n.bitLength(), RNG);
+      sampleValue = new BigInteger(n.bitLength(), RNG.get());
     } while (sampleValue.compareTo(n) >= 0);
 
     return sampleValue;