(*********************************************)
(* (Extended) Krivine Machine implementation *)
(*                                           *)
(* Guillermo Ramos Gutiérrez (2015)          *)
(*********************************************)


(*********)
(* Utils *)
(*********)
let print_with f x = print_endline (f x)
let concat_with sep f xs = String.concat sep (List.map f xs)
let split n xs =
  let rec split_acc n xs ys = match (n, ys) with
    | (0, _) | (_, []) -> (List.rev xs, ys)
    | (n, y :: ys) -> split_acc (n-1) (y::xs) ys in
  split_acc n [] xs
let find_idx a =
  let rec find_acc n = function
    | [] -> raise Not_found
    | (x :: xs) -> if a == x then n else find_acc (n+1) xs in
  find_acc 0
let rec repeat n x = if n == 0 then [] else x :: repeat (n-1) x
let fold_left1 f = function
  | [] -> raise (Invalid_argument "empty string")
  | (x::xs) -> List.fold_left f x xs

let dbg_lev = 1
let debug lev spaces s =
  if dbg_lev >= lev
  then print_endline (" -- " ^ String.concat "" (repeat spaces "  ") ^ s)


(* Compile, decompile and reduce errors *)
exception CpErr of string
exception DcErr of string
exception RdErr of string

(************************************)
(*  Calculus of Constructions terms *)
(************************************)
module CCLambda = struct
    type symbol = string * int
    let symbol : string -> symbol =
      let id = ref 0 in
      let gensym c : symbol =
        let newid = !id in
        id := !id + 1;
        (c, newid)
      in
      gensym

    type t = Var of symbol | App of t * t | Abs of symbol * t
             | Constr of t constr | Case of t * (symbol constr * t) list
             | Fix of symbol * symbol list * symbol * t
             | Acc of t
    and 'a constr = symbol * 'a list

    let show_symbol (s, n) =
      if dbg_lev > 2 then s ^ "/" ^ string_of_int n else s
    let rec show m =
      let show_branch ((x, vs), m) =
        "(" ^ show_symbol x ^ "(" ^ concat_with "," show_symbol vs ^ ") → "
        ^ show m ^ ")" in
      match m with
      | Var x -> show_symbol x
      | App (m1, m2) -> "(" ^ show m1 ^ " " ^ show m2 ^ ")"
      | Abs (x, m) -> "(λ" ^ show_symbol x ^ "." ^ show m ^ ")"
      | Constr (x, ms) -> show_symbol x ^ "(" ^ concat_with ", " show ms ^ ")"
      | Case (m, bs) -> "(case " ^ show m ^ " of "
                       ^ concat_with " " show_branch bs ^ ")"
      | Fix (f, xs, c, m) ->
         "fix_" ^ string_of_int (List.length xs)
         ^ "(λ" ^ concat_with ".λ" show_symbol (f::xs@[c])
         ^ ". " ^ show m ^ ")"
      | Acc m -> "[" ^ show m ^ "]"
    let print = print_with show

    (* Auxiliar term-generating functions *)
    let identity s =
      let x = symbol s in
      Abs (x, Var x)

    let apps = fold_left1 (fun m n -> App (m, n))

    let none = symbol "None"
    let some = symbol "Some"
    let cons = symbol "Cons"
    let nil  = symbol "Nil"

    (* Peano arithmetic helpers *)
    let z = symbol "z"
    let s = symbol "s"

    let rec peano_add n x =
      if n == 0 then x else peano_add (n-1) (Constr (s, [x]))

    let peano_of_int ?(base=Constr (z, [])) n = peano_add n base

    (* Examples *)

    (* (λx.x) (λy. ((λz.z) y) (λt.t) *)
    let ex_m1 =
      let y = symbol "y" in
      App (identity "x",
           Abs (y, App (App (identity "z",
                             Var y),
                        identity "t")))

    (* (λf.λx. f (f x)) (λy.y) (λz.z) *)
    let ex_id_id =
      let f = symbol "f" in
      let x = symbol "x" in
      App (App (Abs (f, Abs (x, App (Var f, App (Var f, Var x)))),
                identity "y"),
           identity "z")

    (* (λc. case c of (Cons(x, xs) → x) (Nil → Nil)) Cons(λx.x, Nil) *)
    let ex_case_head =
      let c = symbol "c" in
      let x = symbol "x" in
      let xs = symbol "xs" in
      App (Abs (c, Case (Var c, [((cons, [x;xs]), Var x);
                                 ((nil, []), Constr (nil, []))])),
           Constr (cons, [identity "m"; Constr (nil, [])]))

    (* fix_0(λf.λc. case c of (s(x) → s(s(f x))) (z → z)) s(s(s(z))) *)
    let ex_fixpt_dup =
      let f = symbol "f" in
      let c = symbol "c" in
      let x = symbol "x" in
      App (Fix (f, [], c, Case (Var c,
                 [((s, [x]), peano_add 2 (App (Var f, Var x)));
                  ((z, []), peano_of_int 0)])),
           peano_of_int 3)
end

module WeakRed : sig
    type instrs
    type mval

    val show_instrs : instrs -> string
    val show_mval   : mval   -> string

    val compile   : CCLambda.t -> instrs
    val decompile : instrs -> CCLambda.t

    val am_reduce : instrs -> mval
    val extract   : mval -> CCLambda.t

    val reduce : CCLambda.t -> CCLambda.t
end = struct
    open CCLambda

    type dbi = symbol * int
    type instr =
      | ACCESS of dbi
      | CLOSURE of instrs
      | ACCLOSURE of symbol
      | PUSHRETADDR of instrs
      | APPLY of int
      | GRAB of symbol
      | RETURN
      | MAKEBLOCK of symbol * int
      | SWITCH of (symbol constr * instrs) list
      | CLOSUREREC of (symbol * symbol list * symbol) * instrs
      | GRABREC of symbol
    and instrs = instr list
    and accum =
      | NoVar of symbol
      | NoApp of accum * mval list
      | NoCase of accum * instrs * env
      | NoFix of accum * instrs * env
    and mval =
      | Accu of accum
      | Cons of mval constr
      | Clos of instrs * env
      | Ret of instrs * env * int
    and env = mval list
    and stack = mval list
    and state = {
        st_c : instrs;
        st_e : env;
        st_s : stack;
        st_n : int;
    }

    let show_dbi (x, i) =
      string_of_int i ^
        if dbg_lev > 2 then " (" ^ show_symbol x ^ ")" else ""
    let rec show_instr = function
      | ACCESS dbi -> "ACCESS(" ^ show_dbi dbi ^ ")"
      | CLOSURE is -> "CLOSURE(" ^ show_instrs is ^ ")"
      | ACCLOSURE x -> "ACCLOSURE([" ^ show_symbol x ^ "])"
      | PUSHRETADDR is -> "PUSHRETADDR(" ^ show_instrs is ^ ")"
      | APPLY i -> "APPLY(" ^ string_of_int i ^ ")"
      | GRAB x -> "GRAB(" ^ show_symbol x ^ ")"
      | RETURN -> "RETURN"
      | MAKEBLOCK (x, n) -> "MAKEBLOCK(#" ^ show_symbol x ^ ", " ^
                             string_of_int n ^ ")"
      | SWITCH bs -> "SWITCH(" ^ concat_with ", " show_branch bs ^ ")"
      | CLOSUREREC (_, is) -> "CLOSUREREC(" ^ show_instrs is ^ ")"
      | GRABREC x -> "GRABREC(" ^ show_symbol x ^ ")"
    and show_branch (((c, _), _), is) = c ^ " → " ^ show_instrs is
    and show_accum = function
      | NoVar x -> "{NoVar: " ^ show_symbol x ^ "}"
      | NoApp (k, mvs) -> "{NoApp: " ^ show_accum k ^ ", "
                           ^ concat_with ", " show_mval mvs ^ "}"
      | NoCase (k, is, e) -> "{NoCase: " ^ show_accum k ^ ", "
                             ^ show_instrs is ^ ", " ^ show_env e ^ "}"
      | NoFix (k, is, e) -> "{NoFix: " ^ show_accum k ^ ", "
                            ^ show_instrs is ^ ", " ^ show_env e ^ "}"
    and show_mval mval = match mval with
      | Accu k -> show_accum k
      | Cons ((s, _), mvs) ->
         "{#" ^ s ^ if List.length mvs == 0 then "}"
                    else ": " ^ concat_with ", " show_mval mvs ^ "}"
      | Clos (is, e) -> "{Tλ: (" ^ show_instrs is ^ "), "
                       ^ (if List.length e > 0 && List.hd e == mval
                          then "{Tλ <fix>}::" ^ show_env (List.tl e)
                          else show_env e)
                       ^ "}"
      | Ret (is, e, n) -> "{<RET>: (" ^ show_instrs is ^ "), " ^ show_env e ^ ", "
                         ^ string_of_int n ^ "}"
    and show_env e = "[" ^ concat_with ", " show_mval e ^ "]"
    and show_instrs is = concat_with "; " show_instr is

    let show_stk stk = "| " ^ concat_with "\n     | " show_mval stk
    let show_st {st_c; st_e; st_s; st_n} =
      "\n/---------------------------------------------------\n"
      ^ "  C: " ^ show_instrs st_c ^ "\n"
      ^ "  E: " ^ show_env st_e ^ "\n"
      ^ "  S: " ^ show_stk st_s ^ "\n"
      ^ "  N: " ^ string_of_int st_n ^ "\n"
      ^ "\\---------------------------------------------------"

    let ret = [RETURN]

    let compile (m : CCLambda.t) : instrs =
      let open CCLambda in

      let e = [] in
      let rec compile' e (is : instrs) m =
        debug 3 3 ("COMPILING: " ^ show m);
        debug 3 3 ("     IN ENV: " ^ concat_with ", " show_symbol e);
        match m with
        | Var x -> begin
            try ACCESS(x, find_idx x e) :: is
            with Not_found -> raise (CpErr ("Var " ^ fst x ^ " not found"))
          end
        | Abs (x, m) -> CLOSURE(GRAB x :: compile' (x :: e) ret m) :: is
        | App (m1, m2) -> let cont = compile' e [APPLY(1)] m1 in
                         PUSHRETADDR(is) :: compile' e cont m2
        | Constr (x, args) -> let f arg cont = compile' e cont arg in
                             let cont = [MAKEBLOCK (x, List.length args)] in
                             List.fold_right f (List.rev args) (cont @ is)
        | Case (m, bs) ->
           let compile_branch ((c, args), m) =
             let dbi' = List.rev args @ e in
             ((c, args), compile' dbi' ret m) in
           let bs' = List.map compile_branch bs in
           PUSHRETADDR(is) :: compile' e [SWITCH(bs')] m
        | Fix (f, xs, c, m) ->
           let cont = compile' (c :: List.rev xs @ f :: e) ret m in
           CLOSUREREC((f, xs, c),
                      List.map (fun x -> GRAB x) xs @ GRABREC c :: cont) :: is
        | Acc (Var x) -> ACCLOSURE x :: is
        | Acc _ -> raise (CpErr "Trying to compile non-var accumulator") in
      compile' e ret m

    let rec decompile' e s is =
      debug 3 3 ("DECOMPILING: " ^ show_instrs is);
      debug 3 3 ("     IN ENV: " ^ concat_with ", " CCLambda.show e);
      debug 3 3 ("     IN STK: " ^ concat_with ", " CCLambda.show s);
      match is with
      | [] -> List.hd s
      | [RETURN] -> List.hd s
      | (ACCESS (_, i) :: is') -> decompile' e (List.nth e i :: s) is'
      | (CLOSURE c_is :: is') -> decompile' e (decompile' e s c_is :: s) is'
      | (PUSHRETADDR r_is :: is') -> decompile' e (decompile' e s is' :: s) r_is
      | (APPLY i :: is) -> begin
          match (i, s) with
          | (1, a::b::_) -> App (a, b)
          | (n, a::s') -> App (a, decompile' e s' [APPLY (n-1)])
          | _ -> raise (DcErr "Unable to decompile APPLY")
        end
      | (GRAB x :: is') -> Abs (x, decompile' (Var x :: e) s is')
      | (MAKEBLOCK (x, n) :: is') -> let (args, st') = split n s in
                                   decompile' e (Constr (x, args) :: s) is'
      | (SWITCH brs :: is') -> begin
          let decompile_br ((c, parms), is) =
            let parms_m = List.map (fun x -> Var x) parms in
            ((c, parms), decompile' (List.rev parms_m @ e) s is) in
          match s with
          | [] -> Case (Var (symbol "_"), List.map decompile_br brs)
          | (a::s') -> Case (a, List.map decompile_br brs)
        end
      | (CLOSUREREC ((f, xs, c), is') :: is) ->
         let m = decompile' (e) s is' in
         decompile' e (Fix (f, xs, c, m) :: s) is
      | (GRABREC x :: is') -> Abs (x, decompile' (Var x :: e) s is')
      | _ -> raise (DcErr "Unable to decompile (unknown instruction)")
    let decompile = decompile' [] []

    let rec extract mv =
      debug 3 3 ("EXTRACTING: " ^ show_mval mv);
      match mv with
      | Clos (is, _) -> decompile is
      | Cons (x, mvs) -> Constr (x, List.map extract mvs)
      | Accu k ->
         let rec extract_accu = function
           | (NoVar x) -> Acc (Var x)
           | (NoApp (k, mvs)) ->
              Acc (apps (extract_accu k :: List.map extract mvs))
           | (NoCase (k, is, e)) -> begin
               match decompile' (List.map extract e) [] is with
               | Case (m, bs) ->
                  let accu = extract_accu k in
                  Acc (Case (accu, bs))
               | _ -> raise (DcErr "Invalid decompilation of CASE accum.")
             end
           | (NoFix (k, is, e)) -> begin
               match decompile' (List.map extract e) [] is with
               | Fix (f, xs, c, m) ->
                  Acc (Fix (f, xs, c, App (m, extract_accu k)))
               | _ -> raise (DcErr "Invalid decompilation of fixpt accum.")
             end in
         extract_accu k
      | _ -> raise (DcErr "Unable to extract")

    let am_reduce is =
      debug 3 3 ("REDUCING: " ^ show_instrs is);
      let rec eval st =
        debug 4 4 (show_st st);
        let {st_c=c; st_e=e; st_s=s; st_n=n} = st in
        match c with
        | [] -> raise (RdErr "Empty code section")
        | (instr :: c) -> begin
            match instr with
            | ACCESS ((x, _), i) -> begin
                try eval {st with st_c=c; st_s=(List.nth e i :: s)}
                with Not_found -> raise (RdErr ("Var " ^ x ^ " not found"))
              end
            | CLOSURE c' -> eval {st with st_c=c; st_s=(Clos(c', e) :: s)}
            | ACCLOSURE x -> eval {st with st_c=c; st_s=(Accu(NoVar x) :: s)}
            | PUSHRETADDR c' -> eval {st with st_c=c; st_s=(Ret(c', e, n) :: s)}
            | APPLY i -> begin
                match s with
                | (Clos (c', e') :: s) ->
                   eval {st_c=c'; st_e=e'; st_s=s; st_n=i}
                | (Accu k :: s) -> begin
                    let (args, s') = split i s in
                    match s' with
                    | (Ret (c', e', n') :: s'') ->
                       eval {st_c=c'; st_e=e';
                             st_s=(Accu(NoApp(k,args)))::s''; st_n=n'}
                    | _ -> raise (RdErr "APPLY over accu with invalid stack")
                  end
                | _ -> raise (RdErr "APPLY over non-closure mval")
              end
            | GRAB _ -> begin
                if n == 0 then
                  match s with
                  | (Ret (c', e', n') :: s) ->
                     let clos = Clos (instr :: c, e) in
                     eval {st_c=c'; st_e=e'; st_s=clos::s; st_n=n'}
                  | [] -> Clos (instr :: c, e)
                  | _ -> raise (RdErr "GRAB (n=0) over non-retval")
                else
                  match s with
                  | (v :: s) -> eval {st_c=c; st_e=v::e; st_s=s; st_n=n-1}
                  | _ -> raise (RdErr "GRAB (n>0) over empty stack")
              end
            | RETURN -> begin
                if n == 0 then
                  match s with
                  | (v :: Ret (c', e', n') :: s) ->
                     eval {st_c=c'; st_e=e'; st_s=v::s; st_n=n'}
                  | [v] -> v
                  | _ -> raise (RdErr "RETURN over empty stack or non-retval")
                else
                  match s with
                  | (Clos (c', e') :: s) ->
                     eval {st_c=c'; st_e=e'; st_s=s; st_n=n}
                  | _ -> raise (RdErr "RETURN over empty stack or non-retval")
              end
            | MAKEBLOCK (x, m) -> let (vs, s') = split m s in
                                 eval {st_c=c; st_e=e;
                                       st_s=(Cons (x, vs)::s'); st_n=n}
            | SWITCH bs -> begin
                match s with
                | (Cons (x, vs) :: s) ->
                   let (_, c') = try List.find (fun ((y, _), _) -> y == x) bs
                                 with Not_found ->
                                   raise (RdErr "SWITCH constr id not found") in
                   let e' = List.rev vs @ e in
                   eval {st_c=c'; st_e=e'; st_s=s; st_n=0}
                | (Accu k :: s) ->
                   let accu = Accu (NoCase (k, instr :: c, e)) in
                   eval {st_c=ret; st_e=e; st_s=accu::s; st_n=0}
                | _ -> raise (RdErr "SWITCH over empty stack or non-constr")
              end
            | CLOSUREREC (_, c') -> let rec v = Clos (c', v::e) in
                                   eval {st with st_c=c; st_s=v::s}
            | GRABREC _ -> begin
                match (s, n) with
                | ([Accu k], 1) -> Accu (NoFix (k, c, e))
                | (Accu k :: s, n) ->
                   let accu = Accu (NoFix (k, instr :: c, e)) in
                   eval {st_c=ret; st_e=e; st_s=accu::s; st_n=n-1}
                | (Ret (c', e', n') :: s, n) ->
                   let clos = Clos (instr :: c, e) in
                   eval {st_c=c'; st_e=e'; st_s=clos::s; st_n=n'}
                | (mval :: s, n) ->
                   eval {st_c=c; st_e=mval::e; st_s=s; st_n=n-1}
                | ([], 0) -> Clos (instr :: c, e)
                | _ -> raise (RdErr "GRABREC over empty stack or invalid mval")
              end
          end in
      let st = {st_c = is; st_e = []; st_s = []; st_n = 0} in
      eval st

    let reduce m = debug 2 1 ("V( " ^ show m ^ " )");
                   match m with
                   | App (m1, m2) -> m |> compile |> am_reduce |> extract
                   | _ -> m
end

module StrongRed = struct
    open CCLambda

    let rec extract_unique = function
      | Var x -> x
      | Acc x -> extract_unique x
      | _ -> raise (RdErr "Trying to extract invalid value")

    let unique x = symbol (fst x)

    let rec reduce m =
      debug 2 0 ("N( " ^ show m ^ " )");
      match m with
      | Fix (f, xs, c, m) ->
         let xs' = List.map unique xs in
         let f' = unique f in
         let c' = unique c in
         let accs = List.map (fun x -> Acc (Var x)) (f' :: xs' @ [c']) in
         let m' = reduce (apps (m :: accs)) in
         Fix (f', xs', c', m')
      | _ -> readback (WeakRed.reduce m)
    and readback m =
      debug 2 1 ("R( " ^ show m ^ " )");
      match m with
      | Abs (x, m) -> let u = unique x in
                     Abs (u, reduce (App (Abs (x, m), Acc (Var u))))
      | Constr (x, vs) -> Constr (x, List.map readback vs)
      | Acc k -> readback_acc k
      | App (Fix (f, xs, c, m), p) ->
         App (reduce (Fix (f, xs, c, m)), readback p)
      | _ -> raise (RdErr "Readback of invalid value")
    and readback_acc m =
      debug 2 2 ("R'( " ^ show m ^ " )");
      match m with
      | Var x -> Var x
      | App (k, v) -> App (readback_acc k, readback v)
      | Acc (Var x) -> Var x
      | Case (k, bs) -> let x = extract_unique k in
                       let b = Abs (x, Case (Var x, bs)) in
                       let rb_branch ((c, xs), m) =
                         let us = List.map unique xs in
                         let acc_us = List.map (fun x -> Acc (Var x)) us in
                         ((c, us), reduce (App (b, Constr (c, acc_us)))) in
                       Case (readback_acc k, List.map rb_branch bs)
      | _ -> raise (RdErr "Readback of invalid accumulator")
end

let weakred m =
  print_endline "WEAK REDUCTION TEST";
  print_endline ("Lambda term:\n    " ^ CCLambda.show m);
  let compiled = WeakRed.compile m in
  let reduced = WeakRed.am_reduce compiled in
  print_endline ("Reduced term:\n    " ^ CCLambda.show (WeakRed.extract reduced));
  print_endline "----------------------------------------------------------\n"

let strongred m =
  print_endline "STRONG REDUCTION TEST";
  print_endline ("Lambda term:\n    " ^ CCLambda.show m);
  print_endline ("Reduced term:\n    " ^ CCLambda.show (StrongRed.reduce m));
  print_endline "----------------------------------------------------------\n"

let () =
  let open CCLambda in
  let examples = [ex_m1; ex_id_id; ex_case_head; ex_fixpt_dup] in
  List.iter weakred examples;
  List.iter strongred examples