Source file madcast.ml

open Ppxlib
open Parsetree
open Asttypes
open Longident
open Ast_helper
open Location

let mkpatvar i = Pat.var (Parsetree_utils.mknoloc ("c"^(string_of_int i)))
let mkident i = Exp.ident (Parsetree_utils.mknoloc (Lident ("c"^(string_of_int i))))

(* FIXME: proper location *)
let loc = Location.none

(* ============================= [ Base types ] ============================= *)

let () =
  [ ( "bool -> float",
      [%type: bool], [%type: float],
      [%expr function
          | false -> 0.
          | true -> 1.] );

    ( "bool -> int",
      [%type: bool], [%type: int],
      [%expr function
          | false -> 0
          | true -> 1] );

    ( "bool -> string",
      [%type: bool], [%type: string],
      [%expr string_of_bool] );

    ( "char -> int",
      [%type: char], [%type: int],
      [%expr int_of_char] );

    ( "char -> string" ,
      [%type: char], [%type: string],
      [%expr String.make 1] );

    ( "float -> string",
      [%type: float], [%type: string],
      [%expr string_of_float] );

    ( "int -> bool",
      [%type: int], [%type: bool],
      [%expr function
          | 0 -> false
          | 1 -> true
          | _ -> failwith "madcast: int -> bool"] );

    ( "int -> char",
      [%type: int], [%type: char],
      [%expr fun i ->
          try
            char_of_int i
          with
            Failure _ -> failwith "madcast: int -> char"] );

    ( "int -> float",
      [%type: int], [%type: float],
      [%expr float_of_int] );

    ( "int -> string",
      [%type: int], [%type: string],
      [%expr string_of_int] );

    ( "string -> bool",
      [%type: string], [%type: bool],
      [%expr fun s ->
          try
            bool_of_string s
          with
            Failure _ -> failwith "madcast: string -> bool"] );

    ( "string -> char",
      [%type: string], [%type: char],
      [%expr fun s ->
          if String.length s = 1 then
            s.[0]
          else
            failwith "madcast: string -> char"] );

    ( "string -> float",
      [%type: string], [%type: float],
      [%expr fun s ->
          try
            float_of_string s
          with
            Failure _ -> failwith "madcast: string -> float"] );

    ( "string -> int",
      [%type: string], [%type: int],
      [%expr fun s ->
          try
            int_of_string s
          with
            Failure _ -> failwith "madcast: string -> int"] ) ]
  |>
    List.iter
      (fun (name, itype, otype, expr) ->
        let matcher (itype', otype') =
          if Parsetree_utils.equal_core_type itype itype'
             && Parsetree_utils.equal_core_type otype otype'
          then Some []
          else None
        in
        let builder casts =
          assert (casts = []);
          expr
        in
        RuleSet.register (Rule.make ~name ~matcher ~builder ()))

(* ============================== [ Options ] =============================== *)

let () =
  let name = "'a option -> 'b option" in
  let matcher = function
    | [%type: [%t? itype] option], [%type: [%t? otype] option] ->
       Some [itype, otype]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr function
        | None -> None
        | Some x -> Some ([%e List.hd casts] x)]
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

let () =
  let name = "'a -> 'b option" in
  let matcher = function
    | itype, [%type: [%t? otype] option] ->
       Some [itype, otype]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr fun x -> Some ([%e List.hd casts] x)]
  in
  RuleSet.(register
    ~applies_after:[lookup "'a option -> 'b option"]
    (Rule.make ~name ~matcher ~builder ()))

let () =
  let name = "'a option -> 'b" in
  let matcher = function
    | [%type: [%t? itype] option], otype ->
       Some [itype, otype]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr function
        | None -> failwith "madcast: 'a option -> 'b"
        | Some x -> [%e List.hd casts] x]
  in
  RuleSet.(register
    ~applies_after:[lookup "'a option -> 'b option"]
    (Rule.make ~name ~matcher ~builder ()))

