;; An expval is 
;;   * number, or
;;   * closure

;; A denval is an expval

;; The language:
;;  <expr> = <number>                        ;; a number
;;         | @<number>                       ;; looks up a value in the env
;;         | [<expr>]                        ;; creates a closure (1 arg)
;;         | call(<expr>, <expr>)            ;; applies a closure (1 arg)
;;         | ifzero(<expr>, <expr>, <expr>)  ;; branch
;;         | +(<expr>, <expr>)               ;; add
;;         | -(<expr>, <expr>)               ;; subtract

;;;;;;;; grammatical specification ;;;;;;;;;;;;;;;;

(define the-lexical-spec
  '((whitespace (whitespace) 
                skip)
    (id (letter (arbno (or letter digit "?")))
        make-symbol)
    (number ((or "" "-" "+") digit (arbno digit))
            make-number)))

(define the-grammar
  '((expression ("*") black-hole-exp)
    (expression (number) lit-exp)
    (expression ("@" number) lookup-exp)
    (expression ("[" expression "]") proc-exp)
    (expression ("call" "(" expression "," expression ")") app-exp)
    (expression ("ifzero" 
		 "(" expression 
		 "," expression 
		 "," expression ")") 
		ifzero-exp)
    (expression ("+" "(" expression "," expression ")") add-exp)
    (expression ("-" "(" expression "," expression ")") subtract-exp)))

(sllgen:make-define-datatypes the-lexical-spec
                              the-grammar)

(define-datatype
  closure closure?
  (a-closure (env environment?)
	     (body expression?)))

(define-datatype type type?
  (number-type)
  (proc-type
    (arg-type type?)
    (result-type type?))
  (tvar-type
    (serial-number integer?)
    (container vector?)))

;;;;;;;;;;;;;;;; top level ;;;;;;;;;;;;;;;;

; typecheck : string -> val
;
; Takes a string respresting concrete syntax and
; infers types while checking it. If it checks, a value
; is returned representing the type, otherwise an error
; is reported.
;
;  (typecheck "5") = _num_
;  (typecheck "[1]") = _(tvar17 -> num)_
;  (typecheck "call([@0], 1)") = _num_
;
(define typecheck
  (lambda (string)
    (type-to-external-form (type-of-expression (scan&parse string) (empty-tenv)))))

; read-eval-print : -> [loops forever]
(define read-eval-print
  (lambda ()
    ((sllgen:make-rep-loop "-->"
			   (lambda (x) (eval-expression x (empty-env)))
                           (sllgen:make-stream-parser
                            the-lexical-spec
                            the-grammar)))))

; run : string -> expval
;
; Takes a string respresting concrete syntax and
; evaluates it, returning the result.
;
;  (run "5") = _5_
;  (run "[1]") = _(a-closure (empty-env) (lit-exp 1))_
;  (run "call([@0], 1)") = _1_
;
(define run
  (lambda (string)
    (eval-expression (scan&parse string) (empty-env))))

(define scan&parse
  (sllgen:make-string-parser the-lexical-spec the-grammar))

;;;;;;;;;;;;;;;; the checker ;;;;;;;;;;;;;;;;

; type-of-expression : expression tenv -> type
(define (type-of-expression exp tenv)
  (cases expression exp
    (black-hole-exp () (make-tvar))
    (lit-exp (sym) (number-type))
    (lookup-exp (cn) (apply-tenv tenv cn))
    (proc-exp (body)
	      (let ((arg-type (make-tvar)))
		(let ((result-type (type-of-expression 
				    body
				    (extend-tenv arg-type tenv))))
		  (proc-type arg-type result-type))))
    (app-exp (rator rand)
	     (let ([rator-type (type-of-expression rator tenv)]
		   [rand-type (type-of-expression rand tenv)])
	       (let ((result-type (make-tvar)))
		 (check-equal-type!
		  rator-type
		  (proc-type rand-type result-type)
		  exp)
		 result-type)))
    (ifzero-exp (test then else)
		(check-equal-type! (number-type) 
                                   (type-of-expression test tenv) 
                                   exp)
                (let ([then-type (type-of-expression then tenv)]
                      [else-type (type-of-expression else tenv)])
                  (check-equal-type! then-type else-type exp)
                  then-type))
    (add-exp (a b) 
	     (check-equal-type! (number-type) (type-of-expression a tenv) exp)
	     (check-equal-type! (number-type) (type-of-expression b tenv) exp)
	     (number-type))
    (subtract-exp (a b)
		  (check-equal-type! (number-type) (type-of-expression a tenv) 
                                     exp)
		  (check-equal-type! (number-type) (type-of-expression b tenv)
                                     exp)
		  (number-type))))

