;; SPDX-FileCopyrightText: 2024 Artyom Bologov
;; SPDX-License-Identifier: MIT

;;; Permission is hereby granted, free of charge, to any person
;;; obtaining a copy of this software and associated documentation
;;; files (the "Software"), to deal in the Software without
;;; restriction, including without limitation the rights to use,
;;; copy, modify, merge, publish, distribute, sublicense, and/or
;;; sell copies of the Software, and to permit persons to whom the
;;; Software is furnished to do so, subject to the following
;;; conditions:
;;;
;;; The above copyright notice and this permission notice shall be
;;; included in all copies or substantial portions of the Software.
;;;
;;; THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
;;; EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
;;; OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
;;; NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
;;; HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
;;; WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
;;; FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
;;; OTHER DEALINGS IN THE SOFTWARE.

(import-for-syntax (chicken type))

(define-syntax assume
  (syntax-rules ()
    ((_ expr rest ...)
     (begin (assert expr rest ...) #t))))

(define-syntax check-arg
  (syntax-rules ()
    ((_ pred val caller)
     (assume (pred val) "argument should match the specification" '(pred val) val caller))
    ((_ pred val)
     (check-arg pred val 'check-arg))))

(define-syntax values-checked
  (syntax-rules (the
                  exact-integer? integer? boolean? char? complex?
                  fixnum? flonum?
                  eof? inexact? real?
                  list? null? number? pair?
                  input-port? output-port?
                  procedure? rational?
                  string? symbol? keyword? vector?
                  fixnum float boolean char cplxnum eof
                  list null number pair input-port output-port
                  procedure ratnum string symbol keyword vector)
    ((_ (fixnum?) value)
     (let ((v value))
       (check-arg fixnum? v 'values-checked)
       (the fixnum value)))
    ((_ (flonum?) value)
     (let ((v value))
       (check-arg float? v 'values-checked)
       (the float value)))
    ((_ (exact-integer?) value)
     (let ((v value))
       (check-arg integer? v 'values-checked)
       (the integer value)))
    ((_ (integer?) value)
     (let ((v value))
       (check-arg number? v 'values-checked)
       (the number value)))
    ((_ (boolean?) value)
     (let ((v value))
       (check-arg boolean? v 'values-checked)
       (the boolean value)))
    ((_ (char?) value)
     (let ((v value))
       (check-arg char? v 'values-checked)
       (the char value)))
    ((_ (complex?) value)
     (let ((v value))
       (check-arg complex? v 'values-checked)
       (the cplxnum value)))
    ((_ (eof-object?) value)
     (let ((v value))
       (check-arg eof-object? v 'values-checked)
       (the eof value)))
    ((_ (inexact?) value)
     (let ((v value))
       (check-arg inexact? v 'values-checked)
       (the float value)))
    ((_ (real?) value)
     (let ((v value))
       (check-arg real? v 'values-checked)
       (the number value)))
    ((_ (list?) value)
     (let ((v value))
       (check-arg list? v 'values-checked)
       (the list value)))
    ((_ (null?) value)
     (let ((v value))
       (check-arg null? v 'values-checked)
       (the null value)))
    ((_ (number?) value)
     (let ((v value))
       (check-arg number? v 'values-checked)
       (the number value)))
    ((_ (pair?) value)
     (let ((v value))
       (check-arg pair? v 'values-checked)
       (the pair value)))
    ((_ (input-port?) value)
     (let ((v value))
       (check-arg input-port? v 'values-checked)
       (the input-port value)))
    ((_ (output-port?) value)
     (let ((v value))
       (check-arg output-port? v 'values-checked)
       (the output-port value)))
    ((_ (procedure?) value)
     (let ((v value))
       (check-arg procedure? v 'values-checked)
       (the procedure value)))
    ((_ (rational?) value)
     (let ((v value))
       (check-arg rational? v 'values-checked)
       (the ratnum value)))
    ((_ (string?) value)
     (let ((v value))
       (check-arg string? v 'values-checked)
       (the string value)))
    ((_ (symbol?) value)
     (let ((v value))
       (check-arg symbol? v 'values-checked)
       (the symbol value)))
    ((_ (keyword?) value)
     (let ((v value))
       (check-arg keyword? v 'values-checked)
       (the keyword value)))
    ((_ (vector?) value)
     (let ((v value))
       (check-arg vector? v 'values-checked)
       (the vector value)))
    ((_ (predicate) value)
     (let ((v value))
       (check-arg predicate v 'values-checked)
       v))
    ((_ (predicate ...) value ...)
     (values (values-checked (predicate) value) ...))))

(define-syntax %check-case
  (syntax-rules (else)
    ((_ val (clause ...) (else body ...))
     (cond
       clause ...
       (else body ...)))
    ((_ val ((clause-check clause-body ...) ...))
     (cond
       (clause-check clause-body ...)
       ...
       (else (assume (or clause-check ...)
                     "at least one branch of check-case should be true"
                     'clause-check ...))))
    ((_ val (clause ...) (pred body ...) rest ...)
     (%check-case
       val
       (clause ... ((pred val) body ...))
       rest ...))))

(define-syntax check-case
  (syntax-rules ()
    ((_ value clause ...)
     (let ((v value))
       (%check-case v () clause ...)))))

(define-syntax %lambda-checked
  (syntax-rules ()
    ((_ name (body ...) args (checks ...))
     (lambda args
       checks ...
       body ...))
    ((_ name body (args ...) (checks ...) (arg pred) rest ...)
     (%lambda-checked
       name body
       (args ... arg) (checks ... (check-arg pred arg 'name)) rest ...))
    ((_ name body (args ...) (checks ...) arg rest ...)
     (%lambda-checked
       name body
       (args ... arg) (checks ...) rest ...))))

(define-syntax lambda-checked
  (syntax-rules ()
    ((_ (args ...) body ...)
     (%lambda-checked lambda-checked (body ...) () () args ...))
    ;; Case of arg->list lambda, no-op.
    ((_ arg body ...)
     (lambda arg body ...))))

(define-syntax %case-lambda-checked
  (syntax-rules ()
    ((_ (clauses-so-far ...)
        ()
        args-so-far (checks-so-far ...) (body ...))
     (case-lambda
       clauses-so-far ...
       (args-so-far
         checks-so-far ...
         body ...)))
    ((_ (clauses-so-far ...)
        ((() body-to-process ...) clauses-to-process ...)
        args-so-far (checks-so-far ...) (body ...))
     (%case-lambda-checked
       (clauses-so-far ... (args-so-far checks-so-far ... body ...))
       (clauses-to-process ...)
       () () (body-to-process ...)))
    ((_ (clauses-so-far ...)
        (((arg . args-to-process) body-to-process ...) clauses-to-process ...)
        args-so-far (checks-so-far ...) (body ...))
     (%case-lambda-checked
       (clauses-so-far ... (args-so-far checks-so-far ... body ...))
       (clauses-to-process ...)
       () () (body-to-process ...) arg . args-to-process))
    ((_ (clauses-so-far ...)
        ((arg-to-process body-to-process ...) clauses-to-process ...)
        args-so-far (checks-so-far ...) (body ...))
     (%case-lambda-checked
       (clauses-so-far ... (args-so-far checks-so-far ... body ...))
       (clauses-to-process ...)
       arg-to-process () (body-to-process ...)))
    ((_ (clauses-so-far ...) (clauses-to-process ...)
        (args-so-far ...) (checks-so-far ...) (body ...) (arg pred) . args)
     (%case-lambda-checked
       (clauses-so-far ...) (clauses-to-process ...)
       (args-so-far ... arg)
       (checks-so-far ... (check-arg pred arg 'case-lambda-checked))
       (body ...) . args))
    ((_ (clauses-so-far ...) (clauses-to-process ...)
        (args-so-far ...) (checks-so-far ...) (body ...) arg . args)
     (%case-lambda-checked
       (clauses-so-far ...) (clauses-to-process ...)
       (args-so-far ... arg) (checks-so-far ...) (body ...) . args))
    ((_ (clauses-so-far ...) (clauses-to-process ...)
        (args-so-far ...) (checks-so-far ...) (body ...) . arg)
     (%case-lambda-checked
       (clauses-so-far ...) (clauses-to-process ...)
       (args-so-far ... . arg) (checks-so-far ...) (body ...)))))

(define-syntax case-lambda-checked
  (syntax-rules ()
    ((_ (() first-body ...) rest-clauses ...)
     (%case-lambda-checked () (rest-clauses ...) () () (first-body ...)))
    ((_ ((first-arg . first-args) first-body ...) rest-clauses ...)
     (%case-lambda-checked () (rest-clauses ...) () () (first-body ...) first-arg . first-args))
    ((_ (args-var first-body ...) rest-clauses ...)
     (%case-lambda-checked () (rest-clauses ...) args-var () (first-body ...)))))

(define-syntax %declare-checked-var
  (syntax-rules (: ->
                   any? integer? exact-integer? boolean? char? complex?
                   fixnum? flonum?
                   eof? inexact? real?
                   list? null? number? pair?
                   input-port? output-port?
                   procedure? rational?
                   string? symbol? keyword? vector?

                   integer boolean char cplxnum eof fixnum float
                   number list null number pair input-port output-port
                   procedure ratnum string symbol keyword vector *)
    ((_ (predicate ...) value ...)
     (when #f #f))
    ((_ name any?)         (: name *))
    ((_ name fixnum?)      (: name fixnum))
    ((_ name flonum?)      (: name float))
    ((_ name integer?)     (: name number))
    ((_ name exact-integer?) (: name integer))
    ((_ name boolean?)     (: name boolean))
    ((_ name char?)        (: name char))
    ((_ name complex?)     (: name cplxnum))
    ((_ name eof?)         (: name eof))
    ((_ name inexact?)     (: name float))
    ((_ name real?)        (: name number))
    ((_ name list?)        (: name list))
    ((_ name null?)        (: name null))
    ((_ name number?)      (: name number))
    ((_ name pair?)        (: name pair))
    ((_ name input-port?)  (: name input-port))
    ((_ name output-port?) (: name output-port))
    ((_ name procedure?)   (: name procedure))
    ((_ name rational?)    (: name ratnum))
    ((_ name string?)      (: name string))
    ((_ name symbol?)      (: name symbol))
    ((_ name keyword?)     (: name keyword))
    ((_ name vector?)      (: name vector))
    ((_ name predicate)
     (when #f #f))))

(define-syntax %declare-checked-fn
  (syntax-rules (: ->
                   any? integer? boolean? char? complex?
                   fixnum? flonum?
                   eof? inexact? real?
                   list? null? number? pair?
                   input-port? output-port?
                   procedure? rational?
                   string? symbol? keyword? vector?

                   integer boolean char cplxnum eof fixnum float
                   number list null number pair input-port output-port
                   procedure ratnum string symbol keyword vector *)
    ((_ name () (type ...))
     (: name (type ... -> *)))
    ((_ name ((arg fixnum?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... fixnum)))
    ((_ name ((arg flonum?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... float)))
    ((_ name ((arg integer?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... number)))
    ((_ name ((arg boolean?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... boolean)))
    ((_ name ((arg char?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... char)))
    ((_ name ((arg complex?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... cplxnum)))
    ((_ name ((arg eof?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... eof)))
    ((_ name ((arg inexact?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... float)))
    ((_ name ((arg real?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... number)))
    ((_ name ((arg list?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... list)))
    ((_ name ((arg null?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... null)))
    ((_ name ((arg number?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... number)))
    ((_ name ((arg pair?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... pair)))
    ((_ name ((arg input-port?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... input-port)))
    ((_ name ((arg output-port?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... output-port)))
    ((_ name ((arg procedure?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... procedure)))
    ((_ name ((arg rational?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... ratnum)))
    ((_ name ((arg string?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... string)))
    ((_ name ((arg symbol?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... symbol)))
    ((_ name ((arg keyword?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... keyword)))
    ((_ name ((arg vector?) check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... vector)))
    ((_ name (arg check ...) (type ...))
     (%declare-checked-fn name (check ...) (type ... *)))))

(define-syntax define-checked
  (syntax-rules ()
    ;; Function
    ((_ (name arg ...) body ...)
     (begin
       (%declare-checked-fn name (arg ...) ())
       (define name (%lambda-checked name (body ...) () () arg ...))))
    ;; Variable
    ((_ name pred value)
     (begin
       (%declare-checked-var name pred)
       (define name (values-checked (pred) value))))))

(define-syntax %define-record-type-checked
  (syntax-rules ()
    ((_ type-name constructor predicate
        (fields ...) (field-wrappers ...))
     (begin
       (define-record-type
         type-name constructor predicate
         fields ...)
       field-wrappers ...))
    ((_ type-name constructor predicate
        (fields ...) (field-wrappers ...) (field pred accessor modifier)
        fields-to-process ...)
     (%define-record-type-checked
       type-name constructor predicate
       (fields ... (field internal-accessor internal-modifier))
       (field-wrappers
         ...
         (define-checked (accessor (record predicate))
           (internal-accessor record))
         (define-checked (modifier (record predicate) (val pred))
           (internal-modifier record val)))
       fields-to-process ...))
    ((_ type-name constructor predicate
        (fields ...) (field-wrappers ...) (field pred accessor)
        fields-to-process ...)
     (%define-record-type-checked
       type-name constructor predicate
       (fields ... (field internal-accessor))
       (field-wrappers
         ...
         (define-checked (accessor (record predicate))
           (internal-accessor record)))
       fields-to-process ...))))

(define-syntax %wrap-constructor
  (syntax-rules ()
    ((_ constructor internal-constructor (arg-names ...) (args ...))
     (define-checked (constructor args ...)
       (internal-constructor arg-names ...)))
    ((_ constructor internal-constructor (arg-names ...) (args ...)
        (name pred rest ...) fields-to-process ...)
     (%wrap-constructor constructor internal-constructor
                        (arg-names ... name) (args ... (name pred))
                        fields-to-process ...))))

(define-syntax define-record-type-checked
  (syntax-rules ()
    ((_ type-name (constructor constructor-args ...) predicate field ...)
     (begin
       (%define-record-type-checked
         type-name
         (internal-constructor constructor-args ...)
         predicate
         () () field ...)
       (%wrap-constructor constructor internal-constructor () () field ...)))))