;;; 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 arbitrary)
  #:use-module (ice-9 match)
  #:use-module (quickcheck generator)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-26)
  #:use-module (srfi srfi-43)
  #:export (arbitrary
            arbitrary?
            arbitrary-gen
            arbitrary-xform
            $const
            $bool
            $integer
            $natural
            $byte
            $char
            $list
            $pair
            $vector
            $tuple
            $record
            $string
            $symbol
            $bytevector
            $choose
            $procedure))

(define-record-type <arbitrary>
  (make-arbitrary gen xform)
  arbitrary?
  (gen arbitrary-gen)
  (xform arbitrary-xform))

(define-syntax arbitrary
  (syntax-rules (gen xform)
    ((arbitrary (gen g) (xform t))
     (make-arbitrary g t))))

(define ($const value)
  (arbitrary
   (gen (generator-return value))
   (xform (lambda (x gen) gen))))

(define $bool
  (arbitrary
   (gen (generator-let* ((n (choose-integer 0 1)))
          (generator-return (= n 1))))
   (xform (lambda (b gen)
            (generator-variant (if b 1 0) gen)))))

(define $integer
  (arbitrary
   (gen (sized-generator (lambda (size) (choose-integer (- size) size))))
   (xform (lambda (n gen)
            (generator-variant (if (>= n 0) (* 2 n) (1- (* -2 n))) gen)))))

(define $natural
  (arbitrary
   (gen (sized-generator (lambda (size) (choose-integer 0 size))))
   (xform generator-variant)))

(define $byte
  (arbitrary
   (gen choose-byte)
   (xform generator-variant)))

(define $char
  (arbitrary
   (gen (choose-char char-set:designated))
   (xform (lambda (chr gen)
            (generator-variant (char->integer chr) gen)))))

(define ($list elem)
  (match-let* ((($ <arbitrary> elem-gen elem-xform) elem))
    (arbitrary
     (gen (sized-generator (cut choose-list elem-gen <>)))
     (xform (lambda (xs gen)
              (fold (lambda (x gen)
                      (generator-variant 1 (elem-xform x gen)))
                    (generator-variant 0 gen)
                    xs))))))

(define ($pair kar kdr)
  (match-let* ((($ <arbitrary> kar-gen kar-xform) kar)
               (($ <arbitrary> kdr-gen kdr-xform) kdr))
    (arbitrary
     (gen (generator-lift cons kar-gen kdr-gen))
     (xform (lambda (p gen)
              (kar-xform (car p) (kdr-xform (cdr p) gen)))))))

(define ($vector elem)
  (match-let* ((($ <arbitrary> elem-gen elem-xform) elem))
    (arbitrary
     (gen (sized-generator (cut choose-vector elem-gen <>)))
     (xform (lambda (vec gen)
              (vector-fold (lambda (x gen) (elem-xform x gen))
                           (generator-variant (vector-length vec) gen)
                           vec))))))

(define ($tuple . elems)
  (arbitrary
   (gen (generator-fold-right cons '() (map arbitrary-gen elems)))
   (xform (lambda (xs gen)
            (fold (lambda (x xform gen)
                    (generator-variant 1 (xform x gen)))
                  (generator-variant 0 gen)
                  xs
                  (map arbitrary-xform elems))))))

(define (%$record constructor accessor+elems)
  (arbitrary
   (gen (generator-let* ((args (generator-fold-right
                                cons '() (map (compose arbitrary-gen cdr)
                                              accessor+elems))))
          (generator-return (apply constructor args))))
   (xform (lambda (rec gen)
            (fold (lambda (accessor+elem gen)
                    (match-let* (((accessor . elem) accessor+elem)
                                 (xform (arbitrary-xform elem)))
                      (generator-variant 1 (xform (accessor rec) gen))))
                  (generator-variant 0 gen)
                  accessor+elems)))))

(define-syntax-rule ($record constructor (accessor arb) ...)
  (%$record constructor (list (cons accessor arb) ...)))

(define ($string chr)
  (match-let* ((($ <arbitrary> chr-gen chr-xform) chr))
    (arbitrary
     (gen (sized-generator (cut choose-string chr-gen <>)))
     (xform (lambda (str gen)
              (string-fold (lambda (chr gen) (chr-xform chr gen))
                           (generator-variant (string-length str) gen)
                           str))))))

(define ($symbol chr)
  (match-let* ((($ <arbitrary> chr-gen chr-xform) chr))
    (arbitrary
     (gen (sized-generator (cut choose-symbol chr-gen <>)))
     (xform (lambda (sym gen)
              (let ((str (symbol->string sym)))
                (string-fold (lambda (chr gen) (chr-xform chr gen))
                             (generator-variant (string-length str) gen)
                             str)))))))

(define $bytevector
  (arbitrary
   (gen (sized-generator choose-bytevector))
   (xform (lambda (bv gen)
            (define len (bytevector-length bv))
            (let loop ((k 0) (acc (generator-variant len gen)))
              (if (>= k len)
                  acc
                  (loop (1+ k)
                        (generator-variant (bytevector-u8-ref bv k)
                                           acc))))))))

(define (%$choose pred+arbs)
  (arbitrary
   (gen (choose-one (map (compose arbitrary-gen cdr) pred+arbs)))
   (xform (lambda (x gen)
            (let loop ((pred+arbs pred+arbs) (k 0))
              (match pred+arbs
                (() (scm-error 'misc-error '%$choose
                               "No predicated matched the given value: ~A"
                               (list x) (list x)))
                (((pred . arb) . rest)
                 (if (pred x)
                     ((arbitrary-xform arb) x (generator-variant k gen))
                     (loop rest (1+ k))))))))))

(define-syntax-rule ($choose (pred arb) ...)
  (%$choose (list (cons pred arb) ...)))

(define ($procedure result arg)
  (arbitrary
   (gen (generator-promote
         (lambda (x)
           ((arbitrary-xform arg) x (arbitrary-gen result)))))
   (xform (lambda (proc gen)
            (generator-let* ((x (arbitrary-gen arg)))
              ((arbitrary-xform result) (proc x) gen))))))
