(* K machine in cbv *)
(* 15-814, Fall 2021, Frank Pfenning *)
(* post-lecture extension *)

(* turn off non-exhaustiveness warnings *)
(* don't do this at home, folks *)
#options --warn=false

type E = $E. ('lam : E -> E) + ('app : E * E)
             + ('left : E) + ('right : E) + ('cases : E * (E -> E) * (E -> E))
             + ('unit : 1) + ('caseu : E * (1 -> E))
             + ('pair : E * E) + ('casep : E * (E -> E -> E))
             + ('fold : E) + ('unfold : E)
             + ('fix : E -> E)

decl eval_ : E -> (E -> E) -> E

defn eval_ = $eval_. \e. \k.
               case (unfold e)
                 of ( 'lam _ => k e
                    | 'app (e1,e2) =>
                         eval_ e1 (\v1. eval_ e2 (\v2.
                           case (unfold v1) of ( 'lam f => eval_ (f v2) k )))
                    | 'left e => eval_ e (\v. k (fold 'left v))
                    | 'right e => eval_ e (\v. k (fold 'right v))
                    | 'cases (e0, b1, b2) =>
                         eval_ e0 (\v0. case (unfold v0)
                                          of ('left v => eval_ (b1 v) k
                                             |'right v => eval_ (b2 v) k))
                    | 'unit _ => k e
                    | 'caseu (e0,b) =>
                         eval_ e0 (\v0. case (unfold v0)
                                          of ('unit u => eval_ (b u) k))
                    | 'pair (e1,e2) =>
                         eval_ e1 (\v1. eval_ e2 (\v2. k (fold 'pair (v1,v2))))
                    | 'casep (e0,b) =>
                         eval_ e0 (\v0. case (unfold v0)
                                          of ('pair (v1,v2) => eval_ (b v1 v2) k))
                    | 'fold e => eval_ e (\v. k (fold 'fold v))
                    | 'unfold e =>
                         eval_ e (\v. case (unfold v)
                                        of ('fold v => k v))
                    | 'fix f => eval_ (f e) k
                    )

(* more constructor function would help in writing code *)
decl app : E -> (E -> E)
decl lam : (E -> E) -> E

defn app = \e1. \e2. fold 'app (e1,e2)
defn lam = \f. fold 'lam f

(* nat = mu alpha. ('l : 1) + ('r : alpha) *)
decl zero : E (* nat *)
decl succ : E (* nat -> nat *)
decl pred : E (* nat -> nat *)
decl double : E (* nat -> nat *)

defn zero = fold 'fold fold 'left fold 'unit ()
defn succ = fold 'lam (\n. fold 'fold fold 'right n)
defn pred = fold 'lam (\n. fold 'cases (fold 'unfold n, (\u. zero), \n'. n'))
defn double = fold 'fix (\double. fold 'lam (\n. fold 'cases (fold 'unfold n, (\u. zero),
                                                   \n'. app succ (app succ (app double n')))))

(* recursion can also be implemented at the metalevel *)
decl double' : E (* nat -> nat *)
defn double' = $double. fold 'lam (\n. fold 'cases (fold 'unfold n, (\u. zero),
                                                   \n'. app succ (app succ (app double n'))))

eval three = eval_ (app succ (app succ (app succ zero))) (\v. v)
eval two = eval_ (app pred three) (\v. v)
eval six = eval_ (app double three) (\v. v)

(* an ill-typed as an object language program *)
(* evaluation fails *)
fail
eval error = eval_ (app zero zero) (\v. v)