;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; type-extensions.scm - Various CHICKEN type system extensions.
;;;
;;; This file should be loaded as a compiler extension with the
;;; `-extend` flag to `csc`(1):
;;;
;;;     $ csc -extend type-extensions.scm <FILE>
;;;
;;; Copyright (c) 2014-2018, Evan Hanson <evhan@foldling.org>
;;; BSD-style license. See LICENSE for details.
;;;

(import (matchable)
        (chicken platform)
        (chicken syntax))

(import-for-syntax (matchable))

;;
;; Hook type validation to expand compound type aliases.
;;
(letrec* ((rewrite-type
           (lambda (t)
             (let loop ((t t))
               (match t
                 (('list . _)
                  (rewrite-list-type t))
                 (((? symbol? a) . d)
                  (cond ((##sys#get a '##compiler#type-extension) =>
                         (lambda (f) (loop (apply f d))))
                        (else
                         (cons (loop a) (loop d)))))
                 ((a . d)
                  (cons (loop a) (loop d)))
                 (else t)))))
          (rewrite-list-type
           (lambda (t)
             (let loop ((xs (cdr t))
                        (ts '()))
               (match xs
                 (()
                  `(list ,@(reverse (map rewrite-type ts))))
                 ((x (? (cut eq? <> '...))) ; sic.
                  (foldl (lambda (d a) `(pair ,a ,d))
                         `(list-of ,(rewrite-type x))
                         (map rewrite-type ts)))
                 ((a . d)
                  (loop d (cons a ts)))
                 (x
                  (loop `(,x ...) ts)))))))

  (set! chicken.compiler.scrutinizer#validate-type
    (let ((validate-type chicken.compiler.scrutinizer#validate-type))
      (lambda (type name)
        (validate-type (rewrite-type type) name))))

  (set! chicken.compiler.scrutinizer#check-and-validate-type
    (let ((check-and-validate-type chicken.compiler.scrutinizer#check-and-validate-type))
      (lambda (type name #!optional location)
        (check-and-validate-type
         (rewrite-type (strip-syntax type))
         name
         location)))))

(let ((m (assq 'define-type ##sys#chicken.type-macro-environment)))
  (set! (caddr m)
    (##sys#ensure-transformer
     (ir-macro-transformer
      (lambda (e i c)
        (when (feature? 'compiling)
          (match e
            ((_ (? symbol? name))
             `(,(i '##core#elaborationtimeonly)
               (,(i '##sys#put/restore!)
                (quote ,(i name))
                (quote ,(i '##compiler#type-abbreviation))
                (quote (struct ,(i name))))))
            ((_ (name vars ...) type)
             `(,(i '##core#elaborationtimeonly)
               (,(i '##sys#put/restore!)
                (quote ,(i name))
                (quote ,(i '##compiler#type-extension))
                (lambda ,vars
                  (let loop ((t (quote ,(i type))))
                    (if (pair? t)
                        (cons (loop (car t))
                              (loop (cdr t)))
                        (case t
                          ,@(map (lambda (v i) `((,i) ,v))
                                 vars
                                 (map i vars))
                          (else t))))))))
            ((_ name type)
             `(,(i '##core#elaborationtimeonly)
               (,(i '##sys#put/restore!)
                (quote ,(i name))
                (quote ,(i '##compiler#type-abbreviation))
                (quote ,(i (chicken.compiler.scrutinizer#check-and-validate-type type 'define-type name))))))
            (else
             (syntax-error 'define-type "invalid type definition" (i e))))))))))