caster.scm

;; Cast

(load "pmatch.scm")
(load "types.scm")

;; Syntactic Extensions
(define-syntax fn-erase
  (syntax-rules (:)
    ((_ (: v type))
     '(v))
    ((_ ((: v type) v2 ...))
     `(v ,@(fn-erase v2 ...)))))

(define-syntax fn
  ;; Can we avoid the use of eval here?
  (syntax-rules (:)
    ((_ (: v type) (: return) body ...)
     (lambda (v) body ...))
    ((_ (: v type) body ...)
     (lambda (v) body ...))
    ((_ ((: v type) v2 ...) (: return) body ...)
     (fn ((: v type) v2 ...) body ...))
    ((_ ((: v type) v2 ...) body ...)
     (let ((env (the-environment)))
       (eval `(lambda ,(fn-erase ((: v type) v2 ...)) body ...) env)))
    ((_ () body ...)
     (lambda () body ...))))

(define-syntax listof
  (syntax-rules (:)
    ((_ (: type) . items)
     (quote items))))

;; Helpers
(define (all-equal? ls)
  (or (or (null? ls) (null? (cdr ls)))
      (and (equal? (car ls) (cadr ls))
	   (all-equal? (cdr ls)))))

;; Type environments are hash tables now
(define (te/new)
  (make-hash-table))

