(*********************************************)
(* (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 rec find_idx_exn x = function
  | [] -> raise Not_found
  | (y::ys) -> if x = y then 0 else 1 + find_idx_exn x ys

let rec map_rev f xs =
  let rec iter acc = function
    | [] -> acc
    | (x::xs) -> iter (f x :: acc) xs in
  iter [] xs

type symbol = string * int
let show_symbol (s, _) = s
let symbol : string -> symbol =
  let id = ref 0 in
  let gensym c : symbol =
    let newid = !id in
    id := !id + 1;
    (c, newid)
  in
  gensym


(**************************************)
(* (Extended) Untyped Lambda Calculus *)
(**************************************)
module Lambda = struct
    type t = Var of symbol | App of t * t | Abs of symbol * t
             | If of t * t * t | True | False
             | Constr of t constr | Case of t * (symbol constr * t) list
             | Fix of symbol * t
     and 'a constr = symbol * 'a list

    let rec show m =
      let show_branch ((x, args), m) =
        "(" ^ show_symbol x ^ "(" ^ concat_with "," show_symbol args ^ ") → "
        ^ 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 ^ ")"
      | If (m1, m2, m3) -> "if " ^ show m1
                          ^ " then " ^ show m2
                          ^ " else " ^ show m3
      | True -> "True" | False -> "False"
      | Constr (x, ms) -> show_symbol x ^ "(" ^ concat_with ", " show ms ^ ")"
      | Case (m, bs) -> "(case " ^ show m ^ " of "
                       ^ concat_with " " show_branch bs ^ ")"
      | Fix (x, m) -> "fix(λ" ^ show_symbol x ^ "." ^ show m ^ ")"
    let print = print_with show

    (* Constants *)
    let none = symbol "None"
    let some = symbol "Some"

    (* 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.((λy.y) x)) *)
    let ex_m1 =
      let x = symbol "x" in
      let y = symbol "y" in
      Abs (x, App (Abs (y, Var y), Var x))

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

    (* (λc. case c of (Some(x) → x) (None → c)) Some(s(z))) *)
    let ex_case_some =
      let c = symbol "c" in
      let x = symbol "x" in
      App (Abs (c, Case (Var c, [((some, [x]), Var x); ((none, []), Var c)])),
           Constr (some, [peano_of_int 1]))

    (* (λc. case c of (Triple(x,y,z) → y)) Triple(1,2,3)) *)
    let ex_case_tuple =
      let c = symbol "c" in
      let x = symbol "x" in
      let y = symbol "y" in
      let z = symbol "z" in
      let triple = symbol "triple" in
      App (Abs (c, Case (Var c, [((triple, [x;y;z]), Var y)])),
           Constr (triple, List.map peano_of_int [1;2;3]))

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

    (* fix(λf.λg.λc. case c of (s(x) → g (f g x)) (z → z)) (λy.s(s(s(y)))) s(s(s(z))) *)
    let ex_fix_scale =
      let f = symbol "f" in
      let g = symbol "g" in
      let c = symbol "c" in
      let x = symbol "x" in
      let y = symbol "y" in
      App (App (Fix (f, Abs (g, Abs (c, Case (Var c,
                 [((s, [x]), App (Var g, (App (App (Var f, Var g), Var x))));
                  ((z, []), peano_of_int 0)])))),
                Abs (y, peano_add 3 (Var y))),
           peano_of_int 3)
end


