Source file HelpersBase.ml

(*
 * Generic Transformers PPX syntax extension.
 * Copyright (C) 2016-2021
 *   Dmitrii Kosarev aka Kakadu
 * St.Petersburg State University, JetBrains Research
 *)

let id x = x

let not_implemented fmt =
  Printf.ksprintf (Ppxlib.Location.raise_errorf "%s are not yet implemented") fmt
;;

module List = struct
  include Base.List

  let split3 xs =
    List.fold_right
      (fun (a, b, c) (ac, bc, cc) -> a :: ac, b :: bc, c :: cc)
      xs
      ([], [], [])
  ;;

  let filter_map ~f xs =
    List.fold_right
      (fun x acc ->
        match f x with
        | Some v -> v :: acc
        | None -> acc)
      xs
      []
  ;;

  let last_exn xs = List.hd @@ List.rev xs
  let pp ~f xs = Printf.sprintf "[ %s ]" (String.concat "; " @@ List.map f xs)

  let fold_left0 f = function
    | [] -> failwith "wrong argument of fold_left0"
    | h :: tl -> fold_left ~f ~init:h tl
  ;;

  let concat_map ~f xs = List.concat @@ List.map f xs

  let empty = function
    | [] -> true
    | _ -> false
  ;;
end

open Ppxlib

let compare_core_type a b =
  String.compare
    (Format.asprintf "%a" Pprintast.core_type a)
    (Format.asprintf "%a" Pprintast.core_type b)
;;

let visit_typedecl
    ~loc
    ?(onrecord = fun _ -> not_implemented ~loc "record types")
    ?(onmanifest = fun _ -> not_implemented ~loc "manifest")
    ?(onvariant = fun _ -> not_implemented ~loc "algebraic types")
    ?(onabstract = fun _ -> not_implemented ~loc "abstract types without manifest")
    ?(onopen = fun () -> not_implemented ~loc "open types")
    tdecl
  =
  match tdecl.ptype_kind with
  | Ptype_record r -> onrecord r
  | Ptype_open -> onopen ()
  | Ptype_variant cds -> onvariant cds
  | Ptype_abstract ->
    (match tdecl.ptype_manifest with
    | None -> onabstract ()
    | Some typ -> onmanifest typ)
;;

let affect_longident ~f = function
  | Lident x -> Lident (f x)
  | Ldot _ as l -> l
  | Lapply (_, _) as l -> l
;;

let rec map_longident ~f = function
  | Lident x -> Lident (f x)
  | Ldot (l, s) -> Ldot (l, f s)
  | Lapply (l, r) -> Lapply (l, map_longident ~f r)
;;

let lident_tail = function
  | Lident _ as l -> l
  | Ldot (_, s) -> Lident s
  | Lapply (_, _) as l -> l
;;

module SS = Stdlib.Set.Make (String)

let vars_from_core_type =
  let rec helper acc typ =
    match typ.ptyp_desc with
    | Ptyp_var s -> SS.add s acc
    | Ptyp_tuple args | Ptyp_constr (_, args) -> List.fold_left args ~init:acc ~f:helper
    | Ptyp_arrow (_, l, r) -> helper (helper acc l) r
    | Ptyp_alias (t, lab) -> SS.remove lab (helper acc t)
    | Ptyp_object (_, _)
    | Ptyp_class (_, _)
    | Ptyp_variant (_, _, _)
    | Ptyp_any
    | Ptyp_poly (_, _)
    | Ptyp_package _ | Ptyp_extension _ -> acc
  in
  fun root -> helper SS.empty root
;;

