(* Lambda term (encoded using de bruijn indices), booleans and delimited continuations *)

datatype builtin_op = Sum | Sub | LessThan | Equal

datatype term =
    Bool of bool
  | Int of Int64.int
  | Branch of term * term * term
  | Var of int
  | Lambda of term
  | App of term * term
  | BuiltinApp of builtin_op * term * term
  | Fix of term

(* Value, the result of a computation *)
datatype value =
    Closure of env * term
  | BoolV of bool
  | IntV of Int64.int
  | Loop

  (* Why a ref? Because of Fix! We want to update this value post-hoc *)
and env = Env of value ref list

fun pretty_print_value (Closure _) = "<closure>"
  | pretty_print_value (BoolV b) = Bool.toString b
  | pretty_print_value (IntV n) = Int64.toString n
  | pretty_print_value Loop = "<loop>"


exception UnboundVar

val empty_env = Env []

fun env_push_ref (r, Env e) = Env (r :: e)
fun env_push_value (v, e) = env_push_ref (ref v, e)

fun env_lookup (idx, Env e) = !(List.nth (e, idx))
  handle Subscript => raise UnboundVar

exception Stuck of string

(* Handle a possibly recursion record *)
fun unpack_recursion Loop = raise Stuck "<infinite loop>"
  | unpack_recursion x = x

fun eval (_, Bool b) = BoolV b
  | eval (_, Int i) = IntV i
  | eval (env, Branch (cond, trueb, falseb)) = eval_branch (env, eval (env, cond), trueb, falseb)
  | eval (env, Var idx) = unpack_recursion (env_lookup (idx, env))
  | eval (env, Lambda t) = Closure (env, t)
  | eval (env, App (lhs, rhs)) = eval_app (env, eval (env, lhs), eval (env, rhs))
  | eval (env, BuiltinApp (oper, lhs, rhs)) = eval_builtin (oper, eval (env, lhs), eval (env, rhs))
  | eval (env, Fix body) = eval_fix (env, body)

and eval_branch (env, BoolV true, trueb, _) = eval (env, trueb)
  | eval_branch (env, BoolV false, _, falseb) = eval (env, falseb)
  | eval_branch _ = raise Stuck "Condition is not a boolean"

and eval_app (env, Closure (cEnv, body), arg) = eval (env_push_value (arg, cEnv), body)
  | eval_app _= raise Stuck "Left term in application is not a lambda"

and eval_builtin (Sum, IntV n1, IntV n2) = IntV (n1 + n2)
  | eval_builtin (Sub, IntV n1, IntV n2) = IntV (n1 - n2)
  | eval_builtin (LessThan, IntV n1, IntV n2) = BoolV (n1 < n2)
  | eval_builtin (Equal, IntV n1, IntV n2) = BoolV (n1 = n2)
  | eval_builtin _ = raise Stuck "Argument of builtin op is not integer"

and eval_fix (env, body) =
  let
  (* Step 1. Augment the body with an infinite loop variable *)
    val self = ref Loop
    val new_env = env_push_ref (self, env)
    val result = eval (new_env, body)
    (* Here we tie the knot and assign self to the result of the evaluation *)
    val _ = self := result
  in
    result
  end
  

fun evaluate term = eval (empty_env, term)

(* 1. Can we implement recursion in terms of shift and reduce? Re recursion is an effect *)

(* The term i want to define is:
fix (\sum. (\x. if x = 0 then 0 else sum (x - 1)))

Let us create the nameless representation
fix (\. (\. if $0 = 0 then 0 else $0 + $1 ($0 - 1)))
*)

val args = CommandLine.arguments()
val input_string = hd args
val input_number = valOf (Int64.fromString input_string)

local
  val cond = BuiltinApp (Equal, Var 0, Int 0)
  fun decrement x = BuiltinApp (Sub, x, Int 1)
  fun add x y = BuiltinApp (Sum, x, y)
  val branch = Branch (cond, Int 0, add (Var 0) (App (Var 1, decrement (Var 0))))
in
  fun sum input_number = App (Fix (Lambda branch), Int input_number)
end


(* The term i want to define is:
fix (\sum. (\acc. \n. if n = 0 then acc else sum (acc + 1) (n - 1)))

Let us create the nameless representation
fix (\. (\. \. if $0 = 0 then $1 else $2 ($1 + $0) ($0 - 1)))
*)

local
  val cond = BuiltinApp (Equal, Var 0, Int 0)
  fun decrement x = BuiltinApp (Sub, x, Int 1)
  fun sum x y = BuiltinApp (Sum, x, y)
  val branch = Branch (cond, Var 1, App (App (Var 2, sum (Var 1) (Var 0)), decrement (Var 0)))
in
  fun iter_sum input_number = App (App (Fix (Lambda (Lambda branch)), Int 0), Int input_number)
end


(* The term i want to define is:
fix (\fib. (\x. if x < 2 then x else fib (x - 1) + fib (x - 2)))

Let us create the nameless representation
fix (\. (\. if $0 < 2 then $0 else $1 ($0 - 1) + $1 ($0 - 2))) *)

local
  val cond = BuiltinApp (LessThan, Var 0, Int 2)
  fun sub x y = BuiltinApp (Sub, x, y)
  val fib_minus_1 = App (Var 1, sub (Var 0) (Int 1))
  val fib_minus_2 = App (Var 1, sub (Var 0) (Int 2))
  val recursive_case = BuiltinApp (Sum, fib_minus_1, fib_minus_2)
  val branch = Branch (cond, Var 0, recursive_case)
in
  fun fib input_number = App (Fix (Lambda branch), Int input_number)
end

val term = sum input_number
val _ = print (pretty_print_value (evaluate term))
val _ = print "\n"