(* =============================== [ Arrays ] =============================== *)

let () =
  let name = "'a array -> 'b array" in
  let matcher = function
    | [%type: [%t? itype] array], [%type: [%t? otype] array] ->
       Some [itype, otype]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr Array.map [%e List.hd casts]]
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

let () =
  let name = "'a -> 'b array" in
  let matcher = function
    | itype, [%type: [%t? otype] array] ->
       Some [itype, otype]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr fun x -> [|[%e List.hd casts] x|]]
  in
  RuleSet.(register
    ~applies_after:[lookup "'a array -> 'b array"]
    (Rule.make ~name ~matcher ~builder ()))

let () =
  let name = "'a array -> 'b" in
  let matcher = function
    | [%type: [%t? itype] array], otype ->
       Some [itype, otype]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr fun a ->
        if Array.length a = 1 then
          [%e List.hd casts] a.(0)
        else
          failwith "madcast: 'a array -> 'b"]
  in
  RuleSet.(register
    ~applies_after:[lookup "'a array -> 'b array"; lookup "'a -> 'b array"]
    (Rule.make ~name ~matcher ~builder ()))

let () =
  let name = "<tuple> -> 'b array" in
  let matcher = function
    | {ptyp_desc=Ptyp_tuple itypes; _}, [%type: [%t? otype] array] ->
       Some (List.map (fun itype -> (itype, otype)) itypes)
    | _ -> None
  in
  let builder casts =
    (* fun (c0,...ck) -> [|cast0 c0; ... castk ck|] *)
    Exp.fun_
      Nolabel None
      (Pat.tuple (List.mapi (fun i _ -> mkpatvar i) casts))
      (Exp.array (List.mapi (fun i cast -> Exp.apply cast [Nolabel, mkident i]) casts))
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

let () =
  let name = "'a array -> <tuple>" in
  let matcher = function
    | [%type: [%t? itype] array], {ptyp_desc=Ptyp_tuple otypes; _} ->
       Some (List.map (fun otype -> (itype, otype)) otypes)
    | _ -> None
  in
  let builder casts =
    (* function
       | [|c0;...ck|] -> (cast0 c0, ... castk ck)
       | _ -> failwith ... *)
    Exp.function_
      [ Exp.case
          (Pat.array (List.mapi (fun i _ -> mkpatvar i) casts))
          (Exp.tuple (List.mapi (fun i cast -> Exp.apply cast [Nolabel, mkident i]) casts)) ;
        Exp.case
          (Pat.any ())
          [%expr failwith "madcast: 'a array -> <tuple>"] ]
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

let () =
  let name = "<tuple> array -> 'a array" in
  let matcher = function
    | [%type: [%t? {ptyp_desc=Ptyp_tuple itypes; _}] array], [%type: [%t? otype] array] ->
       Some [Typ.tuple itypes, [%type: [%t otype] array]]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr fun a ->
        Array.map [%e List.hd casts] a
        |> Array.to_list
        |> Array.concat]
  in
  RuleSet.(register
    ~applies_before:[lookup "'a -> 'b array"; lookup "'a array -> 'b"]
    (Rule.make ~name ~matcher ~builder ()))

let () =
  let name = "'a array -> <tuple> array" in
  let matcher = function
    | [%type: [%t? itype] array], [%type: [%t? {ptyp_desc=Ptyp_tuple otypes; _}] array] ->
       Some (List.map (fun otype -> (itype, otype)) otypes)
    | _ -> None
  in
  let builder casts =
    let l = List.length casts in
    let exp_int n = Exp.constant (Const.int n) in
    [%expr fun a ->
        if Array.length a mod [%e exp_int l] <> 0 then
          failwith "madcast: 'a array -> <tuple> array"
        else
          Array.init (Array.length a / [%e exp_int l])
            (fun i ->
              [%e Exp.tuple
                  (List.mapi
                     (fun j cast ->
                       [%expr [%e cast] a.([%e exp_int j] + i * [%e exp_int l])])
                     casts)])]
  in
  RuleSet.(register
    ~applies_before:[lookup "'a -> 'b array"; lookup "'a array -> 'b"]
    ~applies_after:[lookup "<tuple> array -> 'a array"]
    (Rule.make ~name ~matcher ~builder ()))