(**************************************************)
(* Untyped lambda calculus with De Bruijn indices *)
(**************************************************)
module DBILambda = struct
    type dbi_symbol = int * symbol
    type t = Var of dbi_symbol | App of t * t | Abs of symbol * t
             | If of t * t * t | True | False
             | Constr of t constr | Case of t * (symbol constr * t) list
             | Fix of symbol * t
     and 'a constr = symbol * 'a list

    let dbi dbis x = (find_idx_exn x dbis, x)

    let output_dbi = false
    let show_dbi_symbol (n, x) =
      if output_dbi then string_of_int n else show_symbol x
    let show_dbi_param x =
      if output_dbi then "" else show_symbol x

    let rec show m =
      match m with
      | Var x -> show_dbi_symbol x
      | App (m1, m2) -> "(" ^ show m1 ^ " " ^ show m2 ^ ")"
      | Abs (x, m) -> "(λ" ^ show_dbi_param x ^ "." ^ show m ^ ")"
      | If (m1, m2, m3) -> "if " ^ show m1
                          ^ " then " ^ show m2
                          ^ " else " ^ show m3
      | True -> "True" | False -> "False"
      | Constr (x, ms) -> show_symbol x ^ "(" ^ concat_with ", " show ms ^ ")"
      | Case (m, bs) -> "(case " ^ show m ^ " of "
                       ^ concat_with " " show_branch bs ^ ")"
      | Fix (x, m) -> "fix(λ" ^ show_symbol x ^ "." ^ show m ^ ")"
    and show_branch ((x, args), m) =
      "(" ^ show_symbol x ^ "(" ^ concat_with "," show_symbol args ^ ") → "
      ^ show m ^ ")"
    let print = print_with show

    let of_lambda =
      let rec of_lambda dbis = function
        | Lambda.Var x -> let (n, x) = dbi dbis x in Var (n, x)
        | Lambda.App (m1, m2) -> App (of_lambda dbis m1, of_lambda dbis m2)
        | Lambda.Abs (x, m) -> Abs (x, of_lambda (x :: dbis) m)
        | Lambda.If (m1, m2, m3) -> If (of_lambda dbis m1,
                                       of_lambda dbis m2, of_lambda dbis m3)
        | Lambda.True -> True | Lambda.False -> False
        | Lambda.Constr (x, ms) -> Constr (x, List.map (of_lambda dbis) ms)
        | Lambda.Case (m, bs) -> Case (of_lambda dbis m,
                                      List.map (trans_br dbis) bs)
        | Lambda.Fix (x, m) -> Fix (x, of_lambda (x :: dbis) m)
      and trans_br dbis ((x, args), m) =
        let dbis = List.rev args @ dbis in
        ((x, args), of_lambda dbis m) in
      of_lambda []
end


(*******************)
(* Krivine Machine *)
(*******************)
module KM = struct
    open DBILambda

    type st_elm = Clos of DBILambda.t * stack
                | IfCont of DBILambda.t * DBILambda.t
                | CaseCont of (symbol DBILambda.constr * DBILambda.t) list * stack
                | FixClos of symbol * DBILambda.t * stack
    and stack = st_elm list

    type state = DBILambda.t * stack * stack

    let reduce m =
      let rec reduce (st : state) : DBILambda.t =
        match st with
        (* Pure lambda calculus *)
        | (Var (0, _), s, Clos (m, e') :: e) -> reduce (m, s, e')
        | (Var (0, _), s, FixClos (f, m, e') :: e) ->
           reduce (m, s, FixClos (f, m, e') :: e')
        | (Var (n, x), s, _ :: e) -> reduce (Var (n-1, x), s, e)
        | (App (m1, m2), s, e) -> reduce (m1, Clos (m2, e) :: s, e)
        | (Abs (_, m), c :: s, e) -> reduce (m, s, c :: e)
        (* Conditionals *)
        | (If (m1, m2, m3), s, e) -> reduce (m1, IfCont (m2, m3) :: s, e)
        | (True, IfCont (m2, m3) :: s, e) -> reduce (m2, s, e)
        | (False, IfCont (m2, m3) :: s, e) -> reduce (m3, s, e)
        (* Case expressions (+ constructors) *)
        | (Case (m, bs), s, e) -> reduce (m, CaseCont (bs, e) :: s, e)
        | (Constr (x, ms), CaseCont (((x', args), m) :: bs, e') :: s, e)
             when x == x' && List.length ms == List.length args ->
           reduce (List.fold_left (fun m x -> Abs (x, m)) m args,
                   map_rev (fun m -> Clos (m, e)) ms @ s, e')
        | (Constr (x, ms), CaseCont (_ :: bs, e') :: s, e) ->
           reduce (Constr (x, ms), CaseCont (bs, e') :: s, e)
        | (Constr (x, ms), s, e) ->
           Constr (x, List.map (fun m -> reduce (m, s, e)) ms)
        (* Fixpoints *)
        | (Fix (x, m), s, e) -> reduce (m, s, FixClos (x, m, e) :: e)
        (* Termination checks *)
        | (m, [], []) -> m
        | (_, _, _) -> m in
      reduce (m, [], [])
end


let dbi_and_red m =
  let dbi_m = DBILambda.of_lambda m in
  print_endline ("# Lambda term:\n    " ^ DBILambda.show dbi_m);
  let red_m = KM.reduce dbi_m in
  print_endline ("# Reduced term:\n    " ^ DBILambda.show red_m);
  print_endline "----------------------------------------------------------\n"

let () =
  let open Lambda in
  List.iter dbi_and_red [ex_m1; ex_m2; ex_case_some;
                         ex_case_tuple; ex_fixpt_mul2; ex_fix_scale]