let%test _ =
  let loc = Location.none in
  [ "a" ] = (vars_from_core_type [%type: 'a list] |> SS.elements)
;;

let%test _ =
  let loc = Location.none in
  [] = (vars_from_core_type [%type: int list] |> SS.elements)
;;

let vars_from_tdecl tdecl =
  let ans =
    match tdecl.ptype_manifest with
    | None -> SS.empty
    | Some typ -> vars_from_core_type typ
  in
  let of_labels ls =
    List.fold_left ~init:SS.empty ls ~f:(fun acc { pld_type } ->
        SS.union acc (vars_from_core_type pld_type))
  in
  let ans2 =
    match tdecl.ptype_kind with
    | Ptype_open | Ptype_abstract -> SS.empty
    | Ptype_record ls -> of_labels ls
    | Ptype_variant cds ->
      List.fold_left cds ~init:SS.empty ~f:(fun acc -> function
        | { pcd_args = Pcstr_tuple ts } ->
          List.fold_left ~init:SS.empty ts ~f:(fun acc x ->
              SS.union acc (vars_from_core_type x))
        | { pcd_args = Pcstr_record ls } -> SS.union acc (of_labels ls))
  in
  SS.union ans2 ans
;;

let map_core_type ?(onconstr = fun _ _ -> None) ~onvar t =
  let rec helper t =
    (* Format.printf "map_core_type,helper `%a`\n%!" Pprintast.core_type t; *)
    match t.ptyp_desc with
    | Ptyp_any -> t
    | Ptyp_var name -> Base.Option.value (onvar name) ~default:t
    | Ptyp_constr (name, args) ->
      (match onconstr name.txt args with
      | None -> { t with ptyp_desc = Ptyp_constr (name, List.map ~f:helper args) }
      | Some t -> t)
    | Ptyp_tuple args -> { t with ptyp_desc = Ptyp_tuple (List.map ~f:helper args) }
    | Ptyp_arrow (lab, from, to_) ->
      { t with ptyp_desc = Ptyp_arrow (lab, helper from, helper to_) }
    | Ptyp_variant (rows, flg, opt) ->
      let rows =
        List.map rows ~f:(fun rf ->
            match rf.prf_desc with
            | Rinherit t -> { rf with prf_desc = Rinherit (helper t) }
            | Rtag (name, flg, ps) ->
              (* Format.printf "got tag `%s`\n%!" name.txt; *)
              let params = List.map ps ~f:helper in
              { rf with prf_desc = Rtag (name, flg, params) })
      in
      { t with ptyp_desc = Ptyp_variant (rows, flg, opt) }
    | _ -> failwith "not implemented"
  in
  helper t
;;

let list_first_some ~f xs =
  List.fold_left xs ~init:None ~f:(function
      | None -> f
      | Some ans -> fun _ -> Some ans)
;;

let is_type_used_in ~tdecl lident =
  let exception Found in
  let rec helper t =
    match t.ptyp_desc with
    | Ptyp_constr ({ txt }, _) when lident = txt -> raise Found
    | Ptyp_object _ -> ()
    | Ptyp_any | Ptyp_var _ | Ptyp_class _ | Ptyp_package _ | Ptyp_extension _ -> ()
    | Ptyp_variant (rfs, _, _) ->
      List.iter rfs ~f:(function
          | { prf_desc = Rtag (_, _, xs) } -> List.iter xs ~f:helper
          | { prf_desc = Rinherit t } -> helper t)
    | Ptyp_poly (_, t) | Ptyp_alias (t, _) -> helper t
    | Ptyp_tuple args | Ptyp_constr (_, args) -> List.iter args ~f:helper
    | Ptyp_arrow (_, l, r) ->
      helper l;
      helper r
  in
  let onrecord = List.iter ~f:(fun { pld_type = t } -> helper t) in
  try
    visit_typedecl
      tdecl
      ~loc:()
      ~onopen:(fun () -> ())
      ~onabstract:(fun _ -> ())
      ~onmanifest:helper
      ~onvariant:
        (List.iter ~f:(function
            | { pcd_args = Pcstr_tuple ls } -> List.iter ~f:helper ls
            | { pcd_args = Pcstr_record ls } -> onrecord ls))
      ~onrecord;
    false
  with
  | Found -> true
;;

let maybe_specialiaze ~what where =
  (* Format.printf "maybe specialize: %a\n%!" Pprintast.structure_item
   *   {pstr_desc=(Pstr_type (Nonrecursive, [what])); pstr_loc=what.ptype_name.loc };
   * List.iter where ~f:(Format.printf "\t%a\n%!" Pprintast.core_type);
   * print_endline "=="; *)
  let myfold ~f ~init xs =
    List.fold_left ~init xs ~f:(function
        | Some r -> fun _ -> Some r
        | None -> f)
  in
  let rec loop t =
    (* Format.printf "loop: %a\n%!" Pprintast.core_type t; *)
    match t.ptyp_desc with
    | Ptyp_constr ({ txt = Lident s }, args) when String.equal s what.ptype_name.txt ->
      (* Format.printf "%s %d\n%!" __FILE__ __LINE__; *)
      Some
        (List.map2_exn what.ptype_params args ~f:(fun (param, _) typ ->
             match param.ptyp_desc with
             | Ptyp_var s -> s, typ
             | _ -> failwith "should not happen"))
    | Ptyp_tuple args | Ptyp_constr (_, args) -> myfold args ~init:None ~f:loop
    | Ptyp_var _ -> None
    (* | Ptyp_record (labs, _) ->
     *   myfold labs ~init:None ~f:(fun (_,e) -> loop e) *)
    | _ -> not_implemented "TODO: maybe_specialize %s" (string_of_core_type t)
  in
  list_first_some ~f:loop where
;;

(* There we iterate over type declaration [where] and check is type [what] is used inside
   If yes, it returns Some substitution of type parameters of what.
*)
let specialize_for_tdecl ~what ~where =
  let loc = where.ptype_name.loc in
  visit_typedecl
    ~loc
    where
    ~onrecord:(fun labs ->
      maybe_specialiaze ~what @@ List.map labs ~f:(fun l -> l.pld_type)
      (* not_implemented ~loc "TODO: record types" *))
      (* ~onmanifest:(fun _ -> not_implemented ~loc "TODO: manifest") *)
    ~onmanifest:(fun t -> maybe_specialiaze ~what [ t ])
    ~onvariant:(fun cstrs ->
      list_first_some cstrs ~f:(fun c ->
          match c.pcd_args with
          | Pcstr_tuple ts -> maybe_specialiaze ~what ts
          | _ -> assert false))
    ~onabstract:(fun _ -> None)
  |> function
  | None -> []
  | Some map ->
    (* Format.printf "Found somethig: %s\n------------\n%!"
     *   (List.map map ~f:(fun (s,_) -> Printf.sprintf "(\"%s\",_)" s)
     *   |> String.concat " "); *)
    map
;;

let with_constr_typ typ ~ok ~fail =
  match typ.ptyp_desc with
  | Ptyp_constr (cid, params) -> ok cid params
  | _ -> fail ()
;;

let constr_of_tuple ?(loc = Location.none) ts =
  let new_lident = Ldot (Lident "GT", Printf.sprintf "tuple%d" @@ List.length ts) in
  let open Ppxlib.Ast_builder.Default in
  ptyp_constr ~loc (Located.mk ~loc new_lident) ts
;;

let using_type ~typename root_type =
  let loc = root_type.ptype_loc in
  let open Ppxlib.Ast_builder.Default in
  (* generation type specification by type declaration *)
  ptyp_constr ~loc (Located.lident ~loc typename) (List.map ~f:fst root_type.ptype_params)
;;

let is_polyvariant typ =
  match typ.ptyp_desc with
  | Ptyp_variant (_, _, _) -> true
  | _ -> false
;;

let is_tuple typ =
  match typ.ptyp_desc with
  | Ptyp_tuple _ts -> true
  | _ -> false
;;

let is_polyvariant_tdecl tdecl =
  let loc = tdecl.ptype_loc in
  visit_typedecl
    ~loc
    tdecl
    ~onopen:(fun () -> false)
    ~onrecord:(fun _ -> false)
    ~onvariant:(fun _ -> false)
    ~onabstract:(fun () -> false)
    ~onmanifest:(fun typ -> is_polyvariant typ)
;;

let is_tuple_tdecl tdecl =
  let loc = tdecl.ptype_loc in
  visit_typedecl
    ~loc
    tdecl
    ~onopen:(fun () -> false)
    ~onrecord:(fun _ -> false)
    ~onvariant:(fun _ -> false)
    ~onabstract:(fun () -> false)
    ~onmanifest:(fun typ -> is_tuple typ)
;;

let is_algebraic_tdecl tdecl =
  let loc = tdecl.ptype_loc in
  visit_typedecl
    ~loc
    tdecl
    ~onopen:(fun () -> false)
    ~onrecord:(fun _ -> false)
    ~onvariant:(fun _ -> true)
    ~onabstract:(fun () -> false)
    ~onmanifest:(fun _typ -> false)
;;

let has_many_constructors_tdecl tdecl =
  let loc = tdecl.ptype_loc in
  visit_typedecl
    ~loc
    tdecl
    ~onopen:(fun () -> false)
    ~onrecord:(fun _ -> false)
    ~onvariant:(fun cs -> Base.Int.( > ) (List.length cs) 1)
    ~onabstract:(fun () -> false)
    ~onmanifest:(fun typ ->
      match typ.ptyp_desc with
      | Ptyp_variant (rf, _, _labels) ->
        (* TODO: we don't take to account labels here *)
        List.length rf > 1
      | _ -> false)
;;

let unfold_tuple t =
  match t.ptyp_desc with
  | Ptyp_tuple ts -> ts
  | _ -> [ t ]
;;

let prepare_patt_match_poly ~loc what rows labels ~onrow ~onlabel ~oninherit =
  let open Ppxlib.Ast_builder.Default in
  let k cs = pexp_match ~loc what cs in
  let rs =
    List.map rows ~f:(function
        | Rtag (lab, _, args) ->
          let args =
            match args with
            | [ t ] -> unfold_tuple t
            | [] -> []
            | _ -> failwith "we don't support conjunction types"
          in
          let names = List.map args ~f:(fun _ -> gen_symbol ~prefix:"_" ()) in
          let lhs =
            ppat_variant ~loc lab.txt
            @@
            match args with
            | [] -> None
            | _ ->
              Some
                (ppat_tuple ~loc
                @@ List.map ~f:(fun s -> ppat_var ~loc (Located.mk ~loc s)) names)
          in
          case ~guard:None ~lhs ~rhs:(onrow lab @@ List.zip_exn names args)
        | Rinherit typ ->
          (match typ.ptyp_desc with
          | Ptyp_constr ({ txt; loc }, ts) ->
            let newname = "subj" in
            let lhs =
              ppat_alias
                ~loc
                (ppat_type ~loc (Located.mk ~loc txt))
                (Located.mk ~loc newname)
            in
            case ~guard:None ~lhs ~rhs:(oninherit ts txt newname)
          | _ -> failwith "this inherit field isn't supported"))
  in
  let ls =
    match labels with
    | None -> []
    | Some ls ->
      List.map ls ~f:(fun lab ->
          let newname = "subj" in
          let lhs =
            ppat_alias
              ~loc
              (ppat_type ~loc (Located.mk ~loc (Lident lab)))
              (Located.mk ~loc newname)
          in
          case ~guard:None ~lhs ~rhs:(onlabel lab newname))
  in
  k @@ rs @ ls
;;

let map_type_param_names ~f ps =
  List.map ps ~f:(fun (t, _) ->
      match t.ptyp_desc with
      | Ptyp_var name -> f name
      | _ -> failwith "bad argument of map_type_param_names")
;;

let notify fmt =
  Format.kasprintf
    (fun s ->
      let _cmd = Printf.sprintf "notify-send \"%s\"" s in
      let (_ : int) = Caml.Sys.command _cmd in
      ())
    fmt
;;

let string_after_a n = Base.Char.(to_int 'a' |> ( + ) n |> of_int_exn |> to_string)

external hash_variant : string -> int = "caml_gt_hash_variant"

(* TODO: Don't use and remove this function *)
let failwiths fmt = Format.kasprintf failwith fmt