(* =============================== [ Lists ] ================================ *)
(* using the rules for arrays *)

let () =
  let name = "'a list -> 'a array -> 'b" in
  let matcher = function
    | [%type: [%t? itype] list], otype ->
       Some [[%type: [%t itype] array], otype]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr fun l -> Array.of_list l |> [%e List.hd casts]]
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

let () =
  let name = "'a -> 'b array -> 'b list" in
  let matcher = function
    | itype, [%type: [%t? otype] list] ->
       Some [itype, [%type: [%t otype] array]]
    | _ -> None
  in
  let builder casts =
    assert (List.length casts = 1);
    [%expr fun x -> [%e List.hd casts] x |> Array.to_list]
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

(* =============================== [ Tuples ] =============================== *)

let () =
  let name = "<tuple> -> <tuple>" in
  let matcher = function
    | {ptyp_desc=Ptyp_tuple itypes; _} , {ptyp_desc=Ptyp_tuple otypes; _}
         when List.length itypes = List.length otypes ->
       Some (List.combine itypes otypes)
    | _ -> None
  in
  let builder casts =
    (* fun (c0,...ck) -> (cast0 c0, ... castk ck) *)
    Exp.fun_
      Nolabel None
      (Pat.tuple (List.mapi (fun i _ -> mkpatvar i) casts))
      (Exp.tuple (List.mapi (fun i cast -> Exp.apply cast [Nolabel, mkident i]) casts))
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

(* ============================= [ Functions ] ============================== *)

let () =
  let name = "('a -> 'b) -> ('c -> 'd)" in
  let matcher = function
    | [%type: [%t? iitype] -> [%t? iotype]], [%type: [%t? oitype] -> [%t? ootype]] ->
       Some [(oitype, iitype); (iotype, ootype)]
    | _ -> None
  in
  let builder = function
    | [icast; ocast] ->
       [%expr fun f x -> x |> [%e icast] |> f |> [%e ocast]]
    | _ -> assert false
  in
  RuleSet.register (Rule.make ~name ~matcher ~builder ())

let () =
  let name = "currying" in
  let matcher (itype, otype) =
    match itype with
    | [%type: [%t? {ptyp_desc=Ptyp_tuple iitypes; _}] -> [%t? iotype]] ->
       (
         let rec matcher = function
           | ([], ootype) -> [(iotype, ootype)] (* this is the right order *)
           | (iitype :: iitypes, [%type: [%t? oitype] -> [%t? ootype]]) ->
              (oitype, iitype) :: matcher (iitypes, ootype)
           | _ -> failwith "matcher"
         in
         try Some (matcher (iitypes, otype))
         with Failure _ -> None
       )
    | _ -> None
  in
  let builder casts =
    let ocast = ExtList.ft casts in
    let icasts = ExtList.bd casts in
    [%expr fun f ->
        [%e ExtList.foldi_right
              (* imbricated functions *)
              (fun i _ exp ->
                Exp.fun_ Nolabel None (mkpatvar i) exp)
              icasts
              (
                (* the body of the function *)
                Exp.apply ocast [Nolabel,
                  Exp.apply [%expr f] [Nolabel,
                    Exp.tuple (
                      List.mapi
                        (fun i icast ->
                          Exp.apply icast [Nolabel, mkident i])
                        icasts
                  )]]
              )]]
  in
  RuleSet.register
    ~applies_after:[RuleSet.lookup "('a -> 'b) -> ('c -> 'd)"]
    (Rule.make ~name ~matcher ~builder ())

