(cs441-load "sets.scm")
(cs441-load "parser.scm")

;
; new records for representing MicroScheme values
;
(define-record Prim (name))
(define-record Closure (formals body env))

;
; finite functions
;
(define make-empty
  (lambda ()
    (lambda (x) #f)))

(define extend
  (lambda (f d r)
    (lambda (x)
      (if (eq? x d)
          r
          (f x)))))

(define lookup
  (lambda (f x)
    (f x)))

(define extend-list
  (lambda (env doms rans)
    (if (null? doms)
        env
        (extend-list
          (extend env (car doms) (car rans))
          (cdr doms) 
          (cdr rans)))))

;
; general purpose helpers
;
(define firsts
  (lambda (lst)
    (map car lst)))

(define seconds
  (lambda (lst)
    (map cadr lst)))

(define last 
  (lambda (lst)
    (if (null? (cdr lst))
        (car lst)
        (last (cdr lst)))))

;; Last element of a list.
(define rac
  (lambda (l)
    (if (null? (cdr l))
        (car l)
        (rac (cdr l)))))

;; All but the last element of a list.
(define rdc
  (lambda (l)
    (if (null? (cdr l))
        '()
        (cons (car l) (rdc (cdr l))))))

;; Scheme's map does not guarantee what order it walks the list
(define mapLR
  (lambda (f l)
    (if (null? l)
        '()
        (let ((x (f (car l))))
          (cons x (mapLR f (cdr l))))))) 

;; Map down list, oring results.
(define ormap
  (lambda (f l)
    (if (null? l)
        (or)
        (or (f (car l)) (ormap f (cdr l))))))
;
; box related functions
;
(define-record box (contents))
(define unbox box->contents)
(define set-box! (lambda (a b) (vector-set! a 1 b)))

(define swap-boxes
  (lambda (a b)
    (if (null? a) #f
         (let ((temp (unbox (car b))))
               (set-box! (car b) (unbox (car a)))
               (set-box! (car a) temp)
               (swap-boxes (cdr a) (cdr b))))))

;
; type representations
;
(define make-tvar
  (let ((n 0))
    (lambda x
      (set! n (+ 1 n))
      (make-box n))))

(define tvar?
  (lambda (x) (and (box? x) (number? (unbox x)))))

(define-record list-type (elt))
(define-record arrow-type (params result))

(define parse-type
  (lambda (s)
    (let ((env (make-empty)))
      (letrec ((parse (lambda (s)
                        (cond ((memq s '(num bool sym))
                               s)
                              ((and (list? s) (eq? 'quote (car s)))
                               (or (lookup env (cadr s))
                                   (let ((tv (make-tvar)))
                                     (set! env (extend env (cadr s) tv))
                                     tv)))
                              ((and (list? s) (eq? (cadr s) 'list))
                               (make-list-type (parse (car s))))
                              ((and (list? s) (eq? (rac (rdc s)) '->))
                               (make-arrow-type
                                 (mapLR parse (rdc (rdc s)))
                                 (parse (rac s))))
                              (else (error "incorrect type syntax"))))))
        (parse s)))))

(define unparse-type
  (lambda (t)
    (cond ((tvar? t)
           (string->symbol (string-append "a" (number->string (unbox t)))))
          ((list-type? t)
           (list (unparse-type (list-type->elt t)) 'list))
          ((arrow-type? t)
           (append
             (map unparse-type (arrow-type->params t))
             (list '-> (unparse-type (arrow-type->result t)))))
          ((box? t)
           (unparse-type (unbox t)))
          (else t))))

;;        name     type                         value
(define initial-info
  (list
    (list '+       '(num num -> num)            +)
    (list '-       '(num num -> num)            -)
    (list '*       '(num num -> num)            *)
    (list '/       '(num num -> num)            /)
    (list '=       '(num num -> bool)           =)
    (list 'not     '(bool -> bool)              not)
    (list 'symbol= '(sym sym -> bool)           eq?)
    (list 'cons    '('a ('a list) -> ('a list)) cons)
    (list 'car     '(('a list) -> 'a)           car)
    (list 'cdr     '(('a list) -> ('a list))    cdr)
    (list 'pair?   '(('a list) -> bool)         pair?)
    (list 'null?   '(('a list) -> bool)         null?)
    (list 'display '('a -> bool)                (lambda (x) (display x) f))
    (list 'nil     '('a list)                   ())))

(define initial-type-env
  (extend-list (make-empty)
    (map car initial-info)
    (map (lambda (x) (lambda () (parse-type (cadr x)))) initial-info)))

(define initial-env
  (extend-list (make-empty)
    (map car initial-info)
    (map caddr initial-info)))

;; Type Checker

(define type-check
  (lambda (e tenv)
    (variant-case e
      (Const (value)
        (cond ((number? value) 'num)
              ((symbol? value) 'sym)
              ((boolean? value) 'bool)
              (else (error "bad constant"))))
      (Var (name)
        (let ((t (lookup tenv name)))
          (cond ((not t) (error "unbound name" name))
                ((procedure? t) (t))
                (else t))))
      (Beg (exprs)
        (last (mapLR (lambda (x) (type-check x tenv)) exprs)))
      (If (test then else)
        (unify 'bool (type-check test tenv))
        (let ((r (type-check then tenv)))
          (unify r (type-check else tenv))
          r))
      (Let (bindings body)
        (let ((vars (firsts bindings))
              (types (mapLR (lambda (x) (lambda () (type-check x tenv))) (seconds bindings))))
          (for-each (lambda (x) (type-check x tenv)) (seconds bindings))
          (type-check body (extend-list tenv vars types))))
      (Let* (bindings body)
        (if (null? bindings)
            (type-check body tenv)
            (let ((x (car (car bindings)))
                  (e (cadr (car bindings)))
                  (rest (cdr bindings)))
              (type-check e tenv)
              (type-check (make-Let* rest body)
                (extend tenv x (lambda () (type-check e tenv)))))))
      (Letrec (bindings body)
        (let ((vars (firsts bindings))
              (exps (seconds bindings)))
          (let* ((tvs (map make-tvar bindings))
                 (new-tenv (extend-list tenv vars tvs)))
            (for-each (lambda (e tv)
                        (unify (type-check e new-tenv) tv))
              exps tvs))
          (let ((new-tenv (extend-list tenv vars
                            (map (lambda (x)
                                   (lambda ()
                                     ;; Unlike the let case, we must type check
                                     ;; every letrec binding since the others might
                                     ;; place type constraints on the one we want.
                                     (let* ((tvs (map make-tvar bindings))
                                            (new-tenv (extend-list tenv vars tvs)))
                                       (for-each (lambda (e tv)
                                                   (unify (type-check e new-tenv) tv))
                                         exps tvs)
                                       (lookup new-tenv x))))
                                 vars))))
            (type-check body new-tenv))))
      (Lam (formals body)
        (let* ((tformals (map make-tvar formals))
               (r (type-check body (extend-list tenv formals tformals))))
          (make-arrow-type tformals r)))
      (Ap (fun args)
        (let* ((tfun (type-check fun tenv))
               (targs (map (lambda (x) (type-check x tenv)) args))
               (r (make-tvar)))
          (unify tfun (make-arrow-type targs r))
          r)))))

(define unify
  (lambda (a b)
    (cond ((eq? a b)
           #t)
          ((and (list-type? a) (list-type? b))
           (unify (list-type->elt a) (list-type->elt b)))
          ((and (arrow-type? a) (arrow-type? b))
           (for-each unify (arrow-type->params a) (arrow-type->params b))
           (unify (arrow-type->result a) (arrow-type->result b)))
          ((and (box? a) (not (tvar? a)))
           (unify (unbox a) b))
          ((and (box? b) (not (tvar? b)))
           (unify a (unbox b)))
          ((tvar? a)
           (if (occurs? a b)
               (error "cannot solve " (unparse-type a) '= (unparse-type b))
               (set-box! a b)))
          ((tvar? b)
           (unify b a))
          (else (error "cannot solve " (unparse-type a) '= (unparse-type b))))))

(define occurs?
  (lambda (v t)
    (cond ((eq? v t) #t)
          ((tvar? t) #f)
          ((box? t) (occurs? v (unbox t)))
          ((list-type? t) (occurs? v (list-type->elt t)))
          ((arrow-type? t) (or (occurs? v (arrow-type->result t))
                               (ormap (lambda (t) (occurs? v t))
                                 (arrow-type->params t))))
          (else #f))))

;; Evaluator

(define interp
  (lambda (e env)
    (variant-case e
      (Const (value)
        value)
      (Var (name) (lookup env name))
      (Beg (exprs)
        (last (mapLR (lambda (x) (interp x env)) exprs)))
      (If (test then else)
          (if (interp test env)
              (interp then env)
              (interp else env))) 
      (Let (bindings body)
        (let ((vars (firsts bindings))
              (vals (mapLR (lambda (x) (interp x env)) (seconds bindings))))
          (interp body (extend-list env vars vals))))
      (Let* (bindings body)
        (if (null? bindings)
            (interp body env)
            (let ((x (car (car bindings)))
                  (e (cadr (car bindings)))
                  (rest (cdr bindings)))
              (interp (make-Let* rest body)
                (extend env x (interp e env))))))
      (Letrec (bindings body)
        (let ((vars (firsts bindings))
              (exps (seconds bindings)))
          (letrec ((new-env (lambda (x)
                              ((extend-list env vars
                                 (map (lambda (x) (interp x new-env)) exps))
                               x))))
            (interp body new-env))))
      (Lam (formals body) 
        (lambda args
          (if (= (length args) (length formals))
              (interp body (extend-list env formals args))
              (error "wrong number of args to lambda"))))
      (Ap (fun args)
        (let* ((vfun (interp fun env))
               (vargs (mapLR (lambda (x) (interp x env)) args)))
          (if (procedure? vfun)
              (apply vfun vargs)
              (error "not a procedure: " vfun)))))))

(define eval
  (lambda (expr)
    (let ((a (parse expr)))
      (printf "This expression yields a result of type ~a~%"
        (unparse-type (type-check a initial-type-env)))
      (interp a initial-env))))