From 8b21c7daf3d64506cfab7c2147ff430fa48cf7ad Mon Sep 17 00:00:00 2001
From: Stefan Zetzsche <120379523+stefan-aws@users.noreply.github.com>
Date: Fri, 3 Nov 2023 19:35:35 +0000
Subject: [PATCH] Java integer fix (#110)

By submitting this pull request, I confirm that my contribution is made
under the terms of the [MIT
license](https://github.com/dafny-lang/dafny/blob/master/LICENSE.txt).
---
 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..b962022b 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.get());
+    } while (sampleValue.compareTo(n) >= 0);
+
+    return sampleValue;
   }
-}
+}
\ No newline at end of file