;; A expval is
;;   * number, or
;;   * boolean, or
;;   * proc, or
;;   * reference

;; A denval is a location

;;;;;;;; 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
  '((program (expression) a-program)

    (var (id) cbv-var)
    (var ("&" id) cbr-var)
    
    (expression ("proc"
                 "(" (separated-list var ",") ")"
                 expression)
                proc-exp)
    (expression (number) lit-exp)
    (expression (id) var-exp)
    (expression
      (primitive "(" (separated-list expression ",") ")")
      primapp-exp)
    (expression ("true") true-exp)
    (expression ("false") false-exp)
    (expression ("let"
                 (arbno id "=" expression)
                 "in"
                 expression)
                let-exp)
    (expression ("(" expression (arbno expression) ")")
                app-exp)
    (expression ("if" expression
                 "then" expression
                 "else" expression)
                if-exp)    
    (expression ("set" id "=" expression)
                set-exp)
    (expression ("{" expression ";" expression "}")
                seq-exp)
    (expression ("ref" "(" id ")") ref-exp)

    (primitive ("setref") setref-prim)
    (primitive ("+")     add-prim)
    (primitive ("-")     subtract-prim)
    (primitive ("*")     mult-prim)
    (primitive ("add1")  incr-prim)
    (primitive ("sub1")  decr-prim)
    (primitive ("or")    or-prim)))

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

(define-datatype
 proc proc?
 (closure (vars (list-of var?)) ; vars instead of ids
          (body-exp expression?)
          (env environment?)))

(define-datatype
 reference reference?
 (a-ref (loc location?)))

; New datatype to represent delayed evaluation   <<<<<
(define-datatype
 thunk thunk?
 (a-thunk (expr expression?)
          (env environment?)))

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

; eval-program : program -> expval
(define eval-program 
  (lambda (pgm)
    (reset-trace)
    (cases program pgm
      (a-program (body)
        (eval-expression body (init-env))))))

