Allow supplying a prng (#112)

* Allow supplying a prng

* Optimize random-bits for "small" numbers of bits.

---------

Co-authored-by: Pavel Panchekha <[email protected]>

Author: samth

math-doc/math/scribblings/math-base.scrbl

@@ -82,23 +82,29 @@ Like @racket[(apply + xs)], but incurs rounding error only once when adding inex
 
 @section{Random Number Generation}
 
-@defproc[(random-natural [k Integer]) Natural]{
+@defproc[(random-natural [k Integer] [p Pseudo-Random-Generator (current-pseudo-random-generator)]) Natural]{
 Returns a random natural number less than @racket[k], which must be positive.
 Use @racket[(random-natural k)] instead of @racket[(random k)] when @racket[k]
 could be larger than @racket[4294967087].
+
+Uses @racket[p] for the underlying generation of randomness. 
 }
 
-@defproc[(random-integer [a Integer] [b Integer]) Integer]{
+@defproc[(random-integer [a Integer] [b Integer] [p Pseudo-Random-Generator (current-pseudo-random-generator)]) Integer]{
 Returns a random integer @racket[n] such that @racket[(a . <= . n)] and @racket[(n . < . b)].
+
+Uses @racket[p] for the underlying generation of randomness. 
 }
 
-@defproc[(random-bits [num Integer]) Natural]{
+@defproc[(random-bits [num Integer] [p Pseudo-Random-Generator (current-pseudo-random-generator)]) Natural]{
 Returns a random natural smaller than @racket[(expt 2 num)]; @racket[num] must be non-negative.
 For powers of two, this is faster than using @racket[random-natural], which
 is implemented in terms of @racket[random-bits], using biased rejection sampling.
 
 As an example of use, the significands of the numbers returned by @racket[bfrandom]
 are chosen by @racket[(random-bits (bf-precision))].
+
+Uses @racket[p] for the underlying generation of randomness. 
 }
 
 

math-lib/math/private/base/base-random.rkt

@@ -5,39 +5,49 @@
 (provide random-bits random-natural random-integer)
 
 ;; Random bits are taken in blocks of this size:
-(define block-bits 29)
-(define block-size (arithmetic-shift 1 block-bits))
+(define block-bits : Nonnegative-Fixnum 29)
+(define block-size : Natural (arithmetic-shift 1 block-bits))
 
-(: random-bits (Integer -> Natural))
-(define (random-bits bits)
+(: random-bits (case-> (Integer -> Natural)
+                       (Integer Pseudo-Random-Generator -> Natural)))
+(define (random-bits bits [prng (current-pseudo-random-generator)])
   (cond [(bits . < . 0)  (raise-argument-error 'random-bits "Non-Negative-Integer" bits)]
         [(bits . = . 0) 0]
+        [(not (fixnum? bits))  (raise-argument-error 'random-bits "reasonably sized integer" bits)]
+        [(bits . <= . block-bits)  (random (ann (fxlshift 1 bits) Nonnegative-Fixnum) prng)]
+        [(bits . <= . (fx* 2 block-bits))
+	 ;; this case is not always in the fixnum range on 32-bit platforms
+         (define rem-bits (fx- bits block-bits))
+         (bitwise-ior (arithmetic-shift (random block-size prng) rem-bits)
+		      (random (ann (fxlshift 1 rem-bits) Nonnegative-Fixnum) prng))]
         [else
          (define max-blocks (assert (quotient bits block-bits) index?))
          (define rem-bits (remainder bits block-bits))
          (let: loop : Natural ([blocks : Nonnegative-Fixnum  0]
-                               [r : Natural  (random (fxlshift 1 rem-bits))])
+                               [r : Natural  (random (fxlshift 1 rem-bits) prng)])
            (cond [(blocks . fx< . max-blocks)
                   (loop (fx+ blocks 1)
                         (bitwise-ior (arithmetic-shift r block-bits)
-                                     (random block-size)))]
+                                     (random block-size prng)))]
                  [else  r]))]))
 
 (define random-max 4294967087)
 
-(: random-natural (Integer -> Natural))
+(: random-natural (case-> (Integer -> Natural)
+                          (Integer Pseudo-Random-Generator -> Natural)))
 ;; Returns a random integer in the interval [0..n)
-(define (random-natural n)
+(define (random-natural n [prng (current-pseudo-random-generator)])
   (cond
     [(n . <= . 0)  (raise-argument-error 'random-natural "Positive-Integer" n)]
-    [(n . <= . random-max)  (random n)]
+    [(n . <= . random-max)  (random n prng)]
     [else
      (define bits (integer-length (- n 1)))
      (let loop ()
-       (define r (random-bits bits))
+       (define r (random-bits bits prng))
        (if (r . >= . n) (loop) r))]))
 
-(: random-integer (Integer Integer -> Integer))
-(define (random-integer a b)
+(: random-integer (case-> (Integer Integer -> Integer)
+                          (Integer Integer Pseudo-Random-Generator -> Integer)))
+(define (random-integer a b [prng (current-pseudo-random-generator)])
   (let ([a  (min a b)] [b  (max a b)])
-    (+ a (random-natural (- b a)))))
+    (+ a (random-natural (- b a) prng))))

math-test/math/tests/random-tests.rkt

@@ -2,6 +2,8 @@
 (require math/base)
 (require typed/rackunit)
 
+(define g (current-pseudo-random-generator))
+
 (check-true
   (for/or ([n (in-range 100)])
     (> (random-natural (expt 2 64)) 4294967087)))
@@ -13,3 +15,16 @@
 (for ([n (in-range 100)])
   (define x (random-bits n))
   (check-true (and (<= 0 x) (< x (expt 2 n)))))
+
+
+(check-true
+  (for/or ([n (in-range 100)])
+    (> (random-natural (expt 2 64) g) 4294967087)))
+
+(for ([n (in-range 100)])
+  (define x (random-integer 1 100 g))
+  (check-true (and (<= 1 x) (< x 100))))
+
+(for ([n (in-range 100)])
+  (define x (random-bits n g))
+  (check-true (and (<= 0 x) (< x (expt 2 n)))))