(define (te/lookup te var)
  (hash-table/get te var #f))

(define (te/extend te var type)
  (hash-table/put! te var type)
  te)

(define (te/nextend te vars types)
  (for-each (lambda (k v)
	      (hash-table/put! te k v))
	    vars types)
  te)

(define (te/remove te var)
  (hash-table/remove! te var)
  te)

(define (te/nremove te vars)
  (for-each (lambda (v) (te/remove te v)) vars)
  te)

(define (te/merge te1 te2)
  (define (put-func h1 h2)
    (hash-table/for-each
     h2 (lambda (k v) (hash-table/put! h1 k v))))
  (let ((te (te/new))) (put-func te te1) (put-func te te2) te))

(define (te/nmerge tes)
  (fold-left (lambda (acc x) (te/merge acc x)) (car tes) tes))

(define (te/copy te)
  (let ((te-copy (te/new)))
    (hash-table/for-each
     te (lambda (k v) (hash-table/put! te-copy k v)))
    te-copy))

(define te->alist hash-table->alist)
(define (alist->te alist)
  (let ((te (te/new)))
    (for-each (lambda (e)
		(hash-table/put! te (car e) (cdr e))) alist)
    te))

;; Cast Judgements
(define-structure (cj) te exp type)
;; / is more readable than - at times
(define cj/te cj-te)
(define cj/exp cj-exp)
(define cj/type cj-type)

;; Dealing with casts
(define (cast? exp)
  (and (pair? exp) (eq? (car exp) ':)))
(define (make-cast exp type)
  `(: ,exp ,type))
(define (uncast exp)
  (and (cast? exp) (cadr exp)))
(define cast/exp uncast)
(define (cast/type exp)
  (and (cast? exp) (caddr exp)))

;; Transform to intermediate langugae
;; Returns a cast judgement
(define (transform exp te)
  (pmatch
    exp
    (,e (guard (number? e)) (make-cj te (make-cast e 'number) 'number))

    (,e (guard (string? e)) (make-cj te (make-cast e 'string) 'string))

    (,e (guard (boolean? e)) (make-cj te (make-cast e 'boolean) 'boolean))

    (,e (guard (char? e)) (make-cj te (make-cast e 'char) 'char))

    (,e (guard (symbol? e))
	(let ((binding (te/lookup te e))) ; te/lookup returns #f if e not bound
	  (if binding
	      (make-cj te (make-cast e binding) binding)
	      (make-cj (te/extend te e 'any) (make-cast e 'any) 'any))))

    ((pair ,x ,y)
     (let* ((x-tf (transform x te))
	    (y-tf (transform y te))
	    (new-te (te/nmerge (list te (cj/te x-tf) (cj/te y-tf)))))
       (make-cj te (make-cast `(pair ,(cj/exp x-tf) ,(cj/exp y-tf))
			      (make-pair-type (cj/type x-tf) (cj/type y-tf)))
		(make-pair-type (cj/type x-tf) (cj/type y-tf)))))

    ((listof (: ,type) . ,items)
     (if (null? items)
	 (make-cj te (make-cast `(listof (: ,type)) (make-list-type type)) (make-list-type type))
	 (let* ((items-tf (map (lambda (e) (transform e te)) items))
		(items-types (map (lambda (e) (cj/type e)) items-tf))
		(items-cast (map (lambda (e) (cj/exp e)) items-tf)))
	   (if ;; (and (~ type (car items-types)) (all-equal? items-types))
	    (~ type (car items-types))
	       (let* ( ;; (new-te (te/nmerge (map (lambda (e) (te/extend (cj/te e)
		      ;; 						(uncast (cj/exp e))
		      ;; 						type)) items-tf)))
		      (new-te (te/nmerge (map (lambda (e) (cj/te e)) items-tf)))
		      (cast-items (map (lambda (e) (make-cast (uncast e) type)) items-cast))
		      (cast-exp (make-cast `(listof (: ,type) ,@cast-items) (make-list-type type))))
		 (make-cj (te/merge new-te te) cast-exp (make-list-type type)))
	       (error "TypeError: " 'list 'elements 'should 'type 'consistent)))))

    ((defvar (: ,var ,type) ,val)
     (let ((val-tf (transform val te)))
       (if (~ (cj/type val-tf) type)
	   (make-cj (te/merge (cj/te val-tf) (te/extend te var type))
		    (make-cast `(defvar (: ,var ,type) ,(cj/exp val-tf)) 'unit) 'unit)
	   (error "TypeError: " 'expected type 'in 'defvar var 'but 'got (cj/type val-tf)))))

    ((if ,pred ,clause1 ,clause2)
     (let ((pred-tf (transform pred te)))
       (if (equal? (cj/type pred-tf) 'boolean)
	   (let* ((clause1-tf (transform clause1 te))
		  (clause2-tf (transform clause2 te)))
	     (if ;; (equal? (cj/type clause1-tf) (cj/type clause2-tf))
	      (~ (cj/type clause1-tf) (cj/type clause2-tf))
	      (let ((new-te (te/nmerge (map (lambda (e) (cj/te e)) (list pred-tf clause1-tf clause2-tf))))
		    (cast-exp (make-cast `(if ,(cj/exp pred-tf) ,(cj/exp clause1-tf) ,(cj/exp clause2-tf))
					 (cj/type clause1-tf))))
		(make-cj (te/merge new-te te) cast-exp (cj/type clause1-tf)))
	      (error "TypeError: " 'if 'branches 'should 'have 'consistent 'type '- 'got
		     (cj/type clause1-tf) 'and (cj/type clause2-tf))))
	   (error "TypeError: " pred 'should 'be 'a 'boolean '- 'got (cj/type pred-tf)))))

    ((fn (: ,v ,type) ,body)
     (let* ((new-te (te/extend (te/copy te) v type))
	    (body-tf (transform body new-te))
	    (cast-exp (make-cast `(fn (: ,v ,type) ,(cj/exp body-tf))
				 `(-> ,type ,(cj/type body-tf)))))
       (make-cj (te/merge (te/remove new-te v) te) cast-exp `(-> ,type ,(cj/type body-tf)))))

    ;; ((fn (: ,v ,type) ,body)
    ;;  (let* ((new-te (te/extend (te/copy te) v type))
    ;; 	   (body-tf (transform body new-te)))
    ;;    (begin
    ;; 	(display "Type of ") (display v) (display " in body : ")
    ;; 	(display (te/lookup (cj/te body-tf) v)) (newline)
    ;; 	(if (not (equal? (te/lookup (cj/te body-tf) v) type))
    ;; 	    (error "Cannot use " v 'as (te/lookup (cj/te body-tf) v))
    ;; 	    (let ((cast-exp (make-cast `(fn (: ,v ,type) ,(cj/exp body-tf))
    ;; 				       `(-> ,type ,(cj/type body-tf)))))
    ;; 	      (make-cj (te/merge (te/remove new-te v) te) cast-exp `(-> ,type ,(cj/type body-tf))))))))

    ((fn (: ,v ,type) (: ,ret-type) ,body)
     (let* ((new-te (te/extend (te/copy te) v type))
	    (body-tf (transform body new-te))
	    (co-domain-type (cj/type body-tf)))
       (if (~ co-domain-type ret-type)
	   (let ((cast-exp (make-cast `(fn (: ,v ,type) (: ,ret-type) ,(cj/exp body-tf))
				      `(-> ,type ,ret-type))))
	     (make-cj (te/merge (te/remove new-te v) te) cast-exp
		      `(-> ,type ,ret-type)))
	   (error "TypeError : " 'expected ret-type 'got co-domain-type 'for body))))
   
    ((fn ((: ,v ,type) . ,res) ,body)
     (let* ((bindings (map (lambda (x) (cons (cast/exp x) (cast/type x))) res))
	    (new-te (te/merge (te/copy te) (alist->te (cons (cons v type) bindings))))
	    (body-tf (transform body new-te))
	    (co-domain-type (cj/type body-tf))
	    (cast-exp (make-cast `(fn ((: ,v ,type) ,@res) ,(cj/exp body-tf))
				 `(-> ,(list->pair-type (map (lambda (x) (cast/type x))
							     (cons `(: ,v ,type) res)))
				      ,co-domain-type))))
       (make-cj (te/merge te (te/nremove new-te
					 (map (lambda (x) (car x))
					      (cons (cons v type) bindings))))
		cast-exp (cast/type cast-exp))))
    
    ((defn (: ,name ,type) (,arg1 . ,args) ,body)
     ;; (display "DEFN")
     ;; (newline)
     (if (not (arrow-type? type))
	 (error "TypeError: " name 'should 'be 'a 'function 'type)
	 (if (pair-type? (domain type))
	     (if (not (equal? (pair-arity (domain type)) (length `(,arg1 ,@args))))
		 (error "TypeError: " 'expected (pair-arity (domain type)) 'arguments
			'for name '- 'got (length `(,arg1 ,@args)))
		 (let* ( ;; (bindings (map (lambda (x y) (cons (cdr (domain type))
			;; 				       `(arg1 ,@args)))))
			(new-te (te/nextend te `(,arg1 ,@args) (cdr (domain type))))
			;; (throw
			;;  (begin (display "*defn te: ") (display (te->alist new-te)) (newline) 3))
			(body-tf (transform body (te/extend new-te name type)))
			(cast-exp (make-cast `(defn (: ,name ,type) (,arg1 ,@args)
						,(cj/exp body-tf)) 'unit))
			;; (throw2
			;;  (begin (display "CAST-EXP: ") (display cast-exp) (newline)
			;; 	(display "BODY-TF-TYPE: ") (display (cj/type body-tf)) (newline)
			;; 	(display "**************") (newline) 3))
			)
		   (if (~ (cj/type body-tf) (co-domain type))
		       (make-tj (te/merge (te/extend te name type) (te/nremove new-te `(,arg1 ,@args)))
				cast-exp 'unit)
		       (error "TypeError: " 'inconsistent 'return 'type 'in 'defn name))))
	     ;; if not a pair type
	     (if (not (null? args))
		 (error "TypeError: " 'expected 1 'argument 'for name '-
			'got (length `(,arg1 ,@args)))
		 (let* ((new-te (te/nextend te (list name arg1) (list type (domain type))))
			(body-tf (transform body new-te))
			(cast-exp (make-cast `(defn (: ,name ,type) (,arg1)
						,(cj/exp body-tf)) 'unit)))
		   (if (~ (cj/type body-tf) (co-domain type))
		       (make-tj (te/merge te (te/remove new-te arg1))
				cast-exp 'unit)
		       (error "TypeError: " 'inconsistent 'return 'type 'in 'defn name)))))))

    ((fn ((: ,v ,type) . ,res) (: ,ret-type) ,body)
     (let* ((e-tf (transform `(fn ((: ,v ,type) . ,res) ,body) te))
	    (co-domain-type (cj/type e-tf)))
       (if (~ co-domain-type ret-type)
	   (let ((rt `(-> ,(list->pair-type (map (lambda (x) (cast/type x))
						 (cons `(: ,v ,type) res)))
			  ,ret-type)))
	     (make-cj (cj/te e-tf) (make-cast (uncast (cj/exp e-tf)) rt) rt))
	   (error "TypeError : " 'expected ret-type 'got co-domain-type))))
    
    ((,rator ,rand)			    ; function with arity 1
     (let* ((rator-tf (transform rator te)) ; a cast judgement (cj) object
	    (rand-tf (transform rand te))
	    (rator-type (cj/type rator-tf))
	    (rand-type (cj/type rand-tf)))
       ;; (display "Rator TE: ")
       ;; (display (te->alist (cj/te rator-tf)))
       ;; (newline)
       ;; (display "Rand TE: ")
       ;; (display (te->alist (cj/te rand-tf))) (newline)
       (if (any-type? rator-type)
	   (let ((new-te (te/merge (cj/te rand-tf)
				   (te/extend (cj/te rator-tf) rator `(-> ,rand-type any))))
		 (cast-exp (make-cast (list (make-cast (uncast (cj/exp rator-tf))
						       `(-> ,rand-type any))
					    (cj/exp rand-tf)) 'any)))
	     ;; (display "Final TE: ")
	     ;; (display (te->alist new-te)) (newline)
	     (make-cj new-te cast-exp 'any))
	   (if (arrow-type? rator-type)
	       (if (equal? (domain rator-type) rand-type)
		   (let ((new-te (te/merge (cj/te rator-tf) (cj/te rand-tf)))
			 (cast-exp (make-cast (list (cj/exp rator-tf) (cj/exp rand-tf))
					      (co-domain rator-type))))
		     (make-cj new-te cast-exp (co-domain rator-type)))
		   (if (and  (~ (domain rator-type) rand-type) (not (any-type? rand-type)))
		    ;; (~ (domain rator-type) rand-type)
		       (let ((new-te (te/merge (cj/te rator-tf)
					       (cj/te rand-tf)
					       ;; (te/extend (cj/te rand-tf)
					       ;; 		 rand (domain rator-type))
					       ))
			     (cast-exp (make-cast (list (cj/exp rator-tf)
							(make-cast (uncast (cj/exp rand-tf))
								   (domain rator-type)))
						  (co-domain rator-type))))
			 (make-cj new-te cast-exp (co-domain rator-type)))
		       (error "TypeError : " 'expected (domain rator-type) 'got rand-type
			      'for rand)))
	       (error "TypeError : " rator 'must 'be 'a 'function)))))
    
    ((,rator . ,rands)
     (let* ((rator-tf (transform rator te))
	    (rator-type (cj/type rator-tf))
	    (rand-tf (map (lambda (e) (transform e te)) rands))
	    (rand-cast (map (lambda (e) (cj/exp e)) rand-tf))
	    (rand-type (list->pair-type (fold-right (lambda (x acc)
						      (cons (cj/type x) acc)) '() rand-tf))))
       ;; (display "Rand type ")
       ;; (display rand-type)
       ;; (newline)
       ;; (display "Rator type ")
       ;; (display rator-type) (newline)
       (if (any-type? rator-type)
	   (let* ((new-te (te/nmerge (append (list (te/extend (cj/te rator-tf)
							      rator `(-> ,rand-type any)))
					     (map (lambda (x) (cj/te x)) rand-tf))))
		  (cast-exp (make-cast `((: ,rator (-> ,rand-type any)) ,@rand-cast) 'any)))
	     (make-cj new-te cast-exp 'any))
	   (if (arrow-type? rator-type)
	       (if (or (equal? (domain rator-type) rand-type)
		       (and (equal? (arity rator-type) 'n)
			    (equal? (domain rator-type) (cadr rand-type))
			    (all-equal? (cdr rand-type))))
		   (let* ((new-te (te/nmerge (append (list (cj/te rator-tf))
						     (map (lambda (x) (cj/te x)) rand-tf))))
			  (cast-exp `(: (,(cj/exp rator-tf) ,@rand-cast) ,(co-domain rator-type))))
		     (make-cj new-te cast-exp (co-domain rator-type)))
		   (if (or (~ (domain rator-type) rand-type)
			   (and (equal? (arity rator-type) 'n)
				(~ (domain rator-type) (cadr rand-type))
				(all-equal? (cdr rand-type))))
		       (let* ( ;; (rc (map (lambda (x y) (make-cast (uncast x) y))
			      ;; 	      rand-cast
			      ;; 	      (if (pair? (domain rator-type))
			      ;; 		  (cdr (domain rator-type))
			      ;; 		  (domain rator-type))))
			      (rc (if (pair? (domain rator-type))
				      (map (lambda (x y) (make-cast (uncast x) y))
					   rand-cast (cdr (domain rator-type)))
				      (map (lambda (x) (make-cast (uncast x) (domain rator-type)))
					   rand-cast)))
			      ;; (rc-te (map (lambda (x y z) (te/extend (cj/te x) y z))
			      ;; 		 rand-tf rands (cdr (domain rator-type))))
			      (rc-te (if (pair? (domain rator-type))
					 (map (lambda (x y z) (te/extend (cj/te x) y z))
					      rand-tf rands (cdr (domain rator-type)))
					 (map (lambda (x y) (te/extend (cj/te x) y
								       (domain rator-type)))
					      rand-tf rands)))
			      (new-te (te/nmerge (append (list (cj/te rator-tf)) rc-te)))
			      (cast-exp `(: (,(cj/exp rator-tf) ,@rc) ,(co-domain rator-type))))
			 (make-cj new-te cast-exp (co-domain rator-type)))
		       (begin ;; (display rand-type)
			      ;; (newline)
			      ;; (display rator-type)
			      (error "TypeError: " 'inconsistent 'argument 'types 'for rator))))
	       (error "TypeError: " rator 'must 'be 'a 'function)))))
    (else (error "TRANSFORM -- Unknown type for --" exp))))

(define (c exp alis)
  (let ((res (transform exp (alist->te alis))))
    (display (cj/exp res)) (newline)
    (display "Γ = ") (display (te->alist (cj/te res)))
    (display " => ") (display exp)
    (display " : ") (display (cj/type res)) (newline)
    res))

;; Typing Judgements
(define-structure (tj) te exp type)
(define tj/te tj-te)
(define tj/exp tj-exp)
(define tj/type tj-type)

(define (check expr tenv)
  (define (tc exp type te)
    (pmatch
      exp
      ((: ,e ,t) (guard (and (~ type t) (not (equal? type t))))
       ;; (display "Casting ")
       ;; (display e) (display " from ") (display t) (display " to ")
       ;; (display type)
       ;; (newline)
       (tc `(: ,e ,type) type te))
      ((: ,e ,t) (guard (symbol? e))
       (let ((binding (te/lookup te e)))
	 (if binding
	     (if (~ binding type)
		 ;; (make-tj (te/extend te e type) exp binding)
		 (make-tj te exp binding)
		 (error "TypeError: " 'expected binding 'got type 'for e))
	     ;; The part below never actually runs (?) need to test a bit more. 
	     (make-tj (te/extend e t) exp t))))
      ((: (defvar (: ,v ,s) ,val) ,t)
       (let* ((tc-val (tc val s te)))
	 (if (~ (tj/type tc-val) s)
	     (let ((new-te (te/merge (tj/te tc-val) (te/extend te v s))))
	       (make-tj new-te exp t))
	     (error "TypeError: " 'value 'type (tj/type tc-val) 'is 'not 'consistent 'with s))))

      ((: (pair ,x ,y) ,t)
       (let* ((tc-x (tc x (cadr t) te))
	      (tc-y (tc y (caddr t) te))
	      (new-te (te/merge (tj/te tc-x) (tj/te tc-y))))
	 (make-tj (te/merge te new-te) exp t)))
      
      ((: (listof (: ,s) . ,items) ,t)
       (if (null? items)
	   (make-tj te exp t)
	   (let* ((tc-items (map (lambda (e) (tc e s te)) items))
		  ;; (items-te (te/nmerge (map (lambda (e) (te/extend (tj/te e)
		  ;; 						      (cast/exp (tj/exp e))
		  ;; 						      s)) tc-items)))
		  (items-te (te/nmerge (map (lambda (e) (tj/te e)) tc-items)))
		  (new-te (te/merge te items-te)))
	     (make-tj new-te exp t))))
      ((: (if ,pred ,clause1 ,clause2) ,t)
       (let* ((tc-pred (tc pred 'boolean te))
	      (tc-clause1 (tc clause1 type te))
	      (tc-clause2 (tc clause2 type te))
	      (new-te (te/nmerge (map (lambda (e) (tj/te e)) (list tc-pred tc-clause1 tc-clause2)))))
	 (make-tj (te/merge te new-te) exp t)))
      ((: (defn (: ,v ,s) (,arg1 . ,args) ,body) ,t)
       ;; (display (length `(,arg1 ,@args)))
       ;; (newline)
       (if (pair-type? (domain s))
	   (let* ((new-te (te/extend (te/nextend te `(,arg1 ,@args) (cdr (domain s))) v s))
		  
		  (tc-body (tc body (co-domain s) new-te)))
	     (if (~ (tj/type tc-body) (co-domain s))
		 (make-tj new-te exp t)
		 (error "TypeError: " 'inconsistent 'function 'body 'in v)))
	   (let* ((new-te (te/extend te arg1 (domain s)))
		  (tc-body (tc body (co-domain s) (te/extend new-te v s))))
	     (if (~ (tj/type tc-body) (co-domain s))
		 (make-tj new-te exp t)
		 (error "TypeError: " 'inconsistent 'function 'body 'in v)))))
      ((: (fn (: ,v ,s) ,body) ,t)
       ;; check consistency with t? 
       (let ((tc-body (tc body (co-domain type) (te/extend (te/copy te) v s))))
	 (if (~ (tj/type tc-body) (co-domain type))
	     (make-tj (te/remove (tj/te tc-body) v) exp t)
	     (error "TypeError: " 'expected (co-domain type) 'got (tj/type tc-body)
		    'for body))))
      ((: (fn (: ,v ,s) (: ,ret) ,body) ,t)
       (let ((tc-body (tc body (co-domain type) (te/extend (te/copy te) v s))))
	 (if (~ (tj/type tc-body) ret)
	     (make-tj (te/remove (tj/te tc-body) v) exp t)
	     (error "TypeError: " 'expected (co-domain type) 'got (tj/type tc-body)
		    'for body))))
      ((: (,rator ,rand) ,t)
       (let* ( ;; (tc-rator (tc rator `(-> ,(cast/type rand) ,type) te))
	      (tc-rator (tc rator (cast/type rator) te))
	      (tc-rand (tc rand (domain (cast/type rator)) te))
	      (new-te (te/merge (tj/te tc-rator) (tj/te tc-rand))))
	 (make-tj new-te exp t)))
      ((: (,rator . ,rands) ,t)
       (let* ((tc-rator (tc rator (cast/type rator) te))
	      ;; Check each rand
	      (tc-rands (map (lambda (e) (tc e (cast/type e) te)) rands))
	      (rands-te (te/nmerge (map (lambda (e) (tj/te e)) tc-rands)))
	      (new-te (te/merge (tj/te tc-rator) rands-te)))
	 (make-tj new-te exp t)))
      ((: ,e char) (guard (char? e)) (make-tj te exp 'char))
      ((: ,e boolean) (guard (boolean? e)) (make-tj te exp 'boolean))
      ((: ,e string) (guard (string? e)) (make-tj te exp 'string))
      ((: ,e number) (guard (number? e)) (make-tj te exp 'number))
      ((: ,e any) (make-tj te exp 'any))
      (else (error "Unknown type for --" exp))))
  (let ((c (transform expr ;; (alist->te tenv)
		      tenv)))
    (tc (cj/exp c) (cj/type c) (cj/te c))))

(define (t exp te)
  (let ((res (check exp (alist->te te))))
    ;; (display (tj/exp res))
    ;; (newline)
    (display "Γ{") (display (te->alist (tj/te res)))
    (display "} ⊦ ") (display exp)
    (display " : ") (display (tj/type res)) (newline)
    res))