extract nested match expressions

match.ml

(* types *)

type erl_expr =
  | ErlInt of int
  | ErlVar of string
  | ErlApp of erl_expr * erl_expr (* F(A) *)
  | ErlMatch of erl_pat * erl_expr (* P = E *)
and erl_pat =
  | ErlPatVar of string

type expr =
  | Int of int
  | Var of string
  | App of expr * expr
  | Case of expr * (pat * expr) list
and pat =
  | PatVar of string

(* pp *)

let erl_pat_pp = function
  | ErlPatVar v -> v

let rec erl_pp = function
  | ErlInt n -> string_of_int n
  | ErlVar v -> v
  | ErlApp (f, a) -> erl_pp f ^ "(" ^ erl_pp a ^ ")"
  | ErlMatch (p, e) -> erl_pat_pp p ^ " = " ^ erl_pp e

let rec erl_pp_list es = List.map erl_pp es |> String.concat ", "

let pat_pp = function
  | PatVar v -> v

let rec pp = function
  | Int n -> string_of_int n
  | Var v -> v
  | App (f, a) -> pp f ^ "(" ^ pp a ^ ")"
  | Case (e, cs) ->
     let cs = List.map (fun (p, e) -> pat_pp p ^ " -> " ^ pp e) cs in
     "case (" ^ pp e ^ ") of " ^ String.concat "; " cs ^ " end"

let rec pp_list es = List.map pp es |> String.concat ", "

(* erl_expr -> expr *)

let pat_of_erl_pat = function
  | ErlPatVar v -> PatVar v

(* idea 1 *)
let rec expr_of_erl_expr expr_of = function
  | ErlInt n -> Int n
  | ErlVar v -> Var v
  | ErlApp (f, a) -> App (expr_of f, expr_of a)
  | ErlMatch _ -> failwith "match expr"
let rec expr_of_erl_expr_1 = function
  | ErlMatch (p, e) ->
     (* だめ *)
     let e = expr_of_erl_expr_1 e in
     Case (e, [(pat_of_erl_pat p, e)])
  | e -> expr_of_erl_expr expr_of_erl_expr_1 e

let rec expr_of_body_1 = function
  | [] -> failwith "body does not have any exprs"
  | [e] -> expr_of_erl_expr_1 e
  | e :: es -> Case (expr_of_erl_expr_1 e, [(PatVar "_", expr_of_body_1 es)])

(* idea 2 *)
let rec expr_of_erl_expr_2 = function
  | ErlMatch (p1, (ErlMatch (_, e2) as e1)) ->
     let rec eval = function
       | ErlMatch (_, e0) -> eval e0
       | e0 -> e0
     in
     let ex = expr_of_erl_expr_2 e1 in
     fun c -> ex (Case (expr_of_erl_expr_2 (eval e2) (Var "dummy"), [(pat_of_erl_pat p1, c)]))
  | ErlMatch (p, e) ->
     fun c -> Case (expr_of_erl_expr_2 e (Var "dummy"), [(pat_of_erl_pat p, c)])
  | e -> fun _ -> expr_of_erl_expr (fun e -> expr_of_erl_expr_2 e (Var "dummy")) e

let rec expr_of_body_2 = function
  | [] -> failwith "body does not have any exprs"
  | [ErlMatch (p1, (ErlMatch (_, e2) as e1))] ->
     let rec eval = function
       | ErlMatch (_, e0) -> eval e0
       | e0 -> e0
     in
     let ex = expr_of_erl_expr_2 e1 in
     let value = expr_of_erl_expr_2 (eval e2) (Var "dummy") in
     ex (Case (value, [(pat_of_erl_pat p1, value)]))
  | [ErlMatch (p, e)] ->
     let value = expr_of_erl_expr_2 e (Var "dummy") in
     Case (value, [(pat_of_erl_pat p, value)])
  | [e] -> expr_of_erl_expr (fun e -> expr_of_erl_expr_2 e (Var "dummy")) e
  | ErlMatch (p1, (ErlMatch (_, e2) as e1)) :: es ->
     let rec eval = function
       | ErlMatch (_, e0) -> eval e0
       | e0 -> e0
     in
     let ex = expr_of_erl_expr_2 e1 in
     let value = expr_of_erl_expr_2 (eval e2) (Var "dummy") in
     let next = expr_of_body_2 es in
     ex (Case (value, [(pat_of_erl_pat p1, next)]))
  | ErlMatch (p, e) :: es ->
     let value = expr_of_erl_expr_2 e (Var "dummy") in
     let next = expr_of_body_2 es in
     Case (value, [(pat_of_erl_pat p, next)])
  | e :: es ->
     Case (expr_of_erl_expr (fun e -> expr_of_erl_expr_2 e (Var "dummy")) e, [(PatVar "_", expr_of_body_2 es)])

