;;; Guile-QuickCheck
;;; Copyright 2020 Timothy Sample <samplet@ngyro.com>
;;;
;;; This file is part of Guile-QuickCheck.
;;;
;;; Guile-QuickCheck is free software: you can redistribute it and/or
;;; modify it under the terms of the GNU General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; Guile-QuickCheck is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;; GNU General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with Guile-QuickCheck.  If not, see
;;; <https://www.gnu.org/licenses/>.

(define-module (quickcheck rng)
  #:use-module (ice-9 match)
  #:use-module (ice-9 receive)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-42)
  #:use-module (srfi srfi-43)
  #:export (make-rng-state
            rng-substream
            rng-next
            rng-integer))

;;; Commentary:
;;;
;;; This module is an implementation of the MRG32k3a random number
;;; generator (RNG) created by Pierre L'Ecuyer.  This is in fact the RNG
;;; proposed for SRFI-27, but Guile's implementation does not use it.
;;; Two obvious options for RNGs would be to either use Guile's or to
;;; follow QuickCheck and use Haskell's.  The default RNG provided by
;;; Guile has a relatively short period (not quite 2^62), and thus
;;; cannot provide many independent substreams of random numbers
;;; (something necessary for generating random functions).  Haskell
;;; provides an older RNG created by L'Ecuyer.  Its period is slightly
;;; shorter than that of Guile's.  Nevertheless, QuickCheck makes
;;; extensive use of the "split" function to create multiple substreams.
;;; The split function fiddles with the state of the generator in order
;;; to create two states, and includes the comment "no statistical
;;; foundation for this!"  The opportunity to pay omage to SRFI-27 while
;;; (possibly) improving the statistical basis of QuickCheck is too much
;;; to pass up.
;;;
;;; Note that this could also be built on top of the Mersenne Twister
;;; RNG, but it is slightly harder to implement (particularly
;;; "advance-rng-state").
;;;
;;; Code:


;;; Linear algebra

(define (matrix-ref mat i j)
  "Get the value at the @var{i}th row and @var{j}th column of matrix
@var{mat}."
  (vector-ref mat (+ (* 3 i) j)))

(define (apply-matrix-mod m mat v)
  "Apply matrix @var{mat} to vector @var{v} modulo @var{m}."
  (vector-of-length-ec 3 (: i 3)
                       (modulo (apply + (list-ec (: j 3)
                                                 (* (matrix-ref mat i j)
                                                    (vector-ref v j))))
                               m)))

(define (matrix-product-mod m mat1 mat2)
  "Multiply matricies @var{mat1} and @var{mat2} modulo @var{m}."
  (vector-of-length-ec 9 (: i 3) (: j 3)
                       (modulo (apply + (list-ec (: k 3)
                                                 (* (matrix-ref mat1 i k)
                                                    (matrix-ref mat2 k j))))
                               m)))

(define (matrix-power-mod m mat n)
  "Raise matrix @var{mat} to the power @var{n} modulo @var{m}."
  (let loop ((mat mat) (mat* #(1 0 0 0 1 0 0 0 1)) (n n))
    (cond
     ((<= n 0) mat*)
     ((odd? n)
      (loop mat (matrix-product-mod m mat mat*) (1- n)))
     ((even? n)
      (loop (matrix-product-mod m mat mat) mat* (quotient n 2))))))


;;; MRG32k3a random number generator

(define-record-type <rng-state>
  (%make-rng-state start s1 s2)
  rng-state?
  (start rng-state-start)
  (s1 rng-state-s1)
  (s2 rng-state-s2))

(define m1 4294967087)
(define m2 4294944443)

(define next-s1 #(0 1403580 4294156359 1 0 0 0 1 0))
(define next-s2 #(527612 0 4293573854 1 0 0 0 1 0))

(define (next-rng-state rngs)
  "Get the next RNG state after @var{rngs}."
  (match-let ((($ <rng-state> start s1 s2) rngs))
    (%make-rng-state start
                     (apply-matrix-mod m1 next-s1 s1)
                     (apply-matrix-mod m2 next-s2 s2))))

(define (advance-rng-state rngs k)
  "Get the @var{k}th RNG state after @var{rngs}."
  (match-let ((($ <rng-state> start s1 s2) rngs))
    (let ((advance-s1 (matrix-power-mod m1 next-s1 k))
          (advance-s2 (matrix-power-mod m2 next-s2 k)))
      (%make-rng-state start
                       (apply-matrix-mod m1 advance-s1 s1)
                       (apply-matrix-mod m2 advance-s2 s2)))))

(define (rng-value rngs)
  "Get the current value for @var{rngs}.  This will be a nonnegative
integer less than 4294967087."
  (match-let ((($ <rng-state> _ #(x1 _ _) #(x2 _ _)) rngs))
    (modulo (- x1 x2) m1)))

(define (start->rng-state start)
  "Convert @var{start} (a six element vector) into an RNG state."
  (%make-rng-state start
                   (vector-of-length-ec 3 (: i 3) (vector-ref start i))
                   (vector-of-length-ec 3 (: i 3 6) (vector-ref start i))))

(define (make-rng-state seed)
  "Make a new initial RNG state from the number @var{seed}."
  (let* ((rnds (seed->random-state seed))
         (s1 (vector-of-length-ec 3 (: i 3) (random m1 rnds)))
         (s2 (vector-of-length-ec 3 (: i 3) (random m2 rnds))))
    (%make-rng-state (vector-append s1 s2) s1 s2)))

(define (rng-substream rngs k)
  "Get the RNG state for the @var{k}th substream starting from the
initial state of @var{rngs}."
  (define step (* (expt 2 127) (modulo k (expt 2 64))))
  (advance-rng-state (start->rng-state (rng-state-start rngs)) step))

(define (rng-next rngs)
  "Get the current value and next state of @var{rngs} (returned as two
values)."
  (values (rng-value rngs) (next-rng-state rngs)))


;;; Random integers and reals

(define (rng-next* rngs size)
  "Get a value made up of @var{size} values from @var{rngs}, as well
as its next state.  If @code{rng-next} normally returns values less
than @var{m}, this will return values less than
@math{@var{m}@sup{@var{size}}}."
  (let loop ((rngs rngs) (size size) (scale 1) (acc 0))
    (if (<= size 0)
        (values acc rngs)
        (receive (val rngs*) (rng-next rngs)
          (loop rngs* (1- size) (* scale m1) (+ acc (* val scale)))))))

(define (%rng-integer rngs limit)
  "Get a nonnegative integer less than @var{limit} from @var{rngs}, as
well as its next state."
  (receive (size m1*)
      (let loop ((size 1) (m1* m1))
        (if (< limit m1*)
            (values size m1*)
            (loop (1+ size) (* m1* m1))))
    (let ((remainder (modulo m1* limit)))
      (let loop ((i 0) (rngs rngs))
        (receive (val rngs*) (rng-next* rngs size)
          (if (< val (- m1* remainder))
              (values (modulo val limit) rngs*)
              (loop (1+ i) rngs*)))))))

(define (rng-integer rngs min max)
  "Get an integer and the next state of @var{rngs}.  The integer will
be no smaller than @var{min} and no larger than @var{max}."
  (let ((size (1+ (- max min))))
    (receive (val rngs*) (%rng-integer rngs size)
      (values (+ val min) rngs*))))
