;;; Guile-QuickCheck
;;; Copyright 2020, 2021 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 generator)
  #:use-module (ice-9 receive)
  #:use-module (quickcheck rng)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-26)
  #:export (generator?
            generate
            generator-return
            generator-bind
            generator-lift
            generator-let*
            generator-promote
            generator-variant
            generator-fold-right
            generator-sample
            sized-generator
            resize-generator
            choose-integer
            choose-byte
            choose-list
            choose-vector
            choose-char
            choose-string
            choose-symbol
            choose-bytevector
            choose-one
            choose-one/weighted))

;;; Commentary:
;;;
;;; This module contains the definition of and various constructors
;;; for the generator monad.
;;;
;;; Code:

(define-record-type <generator>
  (make-generator proc)
  generator?
  (proc generator-proc))

(define (generate gen size rng)
  "Run the generator @var{gen} with the given @var{size} and
@var{rng}, a random number generator state."
  ((generator-proc gen) size rng))

(define (generator-return x)
  "Create a generator that returns @var{x}."
  (make-generator
   (lambda (size rng) (values x rng))))

(define (generator-bind gen proc)
  "Create a generator that applies the procedure @var{proc} to the output
of generator @var{gen}."
  (make-generator
   (lambda (size rng)
     (receive (x rng*) (generate gen size rng)
       (generate (proc x) size rng*)))))

(define-syntax generator-let*
  (syntax-rules ()
    "Create a generator that binds the generator values @var{gen} to the
variables @var{var} before evaluating the expressions @var{exp}."
    ((_ ((var gen) (var* gen*) (var** gen**) ...) exp exp* ...)
     (generator-bind gen (lambda (var)
                           (generator-let* ((var* gen*) (var** gen**) ...)
                             exp exp* ...))))
    ((_ ((var gen)) exp exp* ...)
     (generator-bind gen (lambda (var)
                           exp exp* ...)))))

(define (generator-fold-right kons knil gens)
  "Create a generator that folds over the results of each generator in
@var{gen} from right to left."
  (if (null? gens)
      (generator-return knil)
      (generator-let* ((kar (car gens))
                       (kdr (generator-fold-right kons knil (cdr gens))))
        (generator-return (kons kar kdr)))))

(define (generator-lift proc . gens)
  "Create a generator that returns the value of @var{proc} applied to
the results of the generators in @var{gens}."
  (generator-let* ((args (generator-fold-right cons '() gens)))
    (generator-return (apply proc args))))

(define (generator-promote proc)
  "Promote a procedure @var{proc} that returns a generator into a
generator that returns procedures."
  (make-generator
   (lambda (size rng)
     (lambda (x)
       (generate (proc x) size rng)))))

(define (generator-variant n gen)
  "Get the @var{n}th variant of the generator @var{gen}, where @var{n}
is a nonnegative number less than 2^64."
  (make-generator
   (lambda (size rng)
     (let ((rng* (rng-substream rng n)))
       (generate gen size rng*)))))

(define* (generator-sample gen #:key (count 16) (size 8))
  "Take a sample (a list of values) of the results generated by the
generator @var{gen}. The number of values is determined by
@var{count}, and the size of the values by @var{size}."
  (let loop ((rng (make-rng-state 4)) (acc '()) (count count))
    (if (<= count 0)
        acc
        (receive (x rng*) (generate gen size rng)
          (loop rng* (cons x acc) (1- count))))))

(define (sized-generator proc)
  "Create a generator using @var{proc}, a procedure that takes a size
and produces a generator. This is used to convert a generator whose
result size is determined at construction time into a generator whose
result size is determined at execution time."
  (make-generator
   (lambda (size rng)
     (generate (proc size) size rng))))

(define (resize-generator size gen)
  "Create a fixed-size generator by always providing @var{size} to the
generator @var{gen}."
  (make-generator
   (lambda (_ rng)
     (generate gen size rng))))

(define (choose-integer lower upper)
  "Create a generator that generates integers between @var{lower} and
@var{upper} (including both end points)."
  (make-generator
   (lambda (size rng)
     (rng-integer rng lower upper))))

(define choose-byte
  (choose-integer 0 255))

(define (choose-list elem-gen size)
  "Create a generator that generates lists of @var{size} elements taken
from @var{elem-gen}."
  (generator-fold-right cons '() (map (const elem-gen) (iota size))))

(define (choose-vector elem-gen size)
  "Create a generator that generates vectors of @var{size} elements
taken from @var{elem-gen}."
  (generator-lift list->vector (choose-list elem-gen size)))

(define (choose-char cs)
  "Create a generator that generates characters from the character set
@var{cs}."
  (let ((cs-len (char-set-size cs)))
    (when (= 0 cs-len)
      (error "Cannot choose from an empty character set."))
    ;; This will use about 4M of memory to represent the 'char-set:full'
    ;; character set.  That is a lot, but on the other hand, the speed
    ;; and simplicity here is pretty nice.
    (let ((bv (make-bytevector (* 4 cs-len))))
      (char-set-fold (lambda (c k)
                       (bytevector-u32-native-set! bv k (char->integer c))
                       (+ k 4))
                     0
                     cs)
      (generator-let* ((n (choose-integer 0 (1- cs-len))))
        (generator-return (integer->char
                           (bytevector-u32-native-ref bv (* n 4))))))))

(define (choose-string char-gen size)
  "Create a generator that generates strings of length @var{size} with
characters taken from @var{char-gen}."
  (generator-lift list->string (choose-list char-gen size)))

(define (choose-symbol char-gen size)
  "Create a generator that generates symbols of length @var{size} with
characters taken from @var{char-gen}."
  (generator-lift string->symbol (choose-string char-gen size)))

(define (choose-bytevector size)
  "Create a generator that generates bytevectors of length @var{size}."
  (generator-lift u8-list->bytevector (choose-list choose-byte size)))

(define (choose-one gens)
  "Create a generator that selects randomly from the generators in
@var{gens}."
  (when (null? gens)
    (error "Cannot choose from an empty list."))
  (generator-let* ((k (choose-integer 0 (1- (length gens)))))
    (list-ref gens k)))

(define (choose-one/weighted weight+gens)
  "Create a generator that selects randomly from the list of weighted
generators @var{weight+gens}. Each weight is a positive integer such
that a generator with weight 2 is twice as likely to be selected than
a generator with weight 1."
  (choose-one (fold (lambda (weight+gen acc)
                      (let ((weight (car weight+gen))
                            (gen (cdr weight+gen)))
                        (append-reverse (map (const gen) (iota weight)) acc)))
                    '()
                    weight+gens)))