; eval-expression : expression env -> expval
(define eval-expression 
  (lambda (exp env)
    (show-trace exp env)
    (cases expression exp
      (ref-exp (id) (a-ref (apply-env env id)))
      (lit-exp (datum) datum)
      (var-exp (id) 
               ; Look up loc, if it contains a   <<<<<<
               ; thunk, evaluate it (both
               ; call-by-name and call-by-need)
               (let ([loc (apply-env env id)])
                 (let ([r (location-val loc)])
                   (if (thunk? r)
                       (cases thunk r
                         (a-thunk (exp env)
                           (let ([v (eval-expression exp
                                                   env)])
                             ; Here's the part that makes
                             ; it call-by-need: after
                             ; eval-ing the thunk, put
                             ; the result into the loc
                             (location-set! loc v)
                             v)))
                       r))))
      (primapp-exp (prim rands)
        (let ((args (eval-rands rands env)))
          (apply-primitive prim args)))
      (true-exp () #t)
      (false-exp () #f)
      (proc-exp (vars body-exp)
                (closure vars body-exp env))
      (app-exp (rator rands)
               (let ([func (eval-expression rator env)]
                     [args (eval-fun-rands rands env)])
                 (apply-proc func args)))
      (if-exp (test then else)
              (if (zero? (eval-expression test env))
                  (eval-expression else env)
                  (eval-expression then env)))
      (let-exp (ids exps body-exp) 
               (eval-expression
                body-exp
                (extend-env
                 ids
                 ; Use eval-let-rands, which will <<<<<<<
                 ; introduce thunks as necessary
                 (map location (eval-let-rands exps env))
                 env)))
      (seq-exp (expr1 expr2)
               (eval-expression expr1 env)
               (eval-expression expr2 env))
      (set-exp (id exp) 
               (let ([new-val (eval-expression exp env)]
                     [var-loc (apply-env env id)])
                 (location-set! var-loc new-val)
                 1)))))

; apply-proc : proc list-of-expval/locations -> expval
(define (apply-proc func args)
  (cases proc func
    (closure (vars body-exp env)
     (let ([ids (map get-var-id vars)]
           [locations (map make-var-location vars args)])
       (eval-expression
        body-exp
        (extend-env
         ids
         locations
         env))))))

; get-var-id : var -> sym
;  Gets the id of a variable
(define (get-var-id v)
  (cases var v
    (cbv-var (id) id)
    (cbr-var (id) id)))

; make-var-location : var expval/loc -> loc
;  If v is cbv, creates a new location with
;    the given arg value (extracting it from an
;    arg location if necessary).
;  If v is cbr, uses arg if it is a location,
;    otherwise creates a new location for the value.
(define (make-var-location v arg)
  (cases var v
    (cbv-var (id) 
             ; Make a new location
             (location (if (location? arg)
                           (location-val arg)
                           arg)))
    (cbr-var (id)
             ; Reuse old, if it's a location 
             (if (location? arg)
                 arg
                 (location arg)))))

; eval-fun-rands : list-of-exp env -> list-expval/loc/thk
;  Get locations (for var-exp) or values (all other)
;  for rands
(define eval-fun-rands
  (lambda (rands env)
    (map (lambda (x) (eval-fun-rand x env)) rands)))

; eval-fun-rand : expression env -> expval/loc/thunk
;  If rand is a var, return its location, otherwise
;  evaluate the expression.
(define eval-fun-rand
  (lambda (rand env)
    (cases expression rand
      (var-exp (id)          (show-var-trace env id)
               (apply-env env id))
      ; insteda of evaluating, create a thunk    <<<<<<<
      (else (thunkify-expression rand env)))))

; thunkify-expression : expr env -> thunk/expval
;   If the expression is simple enough, evaluate it.
;   Otherwise, delay evaluation by creating a 
;   thunk, saving the expr and env.
(define thunkify-expression
  (lambda (exp env)
    (cases expression exp
      (lit-exp (dk) (eval-expression exp env))
      (proc-exp (dk1 dk2) (eval-expression exp env))
      ; Anything else is "complicated"
      (else (a-thunk exp env)))))

; eval-rands : list-of-expr env -> list-of-expval
;  Normal evaluation
(define eval-rands
  (lambda (rands env)
    (map (lambda (x) (eval-expression x env)) rands)))

; eval-let-rands : listof-exp env -> listof-thk/expval <<<
;  Thunkify, don't evaluate
(define eval-let-rands
  (lambda (rands env)
    (map (lambda (x) (thunkify-expression x env)) rands)))

; apply-primitive : primitive list-of-expval -> expval
(define apply-primitive
  (lambda (prim args)
    (cases primitive prim
      (setref-prim () (let ([ref (car args)]
                            [v (cadr args)])
                        (location-set!
                         (cases reference ref
                           (a-ref (loc) loc))
                         v)
                        1))
      (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))
      (or-prim () (or (car args) (cadr args))))))

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

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

;; Abstract envrionment datatype implemenation.
;; (Restored to the original implementation.)

(define-datatype environment environment?
  (empty-env-record)
  (extended-env-record
    (syms (list-of symbol?))
    (vec vector?)    ; can use this for anything.
    (env environment?)))

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

; extend-env : list-of-sym list-of-expval env -> env 
(define extend-env
  (lambda (syms vals env)
    (extended-env-record syms (list->vector vals) env)))

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

;; Environment helper functions

; env-find-position : sym list-of-symbols -> num-or-#f
(define env-find-position 
  (lambda (sym los)
    (list-find-position sym los)))

; list-find-position : sym list-of-symbols -> num-or-#f
(define list-find-position
  (lambda (sym los)
    (list-index (lambda (sym1) (eqv? sym1 sym)) los)))

; list-index : pred list-of-symbols -> num-or-#f
(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))))))

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

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

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

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

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Tracing

