;; The first three lines of this file were inserted by DrScheme. They record metadata
;; about the language level of this file in a form that our tools can easily process.
(require (planet plai/plai:1:3))
;#reader(lib "reader.ss" "plai" "lang")
;;
;; Lecture XX : Type Checking
;;

;------------------------------------------------------------------
; Note: The code in this file may be used as starter code for
; your assignment. But you must include a clear statement at the top
; of your file acknowledging the source and the extent of copying.
;
; E.g: "This code was based on the CS311 TFAE type checking
; code provided by Kris De Volder. We copied all the TFAE code
; except for the "type-check" function, which we rewrote from
; scratch."

; Alternatively, you may simply develop your own code all from scratch
; and ommit this statement.

; Remember that using code from anywhere (internet / lecture notes)
; without explicitly acknowledging the source is plagiarism.
;
; Using publically available code from anywhere *is* permitted *with*
; acknowledgement. However, if little of "your" code is actually written
; by you, then this may be reflected in your mark.
;
; For the CS311 assignment I recommend the following:
;   - start from this code and use it as a general outline
;   - fill in additional cases for parse / unparse as needed
;   - add your own test cases.
;   - !!!rewrite the code in "type-check" from scratch!!!!
;
; If you do this you will not loose marks for copying.
;
; In other words I recommend that you do NOT copy the "core" of the
; typechecker algorithm. Writing this yourself is the main point
; of doing the assignment.
;
; !!! If you decide to copy code from this file for your assignment,
; make sure it is clearly and explicitly stated whether you copied
; any of the code in the type-check function!!!!
;
; If you do copy the "type-check" code, you may loose marks (for
; not doing important portions of the assignment yourself).
; ------------------------------------------------------------------

; In the book, and in the lectures so far, we have mainly considered
; how to write typing rules / judgements that formally define a
; type system.

; How do we turn rules like these into a type-checker algorithm?

; Today, let's consider this question and write ourselves a checker
; for a simple language: TFAE.

#|
<TFAE> ::= <id>
         | true
         | false
         | <number>
         | {+ <TFAE> <TFAE>}
         | {- <TFAE> <TFAE>}
         | {* <TFAE> <TFAE>}
         | {izero <TFAE>}
         | {if <TFAE> <TFAE> <TFAE>}
         | {fun {<id> : <Type>} : <Type> <TFAE>}
         | {rec {<id> {fun {<id> : <Type>} : <Type> <TFAE>}} <TFAE>}
         | {with {<id> <expr>} <expr>}
         | {<TFAE> <TFAE>}
         | nempty
         | {ncons <TFAE> <TFAE>}
         | {nempty? <TFAE>}
         | {nfirst <TFAE>}
         | {nrest <TFAE>}

First let's recall the type-language (i.e. the BNF grammar
that defines the types of our language).

<Type> ::= number
         | boolean
         | nlist
         | (<Type> -> <Type>)

TODO    add additional cases

|#


; As usual, we will need a representation for abstract
; syntax trees for the expressions in our language.
(define-type Exp
  (num    (v number?))
  (bool   (v boolean?))
  (id     (name symbol?))
  (add    (l Exp?) (r Exp?))
  (sub    (l Exp?) (r Exp?))
  (mul    (l Exp?) (r Exp?))
  (iszero (a Exp?))
  (is-nempty (a Exp?))
  (nfirst (a Exp?))
  (nrest (a Exp?))
  (iff    (tst Exp?) (thn Exp?) (els Exp?))
  (fun    (id symbol?) (arg-type Type?) (ret-type Type?) (body Exp?))
  (rec0   (id symbol?) (func fun?) (body Exp?))
  (with   (id symbol?) (val Exp?) (body Exp?))
  (app    (rator Exp?) (rand Exp?))
  (nempty)
  (nlist (l Exp?) (r Exp?)))

; We will also need an AST representation for (parsed) types!
(define-type Type
  (numT)
  (boolT)
  (nlistT)
  (funT (arg Type?) (ret Type?)))

