#lang plait
(define-type IFLANG
  [Num  (val : Number)]
  [Add  (l : IFLANG) (r : IFLANG)]
  [Div  (l : IFLANG) (r : IFLANG)]
  [Id   (name : Symbol)]
  [Let1 (id : Symbol) (named-expr : IFLANG) (bound-body : IFLANG)]
  [Lam  (param : Symbol) (body : IFLANG)]
  [ILam  (param : Symbol) (zero-body : IFLANG) (non-zero-body : IFLANG)]
  [Call (fun : IFLANG) (val : IFLANG)]) ; first type!
(define-type ENV
  [EmptyEnv]
  [Extend (name : Symbol) (val : VAL) (rest : ENV)])

(define-type VAL
  [NumV (n : Number)]
  [LamV (arg : Symbol) (body : IFLANG) (env : ENV)]
  [ILamV (arg : Symbol) (zero-body : IFLANG)  (non-zero-body : IFLANG) (env : ENV)])

(define (lookup name env)
  (type-case ENV env
    [(EmptyEnv) (error 'lookup "missing binding")]
    [(Extend id val rest-env)
     (if (eq? id name)
         val
         (lookup name rest-env))]))

(define (arith-op op val1 val2) (NumV (op (NumV-n val1) (NumV-n val2))))
;; evaluates IFLANG expressions by reducing them to values
(define (interp expr env)
  (type-case IFLANG expr
    [(Num n) (NumV n)]
    [(Add l r) (arith-op + (interp l env)
                         (interp r env))]
    [(Div l r) (arith-op / (interp l env) (interp r env))]
    [(Let1 bound-id named-expr bound-body)
     (interp bound-body
           (Extend bound-id (interp named-expr env) env))]
    [(Id name) (lookup name env)]
    [(ILam bound-id zero-body non-zero-body)
     ; I thought this might delay the lookup so it would be bound by this point
     (LamV bound-id (type-case VAL (lookup bound-id env)
                      [(NumV n) (if (equal? n 0) zero-body non-zero-body)]
                      [else (error 'ILam "not a number in ILam expression")]) env)
     ; I got confused by why the arguments in the environment for a function are not bound to the identifier in the
     ; lambda
     #;(type-case VAL (lookup bound-id env)
       [(NumV n) (if (equal? n 0)
                     (LamV bound-id zero-body env)
                     (LamV bound-id non-zero-body env))]
       [else (error 'ILam "not a number in ILam expression")])
     ]
    [(Lam bound-id bound-body)
     (LamV bound-id bound-body env)]
    [(Call lam-expr arg-expr)
     (let ([fval (interp lam-expr env)])
       (type-case VAL fval
         [(LamV bound-id bound-body f-env)
          (interp bound-body
                (Extend bound-id (interp arg-expr env) f-env))]
         [else (error 'eval "expects function")]))]))
(define (run iflang)  (NumV-n (interp iflang (EmptyEnv))))

(print-only-errors #t)

(test (run
       (Let1 'x (Num 3)
        (Let1 'f (ILam 'y (Num -1) (Add (Id 'x) (Id 'y)))
         (Let1 'x (Num 5) (Call (Id 'f) (Num 4))))))
      7)

(test (run
       (Call (Let1 'x (Num 3) (ILam 'y (Num -1) (Add (Id 'x) (Id 'y)))) (Num 4)))
      7)

(test (run 
       (Call
        (Call
         (Lam 'x (Call (Id 'x) (Num 1)))
         (Lam 'x (ILam 'y (Num -1) (Add (Id 'x) (Id 'y)))))
        (Num 123)))
       124)

(test (run
       (Let1 'f (ILam 'x (Id 'x) (Div (Num 1) (Id 'x)))
        (Add (Call (Id 'f) (Num 3)) (Call (Id 'f) (Num 0)))))
      1/3)

(test (run
       (Let1 'x (Num 3)
        (Let1 'f  (ILam 'y (Id 'x) (Div (Id 'x) (Id 'y)))
         (Let1 'x (Num 5) (Add (Call (Id 'f) (Num 3)) (Call (Id 'f) (Num 0)))))))
      4)