(* idea 3 *)
let rec extract_match_3 e =
  let rec extract_match_3' acc is_top = function
    | ErlMatch (p, e) ->
       let (acc', r) = extract_match_3' acc false e in
       (ErlMatch (p, r) :: acc', r)
    | ErlApp (f, a) ->
       let (accf, rf) = extract_match_3' acc false f in
       let (acca, ra) = extract_match_3' accf false a in
       if is_top then (ErlApp (rf, ra) :: acca, ErlApp (rf, ra)) else (acca, ErlApp (rf, ra))
    | e -> if is_top then (e :: acc, e) else (acc, e)
  in
  let (l, _) = extract_match_3' [] true e in
  List.rev l

let rec expr_of_body_3 es =
  let es' = List.map extract_match_3 es |> List.flatten in
  (* ここではネストしている match 式はないことが保証されている *)
  let rec expr_of = function
    | [] -> failwith "body does not have any exprs"
    | [ErlInt n] -> Int n
    | [ErlVar v] -> Var v
    | [ErlApp (f, a)] ->
       (* f, a の中には match はない *)
       App (expr_of [f], expr_of [a])
    | [ErlMatch (p, e)] ->
       (* e の中には match はない *)
       let e = expr_of [e] in
       Case (e, [(pat_of_erl_pat p, e)])
    | ErlMatch (p, e) :: es ->
       (* e の中には match はない *)
       let e = expr_of [e] in
       let rest = expr_of es in
       Case (e, [(pat_of_erl_pat p, rest)])
    | e :: es -> Case (expr_of [e], [(PatVar "_", expr_of es)]) (* fialyzer だと let *)
  in
  expr_of es'

(* test *)

let () =
  (* print_endline "case1";
   * let case1 = ErlMatch (ErlPatVar "A", ErlMatch (ErlPatVar "B", ErlVar "C")) in
   * print_endline (erl_pp case1);
   * print_endline (pp (expr_of_erl_expr_1 case1));
   * print_endline (pp (expr_of_erl_expr_2 case1 (Var "dummy"))); *)

  print_endline "case2";
  let case2 = [ ErlMatch (ErlPatVar "A", ErlMatch (ErlPatVar "B", ErlVar "C"));
                ErlVar "A"]
  in
  print_endline (erl_pp_list case2);
  case2 |> List.map extract_match_3 |> List.flatten |> erl_pp_list |> print_endline;
  (* print_endline (pp (expr_of_body_1 case2));
   * print_endline (pp (expr_of_body_2 case2)); *)
  print_endline (pp (expr_of_body_3 case2));

  print_endline "case3";
  let case3 = [ ErlMatch (ErlPatVar "A", ErlApp (ErlVar "F", ErlMatch (ErlPatVar "B", ErlVar "C")));
                ErlApp (ErlVar "G", ErlMatch (ErlPatVar "D", ErlVar "E")) ]
  in
  print_endline (erl_pp_list case3);
  case3 |> List.map extract_match_3 |> List.flatten |> erl_pp_list |> print_endline;
  (* input: A = F(B = C), G(D = E). *)
  (* mid: _C = C, B = _C, A = F(_C), _E = E, D = _E, G(_E) *)
  (* or: B = C, A = F(C), D = E, G(E) *)
  (*   C が副作用を起こすような関数だった場合はだめだけど型検査には関係ないかも *)
  (* expected:
   * case C of
   *   B -> case F(C) of
   *          A -> case E of
   *                 D -> G(E)
   *               end
   *        end
   * end
   *)
  (* print_endline (pp (expr_of_body_1 case3));
   * print_endline (pp (expr_of_body_2 case3)); *)
  print_endline (pp (expr_of_body_3 case3));

  print_endline "case4";
  let case4 = [ ErlMatch (ErlPatVar "A", ErlApp (ErlVar "f", ErlMatch (ErlPatVar "B", ErlMatch (ErlPatVar "C", ErlApp (ErlVar "g", ErlVar "D")))));
                ErlApp (ErlMatch (ErlPatVar "E", ErlVar "F"), ErlMatch (ErlPatVar "G", ErlVar "H")) ]
  in
  case4 |> erl_pp_list |> print_endline;
  case4 |> List.map extract_match_3 |> List.flatten |> erl_pp_list |> print_endline;
  case4 |> expr_of_body_3 |> pp |> print_endline;

  (* input: A = f(B = C = g(D)), (E = F)(G = H) *)
  (* mid: C = g(D), B = g(D), A = f(g(D)), E = F, G = H, F(H) *)