(define (print-expr exp)
  (cases expression exp
    (ref-exp (id) (printf "ref(~a)" id))
    (lit-exp (v) (printf "~a" v))
    (var-exp (id) (printf "~a" id))
    (true-exp () (printf "true"))
    (false-exp () (printf "false"))
    (primapp-exp (prim rands)
                 (printf "~a("
                  (cases primitive prim
                    (setref-prim () 'setref)
                    (add-prim  () '+)
                    (subtract-prim () '-)
                    (mult-prim  () '*)
                    (incr-prim  () 'add1)
                    (decr-prim  () 'sub1)
                    (or-prim    () 'or)))
                 (for-each
                  (lambda (rand)
                    (if (not (eq? rand (car rands)))
                      (printf ", "))
                    (print-expr rand))
                  rands)
                 (printf ")"))
    (proc-exp (vars body)
              (printf "proc(")
              (for-each
               (lambda (v)
                 (if (not (eq? v (car vars)))
                   (printf ", "))
                 (cases var v
                   (cbv-var (id) (printf "~a" id))
                   (cbr-var (id) (printf "&~a" id))))
               vars)
              (printf ")")
              (print-expr body))
    (app-exp (rator rands)
             (printf "(")
             (print-expr rator)
             (for-each (lambda (rand)
                         (printf " ")
                         (print-expr rand))
                       rands)
             (printf ")"))
    (if-exp (test then else)
            (printf "if ")
            (print-expr test)
            (printf " then ")
            (print-expr then)
            (printf " else ")
            (print-expr else))
    (let-exp (ids exps body-exp)
             (printf "let")
             (for-each
              (lambda (id exp)
                (printf " ~a=" id)
                (print-expr exp))
              ids exps)
             (printf " in ")
             (print-expr body-exp))
    (seq-exp (expr1 expr2)
             (printf "{")
             (print-expr expr1)
             (printf ";")
             (print-expr expr2)
             (printf "}"))
    (set-exp (id exp) 
             (printf "set ~a=" id)
             (print-expr exp))))

(define (print-env env)
  (cases environment env
    (empty-env-record () (printf "{}"))
    (extended-env-record 
     (syms locs old-env)
     (for-each name-loc (vector->list locs))
     (printf "{")
     (for-each (lambda (sym loc)
                 (printf "~a=~a, "
                         sym
                         (get-loc-name loc)))
               syms (vector->list locs))
     (print-env old-env)
     (printf "}"))))

(define loc-name-table (#%make-hash-table))

; name-loc : location ->
;  Assigns a name to loc
(define (name-loc loc)
  (#%hash-table-get
   loc-name-table
   loc
   ; proc called when loc is not found:
   (lambda ()
     (#%hash-table-put!
      loc-name-table
      loc
      (#%gensym)))))

; get-loc-name : location -> sym
(define (get-loc-name loc)
  (#%hash-table-get loc-name-table loc))

(define (print-locs)
  (#%hash-table-for-each
   loc-name-table
   (lambda (loc name)
     (printf " ~a=" name)
     (print-value (location-val loc)))))

(define (print-value v)
  (cond
    [(reference? v)
     (cases reference v
       (a-ref (loc) (printf "[~a]" (get-loc-name))))]
    [(thunk? v)
     (cases thunk v
       (a-thunk 
        (expr env) 
        (printf "<")
        (print-expr expr)
        (printf ", ")
        (print-env env)
        (printf ">")))]        
    [(proc? v)
     (cases proc v
       (closure 
        (vars body-expr env)
        (printf "<")
        (for-each
         (lambda (vr)
           (if (not (eq? vr (car vars)))
               (printf " "))
           (cases var vr
             (cbv-var (id) (printf "~a" id))
             (cbr-var (id) (printf "&~a" id))))
         vars)
        (printf ", ")
        (print-expr body-expr)
        (printf ", ")
        (print-env env)
        (printf ">")))]
    [else
     (printf "~a" v)]))

(define (show-trace exp env)
  (printf "eval~n")
  (printf " exp: ")
  (print-expr exp)
  (newline)
  (printf " env: ")
  (print-env env)
  (newline)
  (printf " loc:")
  (print-locs)
  (newline))

(define (show-var-trace env id)
  (printf "lookup~n")
  (printf " var: ~a~n" id)
  (printf " env: ")
  (print-env env)
  (newline)
  (printf " loc:")
  (print-locs)
  (newline))

(define (reset-trace)
  (set! loc-name-table (#%make-hash-table)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Locations

; location : obj -> location
;   Creates location, given its initial value
(define (location val)
  (vector val))

; location? : obj -> bool
;  Returns #t if o is a location, #f otherwise
(define (location? o)
  (and (vector? o) (= 1 (vector-length o))))

; location-val : location -> obj
;  Returns the value in loc
(define (location-val loc)
  (vector-ref loc 0))

; location-set! : location obj -> 
;   Changes loc to contain v
(define (location-set! loc v)
  (vector-set! loc 0 v))