#lang plait
(define-type Exp
  [numE (n : Number)]
  [plusE (left : Exp) (right : Exp)]
  [timesE (left : Exp) (right : Exp)]
  [minusE (left : Exp) (right : Exp)]
  [lamE (var : Symbol) (body : Exp)]
  [appE (fun : Exp) (arg : Exp)]
  [varE (name : Symbol)]
  [if0E (check : Exp) (zero : Exp) (non-zero : Exp)]
  [let1E (var : Symbol) (value : Exp) (body : Exp)]
  [recE (var : Symbol) (value : Exp) (body : Exp)])

(define-type Value
  [numV (the-number : Number)]
  [funV (var : Symbol) (body : Exp) (nv : Env)]
  [undefV])

(define (parse s)
  (local
      [(define (sx n) (list-ref (s-exp->list s) n))
       (define (px n) (parse (sx n)))
       (define (? pat) (s-exp-match? pat s))
       (define (parse-let)
         (let* ([def (sx 1)]
                [parts (s-exp->list def)]
                [var (s-exp->symbol (list-ref parts 0))]
                [val (parse (list-ref parts 1))]
                [body (px 2)])
           (values var val body)))]
    (cond
      [(? `SYMBOL) (varE (s-exp->symbol s))]
      [(? `NUMBER) (numE (s-exp->number s))]
      [(? `(+ ANY ANY)) (plusE (px 1) (px 2))]
      [(? `(- ANY ANY)) (minusE (px 1) (px 2))]
      [(? `(* ANY ANY)) (timesE (px 1) (px 2))]
      [(? `(if0 ANY ANY ANY))
       (if0E (px 1) (px 2) (px 3))]
      [(? `(rec (SYMBOL ANY) ANY))
       (local [(define-values (var val body) (parse-let))]
         (recE var val body))]
      [(? `(let1 (SYMBOL ANY) ANY))
       (local [(define-values (var val body) (parse-let))]
         (let1E var val body))]
      [(? `(lam SYMBOL ANY))
       (lamE (s-exp->symbol (sx 1)) (px 2))]
      [(? `(ANY ANY)) (appE (px 0) (px 1))]
      [else (error 'parse (to-string s))])))

(define (num-op op expr1 expr2)
  (local [(define (unwrap v)
            (type-case Value v
              [(numV n) n]
              [else (error 'num-op "NaN")]))]
    (numV (op (unwrap expr1)
              (unwrap expr2)))))

(define-type-alias Env (Hashof Symbol (Boxof Value)))

(define mt-env (hash empty)) ;; "empty environment"

(define (extend old-env new-name value)
  (hash-set old-env new-name (box value)))

(define (lookup (s : Symbol) (n : Env))
  (type-case (Optionof (Boxof Value)) (hash-ref n s)
    [(none) (error s "not bound")]
    [(some b) (unbox b)]))

(test/exn (lookup 'x mt-env) "not bound")

(interp : (Exp Env -> Value))
(define (interp e nv)
  (type-case Exp e
    [(numE n) (numV n)]
    [(varE s) (lookup s nv)]
    [(plusE l r) (num-op + (interp l nv) (interp r nv))]
    [(minusE l r) (num-op - (interp l nv) (interp r nv))]
    [(timesE l r) (num-op * (interp l nv) (interp r nv))]
    [(lamE v b) (funV v b nv)]
    [(if0E c z nz)
     (if (equal? (numV 0) (interp c nv))
         (interp z nv)
         (interp nz nv))]
     [(appE f a)
     (let ([fv (interp f nv)]
           [av (interp a nv)])
       (type-case Value fv
         [(funV v b f-env)
          (interp b (extend f-env v av))]  ;; changed
         [else (error 'app "not a function")]))]
    [(let1E var val body)
     (let ([new-env (extend nv
                            var
                            (interp val nv))])
       (interp body new-env))]
    [(recE var val body) ....]))

(run : (S-Exp -> Value))
(define (run s)
  (interp (parse s) mt-env))

(test (run `{let1 {f {lam x {+ x 1}}} {f 8}}) (numV 9))
(test (run `{let1 {y 1} {let1 {f {lam x {+ x y}}}
                              {f 8}}})
      (numV 9))
(test (run `{let1 {y 1} {let1 {f {lam x {+ x y}}}
                              {let1 {y 2} {f 8}}}})
      (numV 9))
(test (run `{{let1 {x 3} {lam y {+ x y}}} 4})
      (numV 7))
(test (run `{{let1 {y 3} {lam y {+ y 1}}} 5})
      (numV 6))
(test (run `{if0 0 (* 1 2) 1}) (numV 2))