;;;;;;;;;;;;;;;; type helper functions ;;;;;;;;;;;;;;;;

(define make-tvar
  (let ((serial-number 0))
    (lambda ()
      (set! serial-number (+ 1 serial-number))
      (tvar-type serial-number (vector '() )))))

(define number-type?
  (lambda (ty)
    (cases type ty
      (number-type () #t)
      (else #f))))

(define proc-type?
  (lambda (ty)
    (cases type ty
      (proc-type (arg-types result-type) #t)
      (else #f))))

(define tvar-type?
  (lambda (ty)
    (cases type ty
      (tvar-type (sn cont) #t)
      (else #f))))

(define proc-type->arg-type
  (lambda (ty)
    (cases type ty
      (proc-type (arg-type result-type) arg-type)
      (else (eopl:error 'proc-type->arg-type
              "Not a proc type: ~s" ty)))))

(define proc-type->result-type
  (lambda (ty)
    (cases type ty
      (proc-type (arg-types result-type) result-type)
      (else (eopl:error 'proc-type->result-type
              "Not a proc type: ~s" ty)))))

(define tvar-type->serial-number
  (lambda (ty)
    (cases type ty
      (tvar-type (sn c) sn)
      (else (eopl:error 'tvar-type->serial-number
              "Not a tvar-type: ~s" ty)))))

(define tvar-type->container
  (lambda (ty)
    (cases type ty
      (tvar-type (sn vec) vec)
      (else (eopl:error 'tvar-type->container
              "Not a tvar-type: ~s" ty)))))

(define tvar->contents
  (lambda (ty)
    (vector-ref (tvar-type->container ty) 0)))

(define tvar-set-contents!
  (lambda (ty val)
    (vector-set! (tvar-type->container ty) 0 val)))

(define tvar-non-empty?
  (lambda (ty)
    (not (null? 
          (vector-ref (tvar-type->container ty) 0)))))

(define type-to-external-form
  (lambda (ty)
    (cases type ty
      (number-type () 'num)
      (proc-type (arg-type result-type)
		 (list
		  (type-to-external-form arg-type)
		  '->
		  (type-to-external-form result-type)))
      (tvar-type (serial-number container) 
		 (if (tvar-non-empty? ty)
		     (type-to-external-form (tvar->contents ty))
		     (string->symbol
		      (string-append
		       "tvar"
		       (number->string serial-number))))))))


;;;;;;;;;;;;;;;; type equivalence ;;;;;;;;;;;;;;;;

(define check-equal-type!               
  (lambda (t1 t2 exp)
    (cond
      ((eqv? t1 t2) 'ok)
      ((tvar-type? t1) (check-tvar-equal-type! t1 t2 exp))
      ((tvar-type? t2) (check-tvar-equal-type! t2 t1 exp))
      ((and (number-type? t1) (number-type? t2)) #t)
      ((and (proc-type? t1) (proc-type? t2))
       (let ((arg-type1 (proc-type->arg-type t1))
             (arg-type2 (proc-type->arg-type t2))
             (result-type1 (proc-type->result-type t1))
             (result-type2 (proc-type->result-type t2)))
	 (check-equal-type! arg-type1 arg-type2 exp)
         (check-equal-type! result-type1 result-type2 exp)))
      (else (raise-type-error t1 t2 exp)))))

(define check-tvar-equal-type!
  (lambda (tvar ty exp)
    (if (tvar-non-empty? tvar)
	;; Already set:
        (check-equal-type! (tvar->contents tvar) ty exp)
        ;; Set tvar to ty --- but only if tvar is not in ty
        (begin
          (check-no-occurrence! tvar ty exp)
          (tvar-set-contents! tvar ty)))))

(define check-no-occurrence!
  (lambda (tvar ty exp)
    (letrec
      ((loop
         (lambda (ty1)
           (cases type ty1
             (number-type () 'ok)
             (proc-type (arg-type result-type)
               (begin
                 (loop arg-type)
                 (loop result-type)))
             (tvar-type (num vec)
               (if (eqv? tvar ty1)
		   (raise-occurrence-check tvar ty exp)))))))
      (loop ty))))

(define raise-type-error
  (lambda (t1 t2 exp)
    (eopl:error 'check-equal-type!
      "Type mismatch: ~s doesn't match ~s in ~s~%"
      (type-to-external-form t1)
      (type-to-external-form t2)
      exp)))

(define raise-occurrence-check
  (lambda (tvnum t2 exp)
    (eopl:error 
     'check-equal-type!
     "Can't unify: ~s occurs in type ~s in expr ~s~%" 
      tvnum
      (type-to-external-form t2)
      exp)))

;;;;;;;;;;;;;;;; type environment ;;;;;;;;;;;;;;;;

(define-datatype type-environment type-environment?
  (empty-tenv-record)
  (extended-tenv-record
   (type type?)
   (tenv type-environment?)))

(define empty-tenv empty-tenv-record)
(define extend-tenv extended-tenv-record)

(define apply-tenv 
  (lambda (tenv depth)
    (cases type-environment tenv
      (empty-tenv-record ()
        (eopl:error 'apply-tenv "Too deep"))
      (extended-tenv-record (type tenv)
          (if (zero? depth)
            type
            (apply-tenv tenv (- depth 1)))))))

;;;;;;;;;;;;;;;; the interpreter ;;;;;;;;;;;;;;;;

; eval-expression : expression env -> expval
;
; Evaluates an expression in the given environment.
;
;  (eval-expression (lit-exp 5) (empty-env)) = 5
;  (eval-expression (lookup-exp 1) 
;                   (extend-env 42 (extend-env 43 (empty-env)))) = 43
;  (eval-expression (app-exp (proc-exp (lookup-exp 0)) 
;                            (lit-exp 42)) 
;                   (empty-env)) = 42
;
(define eval-expression 
  (lambda (exp env)
    (cases expression exp
      (black-hole-exp () (black-hole))
      (lit-exp (sym) sym)
      (lookup-exp (cn) (apply-env env cn))
      (proc-exp (body) (a-closure env body))
      (app-exp (rator rand)
	       (let ([func (eval-expression rator env)]
		     [arg (eval-expression rand env)])
		 (apply-proc func arg)))
      (ifzero-exp (test then else)
		  (if (zero? (eval-expression test env))
		      (eval-expression then env)
		      (eval-expression else env)))
      (add-exp (a b) (+ (eval-expression a env)
			(eval-expression b env)))
      (subtract-exp (a b) (- (eval-expression a env)
			     (eval-expression b env))))))

; apply-proc : proc expval -> expval
(define (apply-proc func arg)
  (cases closure func
	 (a-closure (env body)
		    (eval-expression
		     body
		     (extend-env
		      arg
		      env)))))

(define (black-hole)
  (black-hole))

;;;;;;;;;;;;;;;; environments ;;;;;;;;;;;;;;;;

;; Abstract envrionment datatype implementation.

(define (denval? x) 
  (or (number? x) (closure? x)))

(define-datatype environment environment?
  (empty-env-rec)
  (extended-env-rec
    (val denval?)
    (env environment?)))

; empty-env : -> env
(define empty-env
  (lambda ()
    (empty-env-rec)))

; extend-env : denval env -> env 
(define extend-env
  (lambda (val env)
    (extended-env-rec val env)))

; apply-env : env num -> denval
(define apply-env
  (lambda (env cn)
    (cases environment env
      (empty-env-rec ()
        (eopl:error 'apply-env "lookup too deep"))
      (extended-env-rec (val env)
	(if (zero? cn)
            val
            (apply-env env (- cn 1)))))))