let () =
  let name = "uncurrying" in
  let matcher (itype, otype) =
    match otype with
    | [%type: [%t? {ptyp_desc=Ptyp_tuple oitypes; _}] -> [%t? ootype]] ->
       (
         let rec matcher = function
           | (iotype, []) -> [(iotype, ootype)] (* this is the right order *)
           | ([%type: [%t? iitype] -> [%t? iotype]], oitype :: ootypes) ->
              (oitype, iitype) :: matcher (iotype, ootypes)
           | _ -> failwith "matcher"
         in
         try Some (matcher (itype, oitypes))
         with Failure _ -> None
       )
    | _ -> None
  in
  let builder casts =
    let ocast = ExtList.ft casts in
    let icasts = ExtList.bd casts in
    [%expr fun f ->
        [%e Exp.fun_ Nolabel None
              (Pat.tuple (List.mapi (fun i _ -> mkpatvar i) icasts))
              (Exp.apply ocast [Nolabel, (Exp.apply [%expr f]
                 (List.mapi (fun i icast -> (Nolabel, Exp.apply icast [Nolabel, mkident i])) icasts))])]]
  in
  RuleSet.register
    ~applies_after:[RuleSet.lookup "('a -> 'b) -> ('c -> 'd)"]
    (Rule.make ~name ~matcher ~builder ())





(* ======================= [ And now, the main loop ] ======================= *)

let rec reverse_possibles = function
  (* changes a list of possibilities in possibilities of lists *)
  | [] -> [[]]
  | possible_heads :: tail_of_possibles ->
     List.map
       (fun possible_tail ->
         List.map
           (fun possible_head ->
             possible_head :: possible_tail)
           possible_heads)
       (reverse_possibles tail_of_possibles)
     |> List.flatten

let rec find_caster (itype, otype) : Parsetree.expression list =
  RuleSet.fold_by_priority
    (fun rules -> function
      | [] ->
         (* Empty means that the stronger priorities have found
            nothing. We go through all the rules at our priority,
            apply them and see which ones did succeed. *)
         List.fold_left
           (fun casts rule ->
             match Rule.match_ rule (itype, otype) with
             | None -> (* the rule found nothing *) casts
             | Some premises ->
                (
                  List.map find_caster premises
                  |> reverse_possibles
                  |> List.map
                       (fun premises ->
                         Rule.build_ rule premises)
                ) @ casts)
           []
           rules
      | _ as casts ->
         (* Non-empty means that the previous priorities have found
            something already, so we let that and do nothing. *)
         casts)
    []



(* ============================== [ Frontend ] ============================== *)

exception NoCastFound
exception SeveralCastsFound

let find_caster itype otype =
  match find_caster (itype, otype) with
  | [cast] -> cast
  | [] -> raise NoCastFound
  | _ -> raise SeveralCastsFound


let split_arrow = function
  | [%type: [%t? itype] -> [%t? otype]] -> itype, otype
  | _ -> invalid_arg "split_arrow"


let annotate expr ty =
  [%expr let (cast : [%t ty]) = [%e expr] in cast]


let derive ty =
  (* We ask [find_caster] to derive expressions for ty = itype -> otype. We then
     annotate them with that type where type variables are universally
     quantified. Since this can syntactically only happen in a let, we
     return something like:

         let cast : [vars]. [itype -> otype] = [expr] in cast
  *)
  let loc = ty.Parsetree.ptyp_loc in
  try
    let itype, otype = split_arrow ty in
    let cast = find_caster itype otype in
    annotate cast (Parsetree_utils.universal_closure_of_core_type ty)
  with
  | Invalid_argument msg when msg = "split_arrow" ->
      Ppx_deriving.(raise_errorf ~loc
                      "Expected an arrow type, got %s"
                      (string_of_core_type ty))
  | NoCastFound ->
    Ppx_deriving.(raise_errorf ~loc
                    "No cast found for %s"
                    (string_of_core_type ty))
  | SeveralCastsFound ->
    Ppx_deriving.(raise_errorf ~loc
                    "Several casts found for %s"
                    (string_of_core_type ty))