exception UnsupportedAst of string

let raise_unsupported_ast (type a) (f : Format.formatter -> a -> unit) (x : a) =
  let buf = Buffer.create 16 in
  let fmt = Format.formatter_of_buffer buf in
  f fmt x;
  Format.pp_print_flush fmt ();
  raise (UnsupportedAst (Buffer.contents buf))

let raise_unsupported_expr = raise_unsupported_ast (Printast.expression 0)

type prim = [ `Add ] [@@deriving show { with_path = false }]

type expr =
  | App of expr * expr
  | Int of int
  | Lam of string * expr
  | Prim of prim
  | Var of string
[@@deriving show { with_path = false }]

let rec expr_of_parsetree (expr : Parsetree.expression) : expr =
  match expr.pexp_desc with
  | Pexp_apply (f, xs) ->
      List.fold_left
        (fun f -> function
          | Asttypes.Nolabel, x -> App (f, expr_of_parsetree x)
          | _ -> raise_unsupported_expr expr)
        (expr_of_parsetree f) xs
  | Pexp_constant (Pconst_integer (n, None)) -> Int (int_of_string n)
  | Pexp_fun (Nolabel, None, arg_pat, body) -> (
      match arg_pat.ppat_desc with
      | Ppat_var { txt = name; _ } -> Lam (name, expr_of_parsetree body)
      | _ -> raise_unsupported_expr expr)
  | Pexp_ident { txt = Lident "+"; _ } -> Prim `Add
  | Pexp_ident { txt = Lident name; _ } -> Var name
  | _ -> raise_unsupported_expr expr

let rec parsetree_of_expr : expr -> Parsetree.expression =
  let expression (pexp_desc : Parsetree.expression_desc) : Parsetree.expression
      =
    {
      pexp_desc;
      pexp_loc = Location.none;
      pexp_loc_stack = [];
      pexp_attributes = [];
    }
  and ident (name : string) : Longident.t Asttypes.loc =
    { txt = Lident name; loc = Location.none }
  and pattern (ppat_desc : Parsetree.pattern_desc) : Parsetree.pattern =
    {
      ppat_desc;
      ppat_loc = Location.none;
      ppat_loc_stack = [];
      ppat_attributes = [];
    }
  and var (name : string) : string Asttypes.loc =
    { txt = name; loc = Location.none }
  in

  let list_of_apps expr =
    let rec loop = function
      | App (f, x) ->
          let f, xs = loop f in
          (f, x :: xs)
      | f -> (f, [])
    in
    let f, xs = loop expr in
    (f, List.rev xs)
  in

  function
  | App (_, _) as expr ->
      let f, xs = list_of_apps expr in
      let xs = List.map (fun x -> (Asttypes.Nolabel, parsetree_of_expr x)) xs in
      expression (Pexp_apply (parsetree_of_expr f, xs))
  | Int n -> expression (Pexp_constant (Pconst_integer (Int.to_string n, None)))
  | Lam (x, b) ->
      expression
        (Pexp_fun
           (Nolabel, None, pattern (Ppat_var (var x)), parsetree_of_expr b))
  | Prim `Add -> expression (Pexp_ident (ident "+"))
  | Var name -> expression (Pexp_ident (ident name))