; and procedures to parse the concrete notations into ASTs
(define (parse sexp)
  (cond ((memv sexp '(true false))
         (bool (eqv? sexp 'true)))
        ((eqv? sexp 'nempty)
         (nempty))
        ((symbol? sexp)
         (id sexp))
        ((number? sexp)
         (num sexp))
        ((pair? sexp)
         (case (car sexp)
           ((ncons) (apply nlist (map parse (cdr sexp))))
           ((if) (apply iff (map parse (cdr sexp))))
           ((+) (apply add (map parse (cdr sexp))))
           ((-) (apply sub (map parse (cdr sexp))))
           ((*) (apply mul (map parse (cdr sexp))))
           ((nempty?) (apply is-nempty (map parse (cdr sexp))))
           ((nfirst) (apply nfirst (map parse (cdr sexp))))
           ((nrest) (apply nrest (map parse (cdr sexp))))
           ((iszero) (apply iszero (map parse (cdr sexp))))
           ((with) ;{with {<id> <expr>} <expr>}
            (with (caar (cdr sexp))
                  (parse (cadr (cadr sexp)))
                  (parse (caddr sexp))))
           ((rec) ;{rec {<id> {fun {<id> : <Type>} : <Type> <TFAE>}} <TFAE>}
            (rec0 (caar (cdr sexp))
                  (parse (cadr (cadr sexp)))
                  (parse (caddr sexp))))
           ((fun) ;{fun {x : type} : type exp}
            (check! "Need : before arg-type decl" sexp (eq? ': (cadr (cadr sexp))))
            (check! "Need : before ret-type decl" sexp (eq? ': (caddr sexp)))
            (fun (car (cadr sexp))
                 (parse-type (caddr (cadr sexp)))
                 (parse-type (cadddr sexp))
                 (parse      (cadr (cdddr sexp)))))
           (else
            (apply app (map parse sexp)))))))

(define (check! message data bool)
  (if (not bool)
      (error message data)
      'ok))

(define (parse-type sexp)
  (cond ((pair? sexp)
         (check! "Expect list of 3" sexp (= 3 (length sexp)))
         (check! "Expect arrow in second place" sexp (eq? '-> (cadr sexp)))
         (funT (parse-type (car sexp))
               (parse-type (caddr sexp))))
        (else
         (case sexp
           ((number)  (numT))
           ((boolean) (boolT))
           ((list) (nlistT))
           (else
            (error "Syntax error in type" sexp))))))

; - - - - - - ----------------------------------------------------------

; To produce more readable error messages, let's implement some functions
; to convert (unparse) parsed internal representation of types and
; expressions back into their more readable external representations
(define (unparse-type type)
  (type-case Type type
    (numT () 'number)
    (boolT () 'boolean)
    (nlistT () 'list)
    (funT (a r) `(,(unparse-type a) -> ,(unparse-type r)))))

(define (unparse x)
  (cond ((Type? x)
         (unparse-type x))
        ((Exp? x)
         (unparse-exp x))
        (else
         x)))

(define (unparse-exp exp)
  (type-case Exp exp
    (num (v) v)
    (bool (v) (if v 'true 'false))
    (id (name) name)
    (add (l r) `{+ ,(unparse l) ,(unparse r)})
    (sub (l r) `{- ,(unparse l) ,(unparse r)})
    (mul (l r) `{* ,(unparse l) ,(unparse r)})
    (nempty () `nempty)
    (nlist (l r) `{ncons ,(unparse l) ,(unparse r)})
    (iszero (a) `{iszero ,(unparse a)})
    (is-nempty (a) `{nempty? ,(unparse a)})
    (nfirst (a) `{nfirst ,(unparse a)})
    (nrest (a) `{nrest ,(unparse a)})
    (iff (tst thn els)
         `(if ,(unparse tst)
              ,(unparse thn)
              ,(unparse els)))
    (fun (id atype rtype body)
         `{fun {,id : ,(unparse atype)} : ,(unparse rtype)
               ,(unparse body)})
    (rec0 (id func body)
         `{rec {,id ,(unparse func)} ,(unparse body)})
    (with (id val body)
         `{with {,id ,(unparse val)} ,(unparse body)})
    (app (f a)
         `{,(unparse f) ,(unparse a)})))

; -----------------------------------------------------------------------------
; Now let's implement the type checker. Based on the typing rules, which
; we will write on the board.




















; We will need a representation for the type env "Gamma":
(define-type Env
  (anEnv (id symbol?) (type Type?) (more Env?))
  (mtEnv))

;We will write the type-checker in a style similar to
;the assignment. I've already set up a basic structure for this
;below.

; type-of :: Exp Continuation<String> -> Type ... or error
;   Returns either a Type (type checking passes)
;        or signals an error by passing an error
;        message String to the escape continuation.
(define (type-of exp esc)

  (define (err . args)
    (esc (apply format args)))

  (define (lookup name env)
    (type-case Env env
       (mtEnv () (err "Unbound identifier: ~a" name))
       (anEnv (id tp more)
              (if (eq? id name) tp (lookup name more)))))

  (define (check-equals! exp part expect got)
    (if (equal? expect got)
        'ok
        (err "~a: ~a has type ~a but expected ~a"
             (unparse exp) (unparse part) (unparse got) (unparse expect))))

  ; Exp Env -> Type (or error)
  ; rewrite
  (define (type-check exp env)
    (type-case Exp exp
      (nempty () (nlistT))
      (nlist (l r)
             (let ((l-type (type-check l env))
                   (r-type (type-check r env)))
               (nlistT)))
      (num (v) (numT))
      (bool (v) (boolT))
      (add (l r)
           (let ((l-type (type-check l env))
                 (r-type (type-check r env)))
             (check-equals! exp l (numT) l-type)
             (check-equals! exp r (numT) r-type)
             (numT)))
      (sub (l r)
           (let ((l-type (type-check l env))
                 (r-type (type-check r env)))
             (check-equals! exp l (numT) l-type)
             (check-equals! exp r (numT) r-type)
             (numT)))
      (mul (l r)
           (let ((l-type (type-check l env))
                 (r-type (type-check r env)))
             (check-equals! exp l (numT) l-type)
             (check-equals! exp r (numT) r-type)
             (numT)))
      (iszero (a)
              (begin (check-equals! exp a (numT) (type-check a env))
                     (boolT)))
      (is-nempty (a)
               (begin (check-equals! exp a (nlistT) (type-check a env))
                      (boolT)))
      (nfirst (a)
              (begin (check-equals! exp a (nlistT) (type-check a env))
                     (type-check (nlist-l a) env)))
      (nrest (a)
              (begin (check-equals! exp a (nlistT) (type-check a env))
                     (type-check (nlist-r a) env)))
      (id (name)
           (lookup name env))
      (iff (tst thn els)
           (let ((tst-t (type-check tst env))
                 (thn-t (type-check thn env))
                 (els-t (type-check els env)))
             (check-equals! exp tst (boolT) tst-t)
             (check-equals! exp els thn-t els-t)
             thn-t))
      (fun (a a-type r-type body)
           (let ((body-t (type-check body
                                     (anEnv a a-type env))))
             (check-equals! exp body r-type body-t)
             (funT a-type r-type)))
      (rec0 (id func body)
           (let ((f-type (type-check func
                                     (anEnv id (funT (fun-arg-type func) (fun-ret-type func)) env)))
                 (b-type (type-check body
                                     (anEnv id (funT (fun-arg-type func) (fun-ret-type func)) env))))
             b-type))
      (with (id val body)
            (let* ((v-type (type-check val env))
                   (b-type (type-check body (anEnv id v-type env))))
              b-type))
      (app (f a)
           (let ((f-type (type-check f env))
                 (a-type (type-check a env)))
             (if (funT? f-type) 'ok
                 (err "~a: ~a has type ~a but expected a function"
                      (unparse exp) (unparse f) (unparse f-type)))
             (check-equals! exp a (funT-arg f-type) a-type)
             (funT-ret f-type)))))

  ;Assignment specification says we should return an
  ;external representation of type, so call unparse!
  (unparse-type (type-check exp (mtEnv))))

;-------------------------------------------------------------------------------------
; Tests for the parse function:
; Note these tests depend on the shape of the parse tree
; and the precise definition of AST types. So these test
; may not pass for student code (if, as is likely, AST reps are
; chosen a little differently)

(test (parse '4)
      (num 4))
(test (parse  'true)
      (bool #t))
(test (parse  'false)
      (bool #f))
(test (parse  '{+ 1 2})
      (add (num 1) (num 2)))
(test (parse  '{iszero 1})
      (iszero (num 1)))
(test (parse  '{if true 1 2})
      (iff (bool #t) (num 1) (num 2)))
(test (parse  'x)
      (id 'x))
(test (parse  '{fun {x : number} : number {+ x 1}})
      (fun 'x (numT) (numT) (add (id 'x) (num 1))))
(test (parse '{f x})
      (app (id 'f) (id 'x)))

(test (parse-type 'number)
      (numT))
(test (parse-type 'boolean)
      (boolT))
(test (parse-type '(number -> (number -> boolean)))
      (funT (numT) (funT (numT) (boolT))))

;---- tests for the type checking --------------------------------------

(define *precise-errors* #t)
   ; set this to false to only check whether errors are generated, but
   ; not whether the message is matching.
   ; (should work for student code, even if error messages are not
   ; identical to ours).

; We define the following to easily call "type-of" with
; an aproprate escaper continuation:
(define (tc sexp)
  (let/cc k
    (type-of (parse sexp)
             (lambda (msg)
               (k (if *precise-errors*
                      msg
                      'error))))))

; Helper for testing: "normalizes" error messages
; to a single 'error value.
(define (expect x)
  (if (string? x)
      (if *precise-errors* ;#f -> treat all messages as the same.
          x
          'error)
      x))

(define-syntax tc-test
  (syntax-rules ()
    ((_ exp expected)
     (test (tc exp) (expect expected)))))

;Numbers
(tc-test '3 'number)

;Booleans
(tc-test 'true 'boolean)
(tc-test 'false 'boolean)

;Binary operator (what about number of operands???)
(tc-test '{+ 3 4} 'number)
(tc-test '{+ 3 {+ 2 4}} 'number)
(tc-test '{+ true 4} "(+ true 4): true has type boolean but expected number")
(tc-test '{+ 4 false} "(+ 4 false): false has type boolean but expected number")

;Unary operator
(tc-test '{iszero 0} 'boolean)
(tc-test '{iszero false} "(iszero false): false has type boolean but expected number")

;(Unbound) identifier
(tc-test 'x "Unbound identifier: x")

;If
(tc-test '{if {+ 3 4} true true} "(if (+ 3 4) true true): (+ 3 4) has type number but expected boolean")
(tc-test '{if true true true} 'boolean)
(tc-test '{if false 1 2} 'number)
(tc-test '{if false 1 true} "(if false 1 true): true has type boolean but expected number")

; Fun and app
(tc-test '{fun {x : number} : boolean {iszero {+ x -1}}} '(number -> boolean))
(tc-test '{fun {x : boolean} : boolean {iszero {+ x -1}}}
         "(+ x -1): x has type boolean but expected number")
(tc-test '{fun {x : number} : number {iszero {+ x -1}}}
         "(fun (x : number) : number (iszero (+ x -1))): (iszero (+ x -1)) has type boolean but expected number")
(tc-test '{{fun {x : number} : boolean {iszero {+ x -1}}}
           10}
         'boolean)
(tc-test '{{fun {x : number} : boolean {iszero {+ x -1}}}
           true}
         "((fun (x : number) : boolean (iszero (+ x -1))) true): true has type boolean but expected number")
(tc-test '{1 2}
         "(1 2): 1 has type number but expected a function")

(tc-test '{rec {f {fun {x : number} : number {f {- x 1}}}} {f 5}} 'number)
(tc-test '{with {val 5} {iszero val}} 'boolean)