#lang plait
(define-type Exp
  [numE (n : Number)]
  [plusE (lhs : Exp) (rhs : Exp)]
  [minusE (lhs : Exp) (rhs : Exp)]
  [varE (name : Symbol)]
  [lamE (param : Symbol) (body : Exp)]
  [appE (fun-expr : Exp) (arg-expr : Exp)]
  [errorE (msg : String)] ;; New
  [if0E (test : Exp) (then : Exp) (else : Exp)])

(define-type Value
  [numV (n : Number)]
  [errorV (msg : String)]
  [funV (param : Symbol)
            (body : Exp)
            (env : Env)])

(define-type Env
  [emptyEnv]
  [Extend (name : Symbol)
        (value : Value)
        (rest : Env)])

(define (lookup name env)
  (type-case Env env
    [(emptyEnv) (error 'lookup (string-append "no binding for" (to-string name)))]
    [(Extend id val rest-env)
     (if (eq? id name)
         val
         (lookup name rest-env))]))

(define-type Continuation
  [emptyCont]
  [plusSecondK (r : Exp)
              (env : Env)
              (k : Continuation)]
  [doPlusK (v1 : Value)
          (k : Continuation)]
  [minusSecondK (r : Exp)
              (env : Env)
              (k : Continuation)]
  [doMinusK (v1 : Value)
          (k : Continuation)]
  [appArgK (arg-expr : Exp)
            (env : Env)
            (k : Continuation)]
  [doAppK (fun-val : Value)
           (k : Continuation)]
  [doIfK (then-expr : Exp)
         (else-expr : Exp)
         (env : Env)
         (k : Continuation)])

(define (parse-error sx)
  (error 'parse (string-append "parse error: " (to-string sx))))

(define (sx-ref sx n) (list-ref (s-exp->list sx) n))
(define (parse sx)
  (local
      [(define (px i) (parse (sx-ref sx i)))]
    (cond
      [(s-exp-number? sx) (numE (s-exp->number sx))]
      [(s-exp-symbol? sx) (varE (s-exp->symbol sx))]
      [(s-exp-match? `(let1 SYMBOL ANY ANY) sx)
       (let* ([id (s-exp->symbol (sx-ref sx 1))]
              [named (px 2)]
              [body (px 3)])
         (appE (lamE id body) named))]
      [(s-exp-match? `(lam (SYMBOL) ANY) sx)
       (let* ([args (sx-ref sx 1)]
              [varE (s-exp->symbol (sx-ref args 0))]
              [body (px 2)])
         (lamE varE body))]
      [(s-exp-match? `(error STRING) sx) (errorE (s-exp->string (sx-ref sx 1)))]
      [(s-exp-match? `(error ANY) sx) (parse-error sx)]
      [(s-exp-match? `(ANY ANY) sx)
       (appE (px 0) (px 1))]
      [(s-exp-list? sx)
       (case (s-exp->symbol (sx-ref sx 0))
         [(+) (plusE (px 1) (px 2))]
         [(-) (minusE (px 1) (px 2))]
         [(if0) (if0E (px 1) (px 2) (px 3))]
         [else (parse-error sx)])]
      [else (parse-error sx)])))

(define (arith-op op val1 val2)
  (local
      [(define (numV->number v)
         (type-case Value v
           [(numV n) n]
           [else (error 'arith-op
                        (string-append "expects a number, got: " (to-string v)))]))]
    (numV (op (numV->number val1)
              (numV->number val2)))))

(define (numzero? x)
   (zero? (numV-n x)))

(define (interp expr env k)
  (type-case Exp expr
    [(numE n) (continue k (numV n))]
    [(plusE l r) (interp l env (plusSecondK r env k))]
    [(minusE l r) (interp l env (minusSecondK r env k))]
    [(varE name) (continue k (lookup name env))]
    [(errorE msg) (errorV msg)]
    [(lamE param body-expr)
     (continue k (funV param body-expr env))]
    [(appE fun-expr arg-expr)
     (interp fun-expr env (appArgK arg-expr env k))]
    [(if0E test-expr then-expr else-expr)
     (interp test-expr env (doIfK then-expr else-expr env k))]))

(define (continue [k : Continuation] [v : Value]) : Value
  (type-case Continuation k
    [(emptyCont) v]
    [(plusSecondK r env next-k)
     (interp r env (doPlusK v next-k))]
    [(doPlusK v1 next-k)
     (continue next-k (arith-op + v1 v))]
    [(minusSecondK r env next-k)
     (interp r env (doMinusK v next-k))]
    [(doMinusK v1 next-k)
     (continue next-k (arith-op - v1 v))]
    [(appArgK arg-expr env next-k)
     (interp arg-expr env (doAppK v next-k))]
    [(doAppK fun-val next-k)
     (interp (funV-body fun-val)
           (Extend (funV-param fun-val) v (funV-env fun-val))
           next-k)]
    [(doIfK then-expr else-expr env next-k)
     (if (numzero? v)
         (interp then-expr env next-k)
         (interp else-expr env next-k))]))

(define init-k (emptyCont))

(define (run s-exp)
  (interp (parse s-exp) (emptyEnv) (emptyCont)))

(module+ test
  (test (interp (numE 10)
                (emptyEnv)
                init-k)
        (numV 10))
  (test (interp (plusE (numE 10) (numE 7))
                (emptyEnv)
                init-k)
        (numV 17))
  (test (interp (minusE (numE 10) (numE 7))
                (emptyEnv)
                init-k)
        (numV 3))
  (test (interp (appE (lamE 'x (plusE (varE 'x) (numE 12)))
                      (plusE (numE 1) (numE 17)))
                (emptyEnv)
                init-k)
        (numV 30))
  (test (interp (varE 'x)
                (Extend 'x (numV 10) (emptyEnv))
                init-k)
        (numV 10))

  (test (interp (appE (lamE 'x (plusE (varE 'x) (numE 12)))
                      (plusE (numE 1) (numE 17)))
                (emptyEnv)
                init-k)
        (numV 30))


  (test/exn (interp (varE 'x) (emptyEnv) init-k)
            "no binding")

  (test/exn
   (run `{ {lam {x} {+ x y}} 0})
   "no binding")

  (test
   (run
    `{{lam {x}
           {{lam {f} {f 2}}
            {lam {y} {+ x y}}}}
      0})
   (numV 2))


  (test (run `{let1 f {lam {x} {+ x 1}}
                    {+ {f 2} {error "abort!"}}})
        (errorV "abort!"))

  (test (run `{{lam {x}
                    {{lam {f}
                          {+ {f 1}
                             {{lam {x} {f 2}} 3}}}
                     {lam {y} {+ x y}}}}
                 0})
        (numV 3))

  (test (interp (if0E (numE 0)
                     (numE 1)
                     (numE 2))
                (emptyEnv)
                init-k)
        (numV 1))
  (test (interp (if0E (numE 1)
                     (numE 0)
                     (numE 2))
                (emptyEnv)
                init-k)
        (numV 2))

  (test (run
         `{{lam {mkrec}
                {{lam {fib}
                      ;; Call fib on 10:
                      {fib 10}}
                 ;; Create recursive fib:
                 {mkrec
                     {lam {fib}
                          ;; Fib:
                          {lam {n}
                               {if0 n
                                    1
                                    {if0 {- n 1}
                                         {error "reached zero"}
                                         {+ {fib {- n 1}}
                                            {fib {- n 2}}}}}}}}}}
              ;; mkrec:
              {lam {body-proc}
                   {{lam {fX}
                         {fX fX}}
                    {lam {fX}
                         {body-proc {lam {x} {{fX fX} x}}}}}}})
           (errorV "reached zero"))
)