;; We start with a simply typed proc language
;; Then we add a `lexvar' form, which is not used by
;;  source programs.
;; Finally, we implement a program transformer that converts
;;  identifier expressions to lexvar expressions.

;;;;;;;;;;;;;;;; top level interface ;;;;;;;;;;;;;;;;

(define type-check
  (lambda (string)
    (type-to-external-form
      (type-of-program
        (scan&parse string)))))

(define run
  (lambda (string)
    (eval-program (scan&parse string))))

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

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

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

    (expression (number) lit-exp)
    (expression ("true") true-exp)
    (expression ("false") false-exp)
    (expression (identifier) var-exp)

    ;; >>- Here's the new lexvar form: --------------<<<<
    (expression ("lexvar" number number) lexvar-exp)

    (expression
      (primitive "(" (separated-list expression ",") ")")
      primapp-exp)
    (expression
      ("if" expression 
       "then" expression 
       "else" expression)
      if-exp)
   (expression
      ("let" (arbno  identifier "=" expression) 
       "in" expression)
      let-exp)
   (expression
      ("(" expression (arbno expression) ")")
      app-exp)
    
   (expression
     ("proc" 
      "(" (separated-list type-exp identifier ",") ")"
      expression)
      proc-exp)
   (expression
     ("letrec"
      (arbno 
       type-exp identifier
       "(" (separated-list type-exp identifier ",") ")"
       "=" expression) "in" expression) 
     letrec-exp)

    (primitive ("+")     add-prim)
    (primitive ("-")     subtract-prim)
    (primitive ("*")     mult-prim)
    (primitive ("add1")  incr-prim)
    (primitive ("sub1")  decr-prim)
    (primitive ("zero?") zero-test-prim)
    
    (type-exp ("int") int-type-exp)
    (type-exp ("bool") bool-type-exp)
    (type-exp
     ("(" (separated-list type-exp "*") "->" type-exp ")")
     proc-type-exp)
    ))

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


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

;;;;;;;;;;;;;;;; The Type Checker ;;;;;;;;;;;;;;;;

; type-of-program : program -> type
(define type-of-program
  (lambda (pgm)
    (cases program pgm
      (a-program (exp)
                 (type-of-expression exp (empty-tenv))))))

; type-of-expression : expr type-env -> type
(define type-of-expression
  (lambda (exp tenv)
    (cases expression exp
      (lit-exp (number) int-type)
      (true-exp () bool-type)
      (false-exp () bool-type)
      (var-exp (id) (apply-tenv tenv id))
      (if-exp (test-exp then-exp else-exp)
              (let ((test-type
                     (type-of-expression test-exp tenv))
                    (then-type
                     (type-of-expression then-exp tenv))
                    (else-type
                     (type-of-expression else-exp tenv)))
                ; these tests either succeed
                ; or raise an error
                (check-equal-type! test-type
                                   bool-type
                                   test-exp)
                (check-equal-type! then-type
                                   else-type
                                   exp)
                ; then-type and else-type are the same
                then-type))
      (proc-exp (texps ids body)
        (type-of-proc-exp texps ids body tenv))
      (primapp-exp (prim rands)
        (type-of-application
          (type-of-primitive prim)
          (types-of-expressions rands tenv)
          prim rands exp))
      (app-exp (rator rands) 
        (type-of-application
          (type-of-expression rator tenv)
          (types-of-expressions rands tenv)
          rator rands exp))
      (let-exp (ids rands body)
        (type-of-let-exp ids rands body tenv))
      (letrec-exp (result-texps proc-names
                                texpss idss bodies
                                letrec-body)
        (type-of-letrec-exp
          result-texps proc-names texpss idss bodies
          letrec-body tenv))
      (else (eopl:error 'type-of-expression
			"form not allowed in source programs: ~a"
			exp)))))

; check-equal-type! : type type 
;                     expr ; for err msg
;                     -> 
(define check-equal-type!
  (lambda (t1 t2 exp)
    (or (equal? t1 t2) 
        (eopl:error 'check-equal-type!
          "Types didn't match: ~s != ~s in~%~s"
          (type-to-external-form t1)
          (type-to-external-form t2)
          exp))))

; type-of-proc-exp : lof-type-expr lof-sym exp type-env
;                    -> type
(define type-of-proc-exp
  (lambda (texps ids body tenv)
    (let ((arg-types (expand-type-expressions texps)))
      (let ((result-type
              (type-of-expression body
                (extend-tenv ids arg-types tenv))))
        (proc-type arg-types result-type)))))

; type-of-application : type lof-type 
;                       expr lof-expr expr ; for err msgs
;                       -> type
(define type-of-application
  (lambda (rator-type rand-types rator rands exp)
    (cases type rator-type
      (proc-type (arg-types result-type)
        (if (= (length arg-types) (length rand-types))
          (begin 
            (for-each
              check-equal-type!
              rand-types arg-types rands)
            result-type)
          (eopl:error 'type-of-expression
            (string-append
              "Wrong num of arguments in expression ~s:"
              "~%expected ~s~%got ~s")
            exp
            (map type-to-external-form arg-types)
            (map type-to-external-form rand-types))))
      (else
        (eopl:error 'type-of-expression
          "Rator not a proc type:~%~s~%had type ~s"   
          rator (type-to-external-form rator-type))))))

; types-of-expressions : lof-expr type-env -> lof-type
(define types-of-expressions
  (lambda (rands tenv)
    (map (lambda (exp) (type-of-expression exp tenv)) 
         rands)))

; type-of-let-exp : lof-sym lof-expr expr type-env -> type
(define type-of-let-exp
  (lambda (ids rands body tenv)
    (let ((tenv-for-rands
            (extend-tenv
              ids
              (types-of-expressions rands tenv)
              tenv)))
      (type-of-expression body tenv-for-rands))))
          
; type-of-letrec-exp : lof-type-expr lof-sym
;                      lof-lof-type-expr lof-lof-sym
;                      lof-expr expr type-env -> type
(define type-of-letrec-exp
  (lambda (result-texps proc-names texpss idss bodies
            letrec-body tenv) 
    (let ((arg-typess
            (map
              (lambda (texps)
                (expand-type-expressions texps))
              texpss))
          (result-types
            (expand-type-expressions result-texps)))
      (let ((the-proc-types
             (map proc-type arg-typess result-types)))
        (let ((tenv-for-body ; type env for proc-bodies
                (extend-tenv proc-names 
                             the-proc-types tenv)))
          (for-each                     
            (lambda (ids arg-types body result-type)
              (check-equal-type!
                (type-of-expression 
                  body
                  (extend-tenv ids arg-types 
                               tenv-for-body))
                result-type
                body))
            idss arg-typess bodies result-types)
          (type-of-expression letrec-body 
                              tenv-for-body))))))

;;;;;;;;;;;;;;;; types ;;;;;;;;;;;;;;;;

;; There is a 1-to-1 correspondance
;; between type expressions and types

(define-datatype type type?
  (atomic-type
    (name symbol?))
  (proc-type
    (arg-types (list-of type?))
    (result-type type?)))

; expand-type-expression : type-expr -> type
(define expand-type-expression
  (lambda (texp)
    (cases type-exp texp
      (int-type-exp () int-type)
      (bool-type-exp () bool-type)
      (proc-type-exp (arg-texps result-texp)
        (proc-type
          (expand-type-expressions arg-texps)
          (expand-type-expression result-texp))))))

; expand-type-expressions : lof-type-expr -> lof-type
(define expand-type-expressions
  (lambda (texps)
    (map expand-type-expression texps)))

;;; types of primitives

(define int-type (atomic-type 'int))
(define bool-type (atomic-type 'bool))

(define type-of-primitive
  (lambda (prim)
    (cases primitive prim
      (add-prim ()
        (proc-type (list int-type int-type) int-type))
      (subtract-prim ()
        (proc-type (list int-type int-type) int-type))
      (mult-prim ()
        (proc-type (list int-type int-type) int-type))
      (incr-prim ()
        (proc-type (list int-type) int-type))
      (decr-prim ()
        (proc-type (list int-type) int-type))
      (zero-test-prim ()
        (proc-type (list int-type) bool-type)))))


;;;;;;;;;;;;;;;; type environments ;;;;;;;;;;;;;;;;

(define-datatype type-environment type-environment?
  (empty-tenv-record)
  (extended-tenv-record
    (syms (list-of symbol?))
    (vals (list-of type?))
    (tenv type-environment?)))

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

(define apply-tenv 
  (lambda (tenv sym)
    (cases type-environment tenv
      (empty-tenv-record ()
        (eopl:error 'apply-tenv "Unbound variable ~s"
                    sym))
      (extended-tenv-record (syms vals env)
        (let ((pos (list-find-position sym syms)))
          (if (number? pos)
            (list-ref vals pos)
            (apply-tenv env sym)))))))
    
; apply-tenv-lexvar-helper : sym tenv num -> expr
(define apply-tenv-lexvar-helper
  (lambda (tenv sym depth)
    (cases type-environment tenv
      (empty-tenv-record ()
        (eopl:error 'apply-tenv-lexvar
                    "Unbound variable ~s"
                    sym))
      (extended-tenv-record (syms vals env)
        (let ((pos (list-find-position sym syms)))
          (if (number? pos)
            (lexvar-exp depth pos)
            (apply-tenv-lexvar-helper 
             env sym (+ depth 1))))))))

; apply-tenv-lexvar : sym tenv -> expr
(define apply-tenv-lexvar
  (lambda (tenv sym)
    (apply-tenv-lexvar-helper tenv sym 0)))

;;;;;;;;;;;;;;;; external form of types ;;;;;;;;;;;;;;;;

(define type-to-external-form
  (lambda (ty)
    (cases type ty
      (atomic-type (name) name)
      (proc-type (arg-types result-type)
        (append
          (arg-types-to-external-form arg-types)
          '(->)
          (list (type-to-external-form result-type)))))))

(define arg-types-to-external-form
  (lambda (types)
    (if (null? types)
      '()
      (if (null? (cdr types))
        (list (type-to-external-form (car types)))
        (cons
          (type-to-external-form (car types))
          (cons '*
            (arg-types-to-external-form (cdr types))))))))

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

;; Plain old call-by-value interpreter.
;; We implement booleans with 0 and 1 (but
;; there's no possibility of confusion with 
;; numbers, thanks to type-checking)

(define eval-program 
  (lambda (pgm)
    (cases program pgm
      (a-program (body)
        (eval-expression body (empty-env))))))

(define eval-expression
  (lambda (exp env)
    (cases expression exp
      (lit-exp (datum) datum)
      (true-exp () 1)
      (false-exp () 0)
      (var-exp (id) (apply-env env id))
      ;; >>- Here's the new lexvar form: --------------<<<<
      (lexvar-exp (depth pos) 
		  ;; All the work is in apply-env-lexvar
                  (apply-env-lexvar env depth pos))
      (primapp-exp (prim rands)
	(let ((args (eval-primapp-exp-rands rands env)))
	  (apply-primitive prim args)))
      (if-exp (test-exp true-exp false-exp)
        (if (true-value? (eval-expression test-exp env))
          (eval-expression true-exp env)
          (eval-expression false-exp env)))
      (let-exp (ids rands body)
	(let ((args (eval-rands rands env)))
          (eval-expression body 
                           (extend-env ids args env))))
      (proc-exp (texps ids body)
        (closure ids body env))
      (app-exp (rator rands)
        (let ((proc (eval-expression  rator env))
              (args (eval-rands rands env)))
          (if (procval? proc) ; should always be true in
                              ; typechecked code
            (apply-procval proc args)
            (eopl:error 'eval-expression 
              "Attempt to apply non-procedure ~s" proc))))
      (letrec-exp (result-texps proc-names 
                                texpss idss bodies
                                letrec-body)
        (eval-expression letrec-body
          (extend-env-recursively proc-names idss bodies
                                  env)))
      )))


(define eval-primapp-exp-rands
  (lambda (rands env)
    (map (lambda (x) (eval-expression x env)) rands)))

(define eval-rands
  (lambda (rands env)
    (map (lambda (x) (eval-rand x env)) rands)))

(define eval-rand
  (lambda (rand env)
    (eval-expression rand env)))

(define apply-primitive
  (lambda (prim args)
    (cases primitive prim
      (add-prim  () (+ (car args) (cadr args)))
      (subtract-prim () (- (car args) (cadr args)))
      (mult-prim  () (* (car args) (cadr args)))
      (incr-prim  () (+ (car args) 1))
      (decr-prim  () (- (car args) 1))
      (zero-test-prim () (if (zero? (car args)) 1 0))
      )))

;;;;;;;;;;;;;;;; booleans ;;;;;;;;;;;;;;;;

(define true-value?
  (lambda (x)
    (not (zero? x))))

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

(define-datatype procval procval?
  (closure 
    (ids (list-of symbol?)) 
    (body expression?)
    (env environment?)))

(define apply-procval
  (lambda (proc args)
    (cases procval proc
      (closure (ids body env)
        (eval-expression body 
                         (extend-env ids args env))))))

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

(define-datatype environment environment? 
  (empty-env-record)             
  (extended-env-record
    (syms (list-of symbol?))
    (vals vector?)
    (env environment?)))

(define apply-env
  (lambda (env sym)
    (cases environment env
      (empty-env-record ()
        (eopl:error 'empty-env "No binding for ~s" sym))
      (extended-env-record (syms vals old-env)
        (let ((pos (env-find-position sym syms)))
          (if (number? pos)
            (vector-ref vals pos)
            (apply-env old-env sym)))))))

;; apply-env-lexvar : env num num -> expval           <<<<
;;  Like apply-env, but takes a lexical address instead of
;;  a symbol
(define apply-env-lexvar
  (lambda (env depth pos)
    (cases environment env
      (empty-env-record ()
        (eopl:error 'empty-env "No binding for ~s" sym))
      (extended-env-record (syms vals old-env)
	  ;; No need to search syms!
          (if (zero? depth)
            (vector-ref vals pos)
            (apply-env-lexvar old-env (- depth 1)
                              pos))))))

(define empty-env
  (lambda ()
    (empty-env-record)))

(define extend-env
  (lambda (syms vals env)
    (extended-env-record syms (list->vector vals) env)))
 
(define extend-env-recursively
  (lambda (proc-names idss bodies old-env)
    (let ((len (length proc-names)))
      (let ((vec (make-vector len)))
        (let ((env (extended-env-record proc-names vec
                                        old-env)))
          (for-each
            (lambda (pos ids body)
              (vector-set! vec pos
                           (closure ids body env)))
            (iota len) idss bodies)
          env)))))

(define env-find-position
  (lambda (sym los)
    (list-find-position sym los)))

(define list-find-position
  (lambda (sym los)
    (list-index (lambda (sym1) (eqv? sym1 sym)) los)))

(define list-index
  (lambda (pred ls)
    (cond
      ((null? ls) #f)
      ((pred (car ls)) 0)
      (else (let ((list-index-r (list-index pred 
                                            (cdr ls))))
              (if (number? list-index-r)
                (+ list-index-r 1)
                #f))))))

(define iota
  (lambda (end)
    (let loop ((next 0))
      (if (>= next end) '()
        (cons next (loop (+ 1 next)))))))

;;;;;;;;;;;;;;;;;;;; Compiler ;;;;;;;;;;;;;;;;;;;;

;; >>- Here's the compiler -----------------------<<<<

;; The compiler reuses much of the type-checking
;; infrastructure, including type environments.
;; Type environments have the same shape as
;; run-tie environments, so we can use them to
;; compute lexical addresses.

; translate-program : program -> program
(define (translate-program p)
  (cases program p
    (a-program (expr)
               (a-program
                (translate-expr expr (empty-tenv))))))

; translate-expr : expr tenv -> expr
(define (translate-expr expr tenv)
  (cases expression expr
    ;; Convert a varable to a lexvar expression
    (var-exp (id) (apply-tenv-lexvar tenv id))
    ;; For most expression forms, we merely process
    ;;  the subexpressions recursively
    (lit-exp (v) expr)
    (lexvar-exp (depth pos) expr)
    (true-exp () expr)
    (false-exp () expr)
    (if-exp (test then else)
            (if-exp
             (translate-expr test tenv)
             (translate-expr then tenv)
             (translate-expr else tenv)))
    (primapp-exp (prim rands)
                 (primapp-exp
                  prim  ; primitive
                  (translate-rands rands tenv)))
    (app-exp (rator rands)
             (app-exp
              (translate-expr rator tenv) ; expr
              (translate-rands rands tenv))) ; lstof-expr
    (let-exp (ids rands body-expr)
             (let-exp
              ids
              (translate-rands rands tenv) ; lstof-expr
              (translate-expr 
               body-expr  ; expr 
               (extend-tenv
                ids
                (types-of-expressions rands tenv)
                tenv))))
    (letrec-exp (result-texps proc-names 
                              texpss idss bodies
                              letrec-body)
                (translate-letrec-expr
                 result-texps proc-names 
                 texpss idss bodies
                 letrec-body tenv))
    (proc-exp (texprs ids body-expr)
              (proc-exp
               texprs ; lstof-type-expr
               ids    ; lstof-sym
               (translate-expr
                body-expr ; expr
                (extend-tenv
                 ids
                 (expand-type-expressions texprs)
                 tenv))))))

;; translate-letrec-expr : ... -> expr
;;  Translates a letrec expression. There's a lot of
;;  work to do because we have to build an appropriate
;;  type environment.
(define translate-letrec-expr
  (lambda (result-texps proc-names texpss idss bodies
            letrec-body tenv) 
           ;; The types for the procedure arguments
    (let ((arg-typess ; lstof-lstof-type
            (map
              (lambda (texps)
                (expand-type-expressions texps))
              texpss))
	  ;; The types for the procedure results
          (result-types ; lstof-type
            (expand-type-expressions result-texps)))
            ;; The complete types for the procedures
      (let ((the-proc-types ; lstof-type
             (map proc-type arg-typess result-types)))
	      ;; An environment with bindings for the procedurezs:
        (let ((tenv-for-body ; type env for proc-bodies
                (extend-tenv proc-names 
                             the-proc-types tenv)))
          (letrec-exp
           result-texps
           proc-names
           texpss
           idss
           (map
	    ;; For every procedure bound by the
	    ;;  letrec, translate the body, 
	    ;;  starting with tenv-for-body
	    ;;  and adding the parameter types.
            (lambda (body ids types)
              (translate-expr 
               body
               (extend-tenv
                ids
                types
                tenv-for-body)))
            bodies idss arg-typess)
	   ;; Also translate the letrec body
           (translate-expr
            letrec-body tenv-for-body)))))))

; translate-rands : lstof-expr tenv -> lstof-expr
(define (translate-rands exprs tenv)
  (map (lambda (expr) (translate-expr expr tenv))
       exprs))

;;;;;;;;;;;;;;;;;;;; Tests ;;;;;;;;;;;;;;;;;;;;

(define sum-program
  (scan&parse
   "letrec int sum(int n) =
                 if zero?(n)
                    then 1
                    else +(n, (sum -(n, 1)))
      in (sum 100)"))