;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;  The INFERRED language (based on CHECKED)

; A trick to make DrScheme show the full details of data created using
; constructors created using define-datatype
(require (only mzscheme print-struct))
(print-struct #t)

;;;;;;;;;;;;;;;; procedures ;;;;;;;;;;;;;;;;

; The following code implements a procedural representation of procedures.

(define proc? procedure?)

(define procedure
  (lambda (x e env)
    (lambda (v)
      (value-of e (extend-env x v env)))))

(define apply-procedure
  (lambda (p v)
    (p v)))

; The following code implements a data-structure representation of procedures.
; It can be obtained by defunctionalizing the procedural representation above.
; Both of these implementations provide the functions proc?, procedure, and
; apply-procedure, but the second implementation overwrites the first one.

(define-datatype proc proc?
  (procedure
   (x symbol?)
   (e expression?)
   (env environment?)))

(define apply-procedure
  (lambda (p v)
    (cases proc p
      (procedure (x e env)
        (value-of e (extend-env x v env))))))

;;;;;;;;;;;;;;;; expressed values ;;;;;;;;;;;;;;;;

;;; an expressed value is either a number, a boolean, or a procedure.

(define-datatype expval expval?
  (num-val
    (value number?))
  (bool-val
    (boolean boolean?))
  (proc-val
    (p proc?))
  )

;;; extractors:

(define expval->num
  (lambda (v)
    (cases expval v
      (num-val (num) num)
      (else (expval-extractor-error 'num v)))))

(define expval->bool
  (lambda (v)
    (cases expval v
      (bool-val (bool) bool)
      (else (expval-extractor-error 'bool v)))))

(define expval->proc
  (lambda (v)
    (cases expval v
      (proc-val (p) p)
      (else (expval-extractor-error 'proc v)))))

(define expval-extractor-error
  (lambda (variant value)
    (eopl:error 'expval-extractors "Looking for a ~s, found ~s"
      variant value)))

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

;;; example of a data type built without define-datatype

; Any implementation of environments, as shown in Section 2.2, will do.
; For example, we could have used define-datatype and cases instead:
;
; (define-datatype environment environment?
;   (empty-env)
;   (extend-env
;     (saved-sym symbol?)
;     (saved-val expval?)
;     (saved-env environment?))
;   (extend-env-recursively
;     (p symbol?)
;     (x symbol?)
;     (e expression?)
;     (saved-env environment?))
;   )
;
; (define apply-env
;   (lambda (env sym)
;     (cases environment env
;       (empty-env ()
;        (eopl:error 'apply-env "No binding for ~s" sym))
;       (extend-env (saved-sym saved-val saved-env)
;         (if (eq? sym saved-sym)
;           saved-val
;           (apply-env saved-env sym)))
;       (extend-env-recursively (p x e saved-env)
;         (if (eq? sym p)
;           (proc-val
;            (lambda (v)
;              (value-of e (extend-env x v env))))
;           (apply-env saved-env sym)))
;       )))

(define empty-env
  (lambda ()
    '()))

(define extend-env
  (lambda (saved-sym saved-val saved-env)
    (cond
      ((not (symbol? saved-sym))
       (eopl:error 'extend-env "~s is not a symbol" saved-sym))
      ((not (expval? saved-val))
       (eopl:error 'extend-env "~s is not a value" saved-val))
      ((not (environment? saved-env))
       (eopl:error 'extend-env "~s is not an environment" saved-env))
      (else
       (cons (list saved-sym saved-val) saved-env)))))

(define extend-env-recursively
  (lambda (p x e saved-env)
    (cond
      ((not (symbol? p))
       (eopl:error 'extend-env-recursively "~s is not a symbol" p))
      ((not (symbol? x))
       (eopl:error 'extend-env-recursively "~s is not a symbol" x))
      ((not (expression? e))
       (eopl:error 'extend-env-recursively "~s is not an expression" e))
      ((not (environment? saved-env))
       (eopl:error 'extend-env "~s is not an environment" saved-env))
      (else
       (cons (list p x e) saved-env)))))

(define environment?
  (lambda (x)
    (or (null? x)
        (and (pair? x)
             (list? (car x))
             (or (and (= 2 (length (car x)))
                      (symbol? (car (car x)))
                      (expval? (car (cdr (car x)))))
                 (and (= 3 (length (car x)))
                      (symbol? (car (car x)))
                      (symbol? (car (cdr (car x))))
                      (expression? (car (cdr (cdr (car x)))))))
             (environment? (cdr x))))))

(define apply-env
  (lambda (env sym)
    (cond
      ((null? env)
       (eopl:error 'apply-env "No binding for ~s" sym))
      ((eq? sym (car (car env)))
       (if (null? (cdr (cdr (car env))))
           (car (cdr (car env)))
           (proc-val
             (procedure (car (cdr (car env)))
                        (car (cdr (cdr (car env))))
                        env))))
      (else
       (apply-env (cdr env) sym)))))

; A type environment is like an environment but maps a name to a type
; rather than a value of any sort.  We could implement type environments
; just as we implement environments except replacing expval? by type?,
; but just for the heck of it let's implement type environments using
; define-datatype instead.  There's no type-environment equivalent for
; extend-env-recursively.

(define-datatype type-environment type-environment?
  (empty-tenv)
  (extend-tenv
    (saved-sym symbol?)
    (saved-type type?)
    (saved-tenv type-environment?))
  )

(define apply-tenv
  (lambda (tenv sym)
    (cases type-environment tenv
      (empty-tenv ()
       (eopl:error 'apply-tenv "No binding for ~s" sym))
      (extend-tenv (saved-sym saved-typ saved-tenv)
        (if (eq? sym saved-sym)
          saved-typ
          (apply-tenv saved-tenv sym)))
      )))

;;;;;;;;;;;;;;;; initial environment ;;;;;;;;;;;;;;;;

;; init-env : () -> environment

;; (init-env) builds an environment in which i is bound to the
;; expressed value 1, v is bound to the expressed value 5, and x is
;; bound to the expressed value 10.

(define init-env
  (lambda ()
    (extend-env
     'i (num-val 1)
     (extend-env
      'v (num-val 5)
      (extend-env
       'x (num-val 10)
       (empty-env))))))

;; init-tenv : () -> type-environment

;; (init-tenv) builds a type-environment that corresponds to (init-env).

(define init-tenv
  (lambda ()
    (extend-tenv
     'i (int-type)
     (extend-tenv
      'v (int-type)
      (extend-tenv
       'x (int-type)
       (empty-tenv))))))

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

(define the-lexical-spec
  '((whitespace (whitespace) skip)
    (comment ("%" (arbno (not #\newline))) skip)
    (identifier
     (letter (arbno (or letter digit "_" "-" "?")))
     symbol)
    (number (digit (arbno digit)) number)
    (number ("-" digit (arbno digit)) number)
    ))

;;; Use show-the-datatypes (defined below) to show the define-datatype
;;; automatically generated from this grammar.

(define the-grammar
  '((program (expression) a-program)

    (expression (number) const-exp)

    (expression (identifier) var-exp)

    (expression
     ("-" "(" expression "," expression ")")
     diff-exp)

    (expression
     ("zero?" "(" expression ")")
     zero?-exp)

    (expression
     ("if" expression "then" expression "else" expression)
     if-exp)

    (expression
     ("let" identifier "=" expression "in" expression)
     let-exp)

    (expression
     ("proc" "(" identifier ")" expression)
     proc-exp)

    (expression
     ("(" expression expression ")")
     call-exp)

    (expression
     ("letrec" identifier "(" identifier ")" "=" expression "in" expression)
     letrec-exp)

    ))

;;;;;;;;;;;;;;;; sllgen boilerplate ;;;;;;;;;;;;;;;;

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

(define show-the-datatypes
  (lambda () (sllgen:list-define-datatypes the-lexical-spec the-grammar)))

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

(define just-scan
  (sllgen:make-string-scanner the-lexical-spec the-grammar))

;;;;;;;;;;;;;;;; type checking ;;;;;;;;;;;;;;;;

;; For type inference, we allow an unknown type (marked with a serial number)

(define-datatype type type?
  (int-type)
  (bool-type)
  (proc-type
    (arg-type type?)
    (result-type type?))
  (unknown-type
    (tvar type-variable?)))

(define type-variable? integer?)

(define type->sexp
  (lambda (ty)
    (cases type ty
      (int-type () 'int)
      (bool-type () 'bool)
      (proc-type (arg-type result-type)
        (list (type->sexp arg-type) '-> (type->sexp result-type)))
      (unknown-type (tv) tv))))

;; apply-one-subst: type * type-variable * type -> type
;; (apply-one-subst oldty tv newty) returns the type obtained by
;; substituting newty for every occurrence of (unknown-type tv) in oldty.

(define apply-one-subst
  (lambda (oldty tv newty)
    (cases type oldty
      (int-type () oldty)
      (bool-type () oldty)
      (proc-type (arg-type result-type)
        (proc-type (apply-one-subst arg-type tv newty)
                   (apply-one-subst result-type tv newty)))
      (unknown-type (tv1)
        (if (= tv tv1) newty oldty)))))

;; Many operations on types and substitutions return a type-substitution pair.
;; Let us call such a pair an answer.

(define-datatype answer answer?
  (an-answer (ty type?)
             (subst substitution?)))

;; A substitution is a function from types to types, so it is natural
;; to represent it as a Scheme procedure.  But we also need to keep track
;; of the serial number to use for the next fresh unknown type.

(define-datatype substitution substitution?
  (a-subst (application procedure?)
           (fresh type-variable?)))

; empty-subst : () -> substitution
; (empty-subst) returns an empty substitution, in other words, the
; identity function from types to types).

(define empty-subst
  (lambda ()
    (a-subst (lambda (ty) ty)
             0)))

; extend-subst : type-variable * type * substitution -> substitution

(define extend-subst
  (lambda (saved-tv saved-ty saved-subst)
    (cases substitution saved-subst
      (a-subst (saved-application saved-fresh)
        (a-subst
          (lambda (ty)
            (apply-one-subst (saved-application ty) saved-tv saved-ty))
          saved-fresh)))))

; apply-subst-to-type : substitution * type -> type

(define apply-subst-to-type
  (lambda (subst ty)
    (cases substitution subst
      (a-subst (saved-application saved-fresh)
        (saved-application ty)))))

; fresh-unknown-type : substitution -> answer

(define fresh-unknown-type
  (lambda (subst)
    (cases substitution subst
      (a-subst (saved-application saved-fresh)
        (an-answer (unknown-type saved-fresh)
                   (a-subst saved-application (+ 1 saved-fresh)))))))

;; Unification extends a substitution to make two types equivalent,
;; returning either the new substitution if successful, or #f if not.

;; unify : type * type * substitution -> substitution or #f

(define unify
  (lambda (ty1 ty2 subst)
    (let ((ty1 (apply-subst-to-type subst ty1))
          (ty2 (apply-subst-to-type subst ty2))
          (ext (lambda (tv ty)
                 (if (occurs? tv ty) #f (extend-subst tv ty subst)))))
      (cases type ty1
        (int-type ()
          (cases type ty2
            (int-type () subst)
            (unknown-type (tv2) (ext tv2 ty1))
            (else #f)))
        (bool-type ()
          (cases type ty2
            (bool-type () subst)
            (unknown-type (tv2) (ext tv2 ty1))
            (else #f)))
        (proc-type (ty1a ty1r)
          (cases type ty2
            (proc-type (ty2a ty2r)
              (let ((subst-arg (unify ty1a ty2a subst)))
                (if subst-arg (unify ty1r ty2r subst-arg) #f)))
            (unknown-type (tv2) (ext tv2 ty1))
            (else #f)))
        (unknown-type (tv1)
          (cases type ty2
            (unknown-type (tv2)
              (if (= tv1 tv2) subst (ext tv1 ty2)))
            (else (ext tv1 ty2))))))))

;; occurs? : type-variable * type -> boolean
;; (occurs? tv ty) checks whether tv occurs in ty

(define occurs?
  (lambda (tv ty)
    (cases type ty
      (int-type () #f)
      (bool-type () #f)
      (proc-type (arg-type result-type)
        (or (occurs? tv arg-type)
            (occurs? tv result-type)))
      (unknown-type (tv1) (= tv tv1)))))

;; check-type! : expression * type-environment * substitution * type
;;               -> substitution

; (check-type! exp tenv subst ty) checks that (type-of exp tenv subst)
; and ty can be unified.  If so, it returns an extended substitution.
; If not, it issues an error.

(define check-type!
  (lambda (exp tenv subst expected-type)
    (cases answer (type-of exp tenv subst)
      (an-answer (actual-type subst1)
        (let ((subst2 (unify actual-type expected-type subst1)))
          (if subst2
            subst2
            (eopl:error
              'type-of
              "The expression ~s has the type ~s rather than ~s"
              exp
              (type->sexp (apply-subst-to-type subst1 actual-type))
              (type->sexp (apply-subst-to-type subst1 expected-type)))))))))

;; type-of-program : program -> type

(define type-of-program
  (lambda (pgm)
    (cases program pgm
      (a-program (body)
        (cases answer (type-of body (init-tenv) (empty-subst))
          (an-answer (ty subst)
            (apply-subst-to-type subst ty)))))))

;; type-of : expression * type-environment * substitution -> answer

(define type-of
  (lambda (exp tenv subst)
    (cases expression exp

      (const-exp (n) (an-answer (int-type) subst))

      (var-exp (x) (an-answer (apply-tenv tenv x) subst))

      (diff-exp (e1 e2)
        (let ((subst1 (check-type! e1 tenv subst (int-type))))
          (let ((subst2 (check-type! e2 tenv subst1 (int-type))))
            (an-answer (int-type) subst2))))

      (zero?-exp (e1)
        (let ((subst1 (check-type! e1 tenv subst (int-type))))
          (an-answer (bool-type) subst1)))

      (if-exp (e1 e2 e3)
        (let ((subst1 (check-type! e1 tenv subst (bool-type))))
          (cases answer (type-of e2 tenv subst1)
            (an-answer (ty2 subst2)
              (let ((subst3 (check-type! e3 tenv subst2 ty2)))
                (an-answer ty2 subst3))))))

      (let-exp (x e1 e2)
        (cases answer (type-of e1 tenv subst)
          (an-answer (ty1 subst1)
            (type-of e2 (extend-tenv x ty1 tenv) subst1))))

      (proc-exp (x e)
        (cases answer (fresh-unknown-type subst)
          (an-answer (arg-type subst1)
            (cases answer (type-of e (extend-tenv x arg-type tenv) subst1)
              (an-answer (result-type subst2)
                (an-answer (proc-type arg-type result-type) subst2))))))

      (call-exp (e1 e2)
        (cases answer (fresh-unknown-type subst)
          (an-answer (result-type subst1)
            (cases answer (type-of e2 tenv subst1)
              (an-answer (arg-type subst2)
                (an-answer result-type
                  (check-type! e1 tenv subst2
                               (proc-type arg-type result-type))))))))

      (letrec-exp (p x e1 e2)
        (cases answer (fresh-unknown-type subst)
          (an-answer (ty1 subst1)
            (cases answer (fresh-unknown-type subst1)
              (an-answer (ty subst2)
                (let ((body-tenv (extend-tenv p (proc-type ty1 ty) tenv)))
                  (type-of e2 body-tenv
                    (check-type! e1 (extend-tenv x ty1 body-tenv) subst2
                                 ty))))))))

      )))

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

;; value-of-program : program -> expval

(define value-of-program
  (lambda (pgm)
    (cases program pgm
      (a-program (body)
        (value-of body (init-env))))))

;; value-of : expression * environment -> expval

(define value-of
  (lambda (exp env)
    (cases expression exp

      (const-exp (n) (num-val n))

      (var-exp (x) (apply-env env x))

      (diff-exp (e1 e2)
        (let ((n1 (expval->num (value-of e1 env)))
              (n2 (expval->num (value-of e2 env))))
          (num-val (- n1 n2))))

      (zero?-exp (e1)
        (let ((n1 (expval->num (value-of e1 env))))
          (if (zero? n1)
            (bool-val #t)
            (bool-val #f))))

      (if-exp (e1 e2 e3)
        (if (expval->bool (value-of e1 env))
          (value-of e2 env)
          (value-of e3 env)))

      (let-exp (x e1 e2)
        (let ((v1 (value-of e1 env)))
          (value-of e2
            (extend-env x v1 env))))

      (proc-exp (x e)
        (proc-val (procedure x e env)))

      (call-exp (e1 e2)
        (let ((p (expval->proc (value-of e1 env))))
          (apply-procedure p (value-of e2 env))))

      (letrec-exp (p x e1 e2)
        (value-of e2
          (extend-env-recursively p x e1 env)))

      )))

;;;;;;;;;;;;;;;; testing ;;;;;;;;;;;;;;;;

;; run : string -> expval

(define run
  (lambda (string)
    (let ((program (scan&parse string)))
      (begin (type-of-program program)
             (value-of-program program)))))

; A test is a list with three elements: the name of the test (a symbol),
; the program to test (a string), and the desired result (a number or a
; boolean).  The desired result may also be the symbol error, but for
; now we just skip those tests.

(define test->name
  (lambda (test)
    (car test)))

(define test->program
  (lambda (test)
    (car (cdr test))))

(define test->result
  (lambda (test)
    (car (cdr (cdr test)))))

(define result->expval
  (lambda (result)
    (cond
      ((number? result)
       (num-val result))
      ((boolean? result)
       (bool-val result))
      (else
       (eopl:error 'result->expval "Can't convert ~s to expval" result)))))

;; test-one : test -> boolean
; The test-one function returns whether the test succeeded.

(define test-one
  (lambda (test)
    (let ((result (test->result test)))
      (or (eq? result 'error)
          (equal? (run (test->program test))
                  (result->expval result))))))

;; test-list : (list-of test) -> (list-of sym)
; The test-list function returns the names of the tests that failed.

(define test-list
  (lambda (tests)
    (cond
      ((null? tests)
       '())
      ((test-one (car tests))
       (test-list (cdr tests)))
      (else
       (cons (test->name (car tests))
             (test-list (cdr tests)))))))

;; test : () -> (list-of sym)

(define test
  (lambda ()
    (test-list
      '(

        ;; simple arithmetic
        (positive-const "11" 11)
        (negative-const "-33" -33)
        (simple-arith-1 "-(44,33)" 11)

        ;; nested arithmetic
        (nested-arith-left "-(-(44,33),22)" -11)
        (nested-arith-right "-(55, -(22,11))" 44)

        ;; simple variables
        (test-var-1 "x" 10)
        (test-var-2 "-(x,1)" 9)
        (test-var-3 "-(1,x)" -9)

        ;; simple unbound variables
        (test-unbound-var-1 "foo" error)
        (test-unbound-var-2 "-(x,foo)" error)

        ;; simple conditionals
        (if-true "if zero?(0) then 3 else 4" 3)
        (if-false "if zero?(1) then 3 else 4" 4)

        ;; test dynamic typechecking
        (no-bool-to-diff-1 "-(zero?(0),1)" error)
        (no-bool-to-diff-2 "-(1,zero?(0))" error)
        (no-int-to-if "if 1 then 2 else 3" error)

        ;; make sure that the test and both arms get evaluated
        ;; properly.
        (if-eval-test-true "if zero?(-(11,11)) then 3 else 4" 3)
        (if-eval-test-false "if zero?(-(11, 12)) then 3 else 4" 4)

        ;; and make sure the other arm gets type-checked too
        (if-eval-test-true-2 "if zero?(-(11, 11)) then 3 else foo" error)
        (if-eval-test-false-2 "if zero?(-(11,12)) then foo else 4" error)

        ;; simple let
        (simple-let-1 "let x = 3 in x" 3)

        ;; make sure the body and rhs get evaluated
        (eval-let-body "let x = 3 in -(x,1)" 2)
        (eval-let-rhs "let x = -(4,1) in -(x,1)" 2)

        ;; check nested let and shadowing
        (simple-nested-let "let x = 3 in let y = 4 in -(x,y)" -1)
        (check-shadowing-in-body "let x = 3 in let x = 4 in x" 4)
        (check-shadowing-in-rhs "let x = 3 in let x = -(x,1) in x" 2)

        ;; test procedures
        (plus "let plus = proc (x) proc (y) -(x,-(0,y)) in ((plus 3) 4)" 7)
        (shadow "let i = 3
                 in let add-i = proc (j) -(i,-(0,j))
                    in let i = 8
                       in (add-i i)"
                11)
        (lexical-scope-1 "let i = 3
                          in let add-i = proc (j) -(i,-(0,j))
                             in let i = 8
                                in (add-i 0)"
                         3)
        (lexical-scope-2 "let plus = proc (x) proc (y) -(x,-(0,y))
                          in let x = (plus 3)
                             in (x 4)"
                         7)

        ;; test recursive procedures
        (letrec "letrec parity(x) = if zero?(x)
                                    then 0
                                    else -(1,(parity -(x,1)))
                 in (parity 13)"
                1)

        ))))