(*pp camlp4o pa_extend.cmo q_MLast.cmo *)

(* File: pa_sexp_conv.ml

    Copyright (C) 2005-

      Jane Street Holding, LLC
      Author: Markus Mottl
      email: mmottl@janestcapital.com
      WWW: http://www.janestcapital.com/ocaml

   This file is derived from file "pa_tywith.ml" of version 0.45 of the
   library "Tywith".

   Tywith is Copyright (C) 2004, 2005 by

      Martin Sandin  <msandin@hotmail.com>

   This library is free software; you can redistribute it and/or
   modify it under the terms of the GNU Lesser General Public
   License as published by the Free Software Foundation; either
   version 2 of the License, or (at your option) any later version.

   This library is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public
   License along with this library; if not, write to the Free Software
   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*)

(* Pa_sexp_conv: Preprocessing Module for Automated S-expression Conversions *)

open MLast
open Printf

(* Utility functions *)

let snd4 (_, x, _, _) = x

let ( *** ) f g x = f (g x)

let both fa fb (a, b) = fa a, fb b

let list_mapi f lst =
  let rec aux ix = function
    | [] -> []
    | x :: xs -> f ix x :: aux (ix + 1) xs in
  aux 1 lst

let rec list_rev_map_append f l1 l2 =
  match l1 with
  | [] -> l2
  | h :: t -> list_rev_map_append f t (f h :: l2)


(* Module/File path management *)

(* Reference storing the path to the currently preprocessed module *)
let conv_path_ref : string option ref = ref None

(* Get path to the currently preprocessed module *)
let get_conv_path () =
  match !conv_path_ref with
  | None -> failwith "Pa_sexp_conv: path not set";
  | Some conv_path -> conv_path

(* Set path to the currently preprocessed module *)
let set_conv_path conv_path =
  if !conv_path_ref <> None then failwith "Pa_sexp_conv: module name set twice";
  conv_path_ref := Some conv_path


(* Generator registration *)

module GeneratorMap = Map.Make(String)

(* Map of "with"-generators *)
let generators = ref GeneratorMap.empty

(* Register a "with"-generator *)
let add_generator id e = generators := GeneratorMap.add id e !generators

(* Removes a "with"-generator *)
let rem_generator id = generators := GeneratorMap.remove id !generators


(* General purpose code generation module *)

module Gen = struct
  (* [tp_path tp] takes a type. @return a module path (list of strings)
     denoting this type. *)
  let rec tp_path = function
    | <:ctyp< $lid:id$ >> | <:ctyp< $uid:id$ >> -> [id]
    | <:ctyp< $tp1$ . $tp2$ >> ->
       (match tp_path tp2 with [n] -> n | _ -> assert false) :: tp_path tp1
    | _ -> invalid_arg "tp_path"

  (* [path_expr _loc path] takes a location [_loc] and a [path] denoting
     a type (list of strings).  @return an expression denoting the value
     given by the path. *)
  let rec path_expr _loc = function
    | [n] when n.[0] = Char.uppercase n.[0] -> <:expr< $uid:n$ >>
    | [n] -> <:expr< $lid:n$ >>
    | n :: ns -> <:expr< $path_expr _loc ns$ . $path_expr _loc [n]$ >>
    | _ -> invalid_arg "path_expr"

  (* [fold_tp _loc ~appl ~arrw ~tupl ~parm ~path ~vrow tp] takes a type and
     conversion functions.  @return the value resulting from applying the
     conversion functions to the location [_loc] and the kind of type. *)
  let rec fold_tp _loc ~appl ~arrw ~tupl ~parm ~path ~vrow tp =
    let aux = fold_tp _loc ~appl ~arrw ~tupl ~parm ~path ~vrow in
    match tp with
    | <:ctyp< $tp1$ $tp2$ >> -> appl _loc (aux tp1) (aux tp2)
    | <:ctyp< ( $list:tys$ ) >> -> tupl _loc (List.map aux tys)
    | <:ctyp< $lid:tname$ >> -> path _loc [tname]
    | <:ctyp< '$tpname$ >> -> parm _loc tpname
    | <:ctyp< $_$ . $_$ >> -> path _loc (tp_path tp)
    | <:ctyp< $tp1$ -> $tp2$ >> -> arrw _loc (aux tp1) (aux tp2)
    | <:ctyp< [= $list:row_fields$ ] >> -> vrow _loc row_fields
    | _ -> failwith "Unknown type"

  (* [abstract _loc patts body] takes a location [_loc], a pattern list
     [patts], and an expression [body].  @return a function that takes
     the patterns as arguments, and binds them in [body]. *)
  let abstract _loc patts body =
    List.fold_right (fun n e -> <:expr< fun $n$ -> $e$ >>) patts body

  let idp _loc id = <:patt< $lid:id$ >>
  let ide _loc id = <:expr< $lid:id$ >>

  (* [apply_new_args _loc id exprs] takes a location [_loc], a string
     identifier [id] and a list of expressions [expr].  @return a pair
     [(res_patts, res_exprs)], where [res_patts] is a list of patterns,
     and [res_exprs] is a list of expressions.  The expressions in [exprs]
     are applied to [id ^ ix] in [res_exprs], where [ix] is the index
     of the expression in the list, while the patterns in [res_exprs]
     correspond to the pattern [id ^ ix]. *)
  let apply_new_args _loc id =
    List.split *** list_mapi (fun ix e ->
      let n = id ^ string_of_int ix in
      idp _loc n, <:expr< $e$ $lid:n$ >>)

  (* [apply_patts _loc pat patts] takes a location [_loc], a pattern
     [pat], which stands for a constructor, and a list of patterns
     [patts], which stand for its arguments.  @return a pattern that
     matches the whole term. *)
  let apply_patts _loc = List.fold_left (fun pacc p -> <:patt< $pacc$ $p$ >>)

  (* [switch_tp_def common type_decl] takes a handler function [common]
     and a type declaration [type_decl], and apply [common] to common
     parameters of the type declaration and a function, which takes
     [~alias], [~sum], [~record] and [~variants] as handlers for the
     subcases of the type declaration. *)
  let switch_tp_def common ((_loc, name), tparams, tp_def, _cl) =
    common _loc name tparams (fun ~alias ~sum ~record ~variants ->
      match tp_def with
      | <:ctyp< [ $list:alts$ ] >> -> sum _loc alts
      | <:ctyp< [= $list:row_fields$ ] >> -> variants _loc row_fields
      | <:ctyp< $lid:_$ >>
      | <:ctyp< ( $list:_$ ) >>
      | <:ctyp< $_$ -> $_$ >>
      | <:ctyp< $_$ . $_$ >>
      | <:ctyp< '$_$ >>
      | <:ctyp< $_$ $_$ >> -> alias _loc tp_def
      | <:ctyp< { $list:flds$ } >> -> record _loc flds
      | _ -> failwith "Unknown type")

  (* [mk_expr_lst _loc expr_list] takes a list of expressions.
     @return an expression representing a list of expressions. *)
  let rec mk_expr_lst _loc = function
    | [] -> <:expr< [] >>
    | e :: es -> <:expr< [$e$ :: $mk_expr_lst _loc es$] >>

  (* [mk_patt_lst _loc patt_list] takes a list of patterns.
     @return a pattern representing a list of patterns. *)
  let rec mk_patt_lst _loc = function
    | [] -> <:patt< [] >>
    | p :: ps -> <:patt< [$p$ :: $mk_patt_lst _loc ps$] >>

  (* [mk_rev_bindings _loc fps] takes a list of values of the form
     [`Fun fun_expr] and [`Match matching].  [fun_expr] is an expression
     denoting a function, and [matching] is a list of bindings. @return
     the tuple [(bindings, patts, var_exprs)], where [bindings] is a list
     of [(pattern, expression)] tuples, [patts] is the list of those
     patterns, and [var_exprs] is the expression (variable) associated
     with each of those patterns.  The resulting lists are reversed. *)
  let mk_rev_bindings _loc fps =
    let coll (i, bindings, patts, vars) fp =
      let name = "v" ^ string_of_int i in
      let var_expr = ide _loc name in
      let expr =
        match fp with
        | `Fun fun_expr -> <:expr< $fun_expr$ $var_expr$ >>
        | `Match matchings ->
            <:expr< match $var_expr$ with [ $list:matchings$ ] >>
      in
      let patt = idp _loc name in
      i - 1, (patt, expr) :: bindings, patt :: patts, var_expr :: vars in
    let n = List.length fps in
    let _, bindings, patts, vars =
      List.fold_left coll (n, [], [], []) fps in
    bindings, patts, vars

  (* [mk_bindings _loc fps] same as [mk_rev_bindings] but the resulting
     lists are in order. *)
  let mk_bindings _loc fps = mk_rev_bindings _loc (List.rev fps)
end


(* Generators for S-expressions *)

(* Generator for converters of OCaml-values to S-expressions *)
module Generate_sexp_of = struct
  let mk_abst_call _loc tn path = Gen.path_expr _loc (("sexp_of_" ^ tn) :: path)

  (* Conversion of type paths *)
  let sexp_of_path _loc path =
    let expr =
      match path with
      | ["unit"] -> <:expr< Sexplib.Conv.sexp_of_unit >>
      | ["bool"] -> <:expr< Sexplib.Conv.sexp_of_bool >>
      | ["string"] -> <:expr< Sexplib.Conv.sexp_of_string >>
      | ["char"] -> <:expr< Sexplib.Conv.sexp_of_char >>
      | ["int"] -> <:expr< Sexplib.Conv.sexp_of_int >>
      | ["float"] -> <:expr< Sexplib.Conv.sexp_of_float >>
      | ["int32"] -> <:expr< Sexplib.Conv.sexp_of_int32 >>
      | ["int64"] -> <:expr< Sexplib.Conv.sexp_of_int64 >>
      | ["nativeint"] -> <:expr< Sexplib.Conv.sexp_of_nativeint >>
      | ["ref"] -> <:expr< Sexplib.Conv.sexp_of_ref >>
      | ["t"; "Lazy"] | ["lazy_t"] -> <:expr< Sexplib.Conv.sexp_of_lazy >>
      | ["option"] -> <:expr< Sexplib.Conv.sexp_of_option >>
      | ["list"] -> <:expr< Sexplib.Conv.sexp_of_list >>
      | ["array"] -> <:expr< Sexplib.Conv.sexp_of_array >>
      | ["t"; "Hashtbl"] -> <:expr< Sexplib.Conv.sexp_of_hashtbl >>
      | tn :: path -> mk_abst_call _loc tn path
      | [] -> assert false  (* impossible *)
    in
    `Fun expr

  (* Conversion of tuples *)
  let sexp_of_tuple _loc fps =
    let bindings, patts, vars = Gen.mk_bindings _loc fps in
    let expr =
      <:expr<
        let $list:bindings$ in
        Sexplib.Sexp.List $Gen.mk_expr_lst _loc vars$
      >>
    in
    let matching =
      (
        <:patt< ( $list:patts$ ) >>,
        None,
        expr
      )
    in
    `Match [matching]

  (* Conversion of polymorphic types *)

  let sexp_of_appl_fun _loc fp1 fp2 =
    match fp1, fp2 with
    | `Fun fun_expr1, `Fun fun_expr2 -> <:expr< $fun_expr1$ $fun_expr2$ >>
    | `Fun fun_expr, `Match matchings ->
        <:expr< $fun_expr$ (fun [ $list:matchings$ ]) >>
    | _ -> assert false  (* impossible *)

  let sexp_of_appl _loc fp1 fp2 = `Fun (sexp_of_appl_fun _loc fp1 fp2)

  (* Conversion of type parameters *)
  let sexp_of_parm _loc parm = `Fun (Gen.ide _loc ("of__" ^ parm))

  (* Conversion of function types *)
  let sexp_of_arrow _loc _ _ =
    failwith "sexp_of_arrow: cannot convert functions to S-expressions"

  (* Get the path associated with a polymorphic type *)
  let rec get_appl_path _loc = function
    | <:ctyp< $lid:_$ >> as tp -> tp
    | <:ctyp< $_$ . $_$ >> as tp -> tp
    | <:ctyp< $tp$ $_$ >> -> get_appl_path _loc tp
    | _ -> failwith "get_appl_path: unknown type"

  (* Conversion of types *)
  let rec sexp_of_type _loc tp =
    Gen.fold_tp
      _loc
      ~appl:sexp_of_appl ~arrw:sexp_of_arrow
      ~tupl:sexp_of_tuple ~parm:sexp_of_parm
      ~path:sexp_of_path ~vrow:sexp_of_variant
      tp

  (* Conversion of variant types *)

  and sexp_of_variants_loop _loc acc = function
    | RfTag (cnstr, _, []) ->
        (
          <:patt< `$cnstr$ >>,
          None,
          <:expr< Sexplib.Sexp.Atom $str:cnstr$ >>
        ) :: acc
    | RfTag (cnstr, _, tps) ->
        let fps = List.map (sexp_of_type _loc) tps in
        let bindings, patts, vars = Gen.mk_bindings _loc fps in
        let arg_patt =
          match patts with
          | [patt] -> patt
          | _ -> <:patt< ( $list:patts$ ) >>
        in
        let cnstr_expr = <:expr< Sexplib.Sexp.Atom $str:cnstr$ >> in
        (
          <:patt< `$cnstr$ $arg_patt$ >>,
          None,
          <:expr<
            let $list:bindings$ in
            Sexplib.Sexp.List $Gen.mk_expr_lst _loc (cnstr_expr :: vars)$
          >>
        ) :: acc
    | RfInh (TyVrn (_loc, row_fields, _)) ->
        List.fold_left (sexp_of_variants_loop _loc) acc row_fields
    | RfInh tp ->
        let tp_path, call =
          match tp with
          | <:ctyp< $tp1$ $tp2$ >> ->
              let fp1 = sexp_of_type _loc tp1 in
              let fp2 = sexp_of_type _loc tp2 in
              let tp_path = List.rev (Gen.tp_path (get_appl_path _loc tp1)) in
              let expr = sexp_of_appl_fun _loc fp1 fp2 in
              tp_path, expr
          | _ ->
              let tp_path = Gen.tp_path tp in
              match tp_path with
              | tn :: path -> tp_path, mk_abst_call _loc tn path
              | [] -> assert false  (* impossible *)
        in
        (
          <:patt< (#$tp_path$ as v) >>,
          None,
          <:expr< $call$ v >>
        ) :: acc

  and sexp_of_variant _loc row_fields =
    `Match (
      List.rev (List.fold_left (sexp_of_variants_loop _loc) [] row_fields))

  (* Conversion of sum types *)

  let branch_sum (_loc, cnstr, tps) =
    let fps = List.map (sexp_of_type _loc) tps in
    let cnstr_expr = <:expr< Sexplib.Sexp.Atom $str:cnstr$ >> in
    let cnstr_patt = <:patt< $uid:cnstr$ >> in
    if fps = [] then
      (
        cnstr_patt,
        None,
        cnstr_expr
      )
    else
      let bindings, patts, vars = Gen.mk_bindings _loc fps in
      let expr =
        <:expr<
          let $list:bindings$ in
          Sexplib.Sexp.List $Gen.mk_expr_lst _loc (cnstr_expr :: vars)$
        >>
      in
      let arg_patt =
        match patts with
        | [patt] -> patt
        | _ -> <:patt< ( $list:patts$ ) >>
      in
      (
        <:patt< $cnstr_patt$ $arg_patt$ >>,
        None,
        expr
      )

  let sexp_of_sum _ alts = `Match (List.map branch_sum alts)

  (* Conversion of record types *)
  let sexp_of_record _loc flds =
    let names = List.map snd4 flds in
    let name_patts = List.map (Gen.idp _loc) names in
    let cnv_fld (_loc, _, _, tp) = sexp_of_type _loc tp in
    let fps = List.map cnv_fld flds in
    let bindings, patts, vars = Gen.mk_bindings _loc fps in
    let cnv_fld_lst name var =
      let name_str = <:expr< Sexplib.Sexp.Atom $str:name$ >> in
      <:expr< Sexplib.Sexp.List $Gen.mk_expr_lst _loc [name_str; var]$ >>
    in
    let fld_lst = List.map2 cnv_fld_lst names vars in
    let binding =
      (
        <:patt< { $list:List.combine name_patts patts$ } >>,
        None,
        <:expr<
          let $list:bindings$ in
          Sexplib.Sexp.List $Gen.mk_expr_lst _loc fld_lst$
        >>
      )
    in
    `Match [binding]

  (* Generate code from type definitions *)
  let generate =
    Gen.switch_tp_def (fun _loc type_name tps case ->
      let is_alias_ref = ref false in
      let handle_alias _loc tp =
        is_alias_ref := true;
        sexp_of_type _loc tp in
      let body =
        match
          case
            ~alias:handle_alias
            ~sum:sexp_of_sum
            ~variants:sexp_of_variant
            ~record:sexp_of_record
        with
        | `Fun fun_expr ->
            (* Prevent violation of value restriction *)
            if !is_alias_ref && tps = [] then
              let matching =
                (
                  <:patt< v >>,
                  None,
                  <:expr< $fun_expr$ v >>
                )
              in
              <:expr< fun [ $list:[matching]$ ] >>
            else <:expr< $fun_expr$ >>
        | `Match matchings -> <:expr< fun [ $list:matchings$ ] >>
      in
      let patts = List.map (Gen.idp _loc *** (^) "of__" *** fst) tps in
      Gen.idp _loc ("sexp_of_" ^ type_name), (Gen.abstract _loc patts body)
    )

  (* Add code generator to the set of known generators *)
  let () = add_generator "sexp_of" (fun td -> ([generate td], []))
end


(* Generator for converters of S-expressions to OCaml-values *)
module Generate_of_sexp = struct
  let mk_abst_call _loc tn ?(internal = false) path =
    let suff =
      if internal then "_of_sexp__"
      else "_of_sexp"
    in
    Gen.path_expr _loc ((tn ^ suff) :: path)

  (* Conversion of type paths *)

  let path_of_sexp_fun _loc path =
    match path with
    | ["unit"] -> <:expr< Sexplib.Conv.unit_of_sexp >>
    | ["string"] -> <:expr< Sexplib.Conv.string_of_sexp >>
    | ["int"] -> <:expr< Sexplib.Conv.int_of_sexp >>
    | ["float"] -> <:expr< Sexplib.Conv.float_of_sexp >>
    | ["bool"] -> <:expr< Sexplib.Conv.bool_of_sexp >>
    | ["int32"] -> <:expr< Sexplib.Conv.int32_of_sexp >>
    | ["int64"] -> <:expr< Sexplib.Conv.int64_of_sexp >>
    | ["nativeint"] -> <:expr< Sexplib.Conv.nativeint_of_sexp >>
    | ["list"] -> <:expr< Sexplib.Conv.list_of_sexp >>
    | ["array"] -> <:expr< Sexplib.Conv.array_of_sexp >>
    | ["option"] -> <:expr< Sexplib.Conv.option_of_sexp >>
    | ["char"] -> <:expr< Sexplib.Conv.char_of_sexp >>
    | ["t"; "Lazy"] | ["lazy_t"] -> <:expr< Sexplib.Conv.lazy_of_sexp >>
    | ["t"; "Hashtbl"] -> <:expr< Sexplib.Conv.hashtbl_of_sexp >>
    | ["ref"] -> <:expr< Sexplib.Conv.ref_of_sexp >>
    | tn :: path -> mk_abst_call _loc tn path
    | [] -> assert false (* no empty paths *)

  let path_of_sexp _loc path = `Fun (path_of_sexp_fun _loc path)

  (* Conversion of tuples *)
  let tuple_of_sexp _loc fps =
    let bindings, patts, vars = Gen.mk_bindings _loc fps in
    let matching =
      [
        (
          <:patt< Sexplib.Sexp.List $Gen.mk_patt_lst _loc patts$ >>,
          None,
          <:expr<
            let $list:bindings$ in
            ( $list:vars$ )
          >>
        );
        (
          <:patt< sexp >>,
          None,
          let n = string_of_int (List.length fps) in
          <:expr<
            Sexplib.Conv_error.tuple_of_size_n_expected _loc $int:n$ sexp
          >>
        )
      ]
    in
    `Match matching

  (* Conversion of applications *)

  let appl_of_sexp_fun _loc fp1 fp2 =
    match fp1, fp2 with
    | `Fun fun_expr1, `Fun fun_expr2 -> <:expr< $fun_expr1$ $fun_expr2$ >>
    | `Fun fun_expr, `Match matchings ->
        <:expr< $fun_expr$ (fun [ $list:matchings$ ]) >>
    | _ -> assert false  (* impossible *)

  let appl_of_sexp _loc fp1 fp2 = `Fun (appl_of_sexp_fun _loc fp1 fp2)

  (* Conversion of type parameters *)
  let parm_of_sexp _loc parm = `Fun (Gen.ide _loc ("of__" ^ parm))

  (* Conversion of function types *)
  let arrow_of_sexp _loc _ _ =
    failwith "arrow_of_sexp: cannot convert S-expressions to functions"

  (* Handle backtracking when variants do not match *)
  let handle_no_variant_match _loc expr =
    (
      <:patt< Sexplib.Conv_error.No_variant_match _ >>,
      None,
      expr
    )

  let is_wildcard = function [_] -> true | _ -> false

  (* Generate code depending on whether to generate a match for the last
     case of matching a variant *)
  let handle_variant_match_last _loc match_last matches =
    if match_last || is_wildcard matches then
      let _, _, expr = List.hd matches in
      expr
    else <:expr< match atom with [ $list:matches$ ] >>

  (* Generate code for matching malformed S-expressions *)
  let mk_variant_other_matches _loc rev_els call =
    let coll_structs acc (_loc, cnstr) =
      let el =
        (
          <:patt< $str:cnstr$ >>,
          None,
          <:expr< Sexplib.Conv_error.$lid:call$ _loc sexp >>
        )
      in
      el :: acc in
    let exc_no_variant_match =
      (
        <:patt< _ >>,
        None,
        <:expr< Sexplib.Conv_error.no_variant_match _loc sexp >>
      )
    in
    List.fold_left coll_structs [exc_no_variant_match] rev_els

  (* Split the row fields of a variant type into lists of atomic variants,
     structured variants, atomic variants + included variant types,
     and structured variants + included variant types. *)
  let rec split_row_fields _loc (atoms, structs, ainhs, sinhs as acc) = function
    | RfTag (cnstr, _, []) ->
        let tpl = _loc, cnstr in
        (
          tpl :: atoms,
          structs,
          `A tpl :: ainhs,
          sinhs
        )
    | RfTag (cnstr, _, tps) ->
        (
          atoms,
          (_loc, cnstr) :: structs,
          ainhs,
          `S (_loc, cnstr, tps) :: sinhs
        )
    | RfInh (TyVrn (_loc, row_fields, _)) ->
        List.fold_left (split_row_fields _loc) acc row_fields
    | RfInh inh ->
        let iinh = `I (_loc, inh) in
        (
          atoms,
          structs,
          iinh :: ainhs,
          iinh :: sinhs
        )

  (* Generate internal call *)
  let rec mk_internal_call type_name _loc = function
    | <:ctyp< $lid:tname$ >> -> mk_abst_call _loc tname ~internal:true []
    | <:ctyp< $path$ . $lid:tname$ >> ->
        mk_abst_call _loc tname ~internal:true (Gen.tp_path path)
    | <:ctyp< $tp1$ $tp2$ >> ->
        let fp1 = `Fun (mk_internal_call type_name _loc tp1) in
        let fp2 = type_of_sexp type_name _loc tp2 in
        appl_of_sexp_fun _loc fp1 fp2
    | _ -> assert false  (* impossible *)

  (* Conversion of types *)
  and type_of_sexp type_name _loc tp =
    Gen.fold_tp
      _loc
      ~appl:appl_of_sexp ~arrw:arrow_of_sexp
      ~tupl:tuple_of_sexp ~parm:parm_of_sexp
      ~path:path_of_sexp ~vrow:(variant_of_sexp type_name ?full_type:None)
      tp

  (* Generate code for matching included variant types *)
  and handle_variant_inh
            type_name _loc full_type match_last other_matches inh =
    let match_exc =
      handle_no_variant_match _loc (
        handle_variant_match_last _loc match_last other_matches) in
    let fun_expr = mk_internal_call type_name _loc inh in
    let new_other_matches =
      [
        (
          <:patt< _ >>,
          None,
          <:expr<
            try ($fun_expr$ sexp :> $full_type$)
            with [ $list:[match_exc]$ ]
          >>
        )
      ] in
    new_other_matches, true

  (* Generate code for matching atomic variants *)
  and mk_variant_match_atom
        type_name _loc full_type rev_atoms_inhs rev_structs =
    let coll (other_matches, match_last) = function
      | `A (_loc, cnstr) ->
          let new_match =
            (
              <:patt< $str:cnstr$ >>,
              None,
              <:expr< `$cnstr$ >>
            )
          in
          new_match :: other_matches, false
      | `I (_loc, inh) ->
          handle_variant_inh
            type_name _loc full_type match_last other_matches inh
    in
    let other_matches =
      mk_variant_other_matches _loc rev_structs "variant_needs_arguments" in
    let match_atoms_inhs, match_last =
      List.fold_left coll (other_matches, false) rev_atoms_inhs in
    handle_variant_match_last _loc match_last match_atoms_inhs

  (* Match arguments of constructors (variants or sum types) *)
  and mk_cnstr_args_match type_name _loc ~is_variant cnstr tps =
    let arg_fps = List.map (type_of_sexp type_name _loc) tps in
    let bindings, patts, vars = Gen.mk_bindings _loc arg_fps in
    let good_arg_match_expr =
        match vars with
        | [var] when is_variant -> <:expr< `$cnstr$ $var$ >>
        | [var] -> <:expr< $uid:cnstr$ $var$ >>
        | _ when is_variant -> <:expr< `$cnstr$ ( $list:vars$ ) >>
        | _ -> <:expr< $uid:cnstr$ ( $list:vars$ ) >> in
    let good_arg_match =
      (
        <:patt< $Gen.mk_patt_lst _loc patts$ >>,
        None,
        <:expr<
          let $list:bindings$ in
          $good_arg_match_expr$
        >>
      )
    in
    let bad_arg_match =
      let handle_exc =
        if is_variant then "variant_incorrect_number_arguments"
        else "sum_incorrect_number_arguments" in
      (
        <:patt< _ >>,
        None,
        <:expr< Sexplib.Conv_error.$lid:handle_exc$ _loc tag sexp >>
      )
    in
    <:expr< match sexp_args with [ $list:[good_arg_match; bad_arg_match]$ ] >>


  (* Variant conversions *)

  (* Generate code for matching structured variants *)
  and mk_variant_match_struct
        type_name _loc full_type rev_structs_inhs rev_atoms =
    let has_structs_ref = ref false in
    let coll (other_matches, match_last) = function
      | `S (_loc, cnstr, tps) ->
          has_structs_ref := true;
          let expr =
            mk_cnstr_args_match type_name _loc ~is_variant:true cnstr tps in
          let new_match =
            (
              <:patt< ($str:cnstr$ as tag) >>,
              None,
              expr
            )
          in
          new_match :: other_matches, false
      | `I (_loc, inh) ->
          handle_variant_inh
            type_name _loc full_type match_last other_matches inh
    in
    let other_matches =
      let call = "variant_does_not_take_arguments" in
      mk_variant_other_matches _loc rev_atoms call in
    let match_structs_inhs, match_last =
      List.fold_left coll (other_matches, false) rev_structs_inhs in
    (
      handle_variant_match_last _loc match_last match_structs_inhs,
      !has_structs_ref
    )

  (* Generate code for handling atomic and structured variants (i.e. not
     included variant types) *)
  and handle_variant_rftag type_name _loc full_type row_fields =
    let rev_atoms, rev_structs, rev_atoms_inhs, rev_structs_inhs =
      List.fold_left (split_row_fields _loc) ([], [], [], []) row_fields in
    let match_atom =
      mk_variant_match_atom
        type_name _loc full_type rev_atoms_inhs rev_structs in
    let match_struct, has_structs =
      mk_variant_match_struct
        type_name _loc full_type rev_structs_inhs rev_atoms in
    let maybe_sexp_args_patt =
      if has_structs then <:patt< sexp_args >>
      else <:patt< _ >>
    in
    [
      (
        <:patt< (Sexplib.Sexp.Atom atom as sexp) >>,
        None,
        <:expr< $match_atom$ >>
      );
      (
        <:patt<
          (Sexplib.Sexp.List
            [Sexplib.Sexp.Atom atom :: $maybe_sexp_args_patt$] as sexp)
        >>,
        None,
        <:expr< $match_struct$ >>;
      );
      (
        <:patt<
          (Sexplib.Sexp.List [Sexplib.Sexp.List _ :: _] as sexp)
        >>,
        None,
        <:expr< Sexplib.Conv_error.list_in_list_invalid_variant _loc sexp >>
      );
      (
        <:patt< (Sexplib.Sexp.List [] as sexp) >>,
        None,
        <:expr< Sexplib.Conv_error.empty_list_invalid_variant _loc sexp >>
      )
    ]

  (* Generate matching code for variants *)
  and variant_of_sexp type_name _loc ?full_type row_fields =
    let full_type =
      match full_type with
      | None -> <:ctyp< [= $list:row_fields$ ] >>
      | Some full_type -> full_type in
    let top_match =
      match row_fields with
      | RfTag _ :: _ -> handle_variant_rftag type_name _loc full_type row_fields
      | RfInh inh :: rest ->
          let rec loop inh row_fields =
            let call =
              <:expr<
                ( $mk_internal_call type_name _loc inh$ sexp :> $full_type$ )
              >> in
            match row_fields with
            | [] -> call
            | h :: t ->
                let expr =
                  match h with
                  | RfTag _ ->
                     let rftag_matches =
                       handle_variant_rftag
                         type_name _loc full_type row_fields in
                     <:expr< match sexp with [ $list:rftag_matches$ ] >>
                  | RfInh inh2 -> loop inh2 t
                in
                <:expr<
                  try $call$ with
                  [ $list:[handle_no_variant_match _loc expr]$ ]
                >>
          in
          [ (<:patt< sexp >>, None, loop inh rest) ]
      | [] -> assert false  (* impossible *)
    in
    `Match top_match


  (* Sum type conversions *)

  (* Generate matching code for well-formed S-expressions wrt. sum types *)
  let mk_good_sum_match type_name (_loc, cnstr, tps) =
    let lccnstr = String.copy cnstr in
    lccnstr.[0] <- Char.lowercase lccnstr.[0];
    if tps = [] then  (* Sum type takes no argument *)
      (
        <:patt< Sexplib.Sexp.Atom ($str:lccnstr$ | $str:cnstr$) >>,
        None,
        <:expr< $uid:cnstr$ >>
      )
    else  (* Sum type takes arguments *)
      let match_patt =
        <:patt<
          (Sexplib.Sexp.List
            [Sexplib.Sexp.Atom ($str:lccnstr$ | $str:cnstr$ as tag) ::
              sexp_args] as sexp)
        >> in
      let match_body =
        mk_cnstr_args_match type_name _loc ~is_variant:false cnstr tps in
      (
        match_patt,
        None,
        match_body
      )

  (* Generate matching code for malformed S-expressions with good tags
     wrt. sum types *)
  let mk_bad_sum_match (_loc, cnstr, tps) =
    let lccnstr = String.copy cnstr in
    lccnstr.[0] <- Char.lowercase lccnstr.[0];
    if tps = [] then  (* Sum type takes no argument *)
      (
        <:patt<
          (Sexplib.Sexp.List
            [Sexplib.Sexp.Atom ($str:lccnstr$ | $str:cnstr$) ::
              _] as sexp)
        >>,
        None,
        <:expr< Sexplib.Conv_error.sum_tag_must_be_atomic _loc sexp >>
      )
    else (* Sum type takes arguments *)
      (
        <:patt<
          (Sexplib.Sexp.Atom ($str:lccnstr$ | $str:cnstr$) as sexp)
        >>,
        None,
        <:expr< Sexplib.Conv_error.sum_tag_must_be_struct _loc sexp >>
      )

  (* Generate matching code for sum types *)
  let sum_of_sexp type_name _loc alts =
    let rev_good_sum_matches =
      List.rev_map (mk_good_sum_match type_name) alts in
    let rev_bad_sum_matches = List.rev_map mk_bad_sum_match alts in
    let bad_lst_match =
      (
        <:patt<
          (Sexplib.Sexp.List [Sexplib.Sexp.List _ :: _] as sexp)
        >>,
        None,
        <:expr< Sexplib.Conv_error.list_in_list_invalid_sum _loc sexp >>
      ) in
    let bad_tag_matches =
      [
        (
          <:patt< (Sexplib.Sexp.List [] as sexp) >>,
          None,
          <:expr< Sexplib.Conv_error.empty_list_invalid_sum _loc sexp >>
        );
        (
          <:patt< sexp >>,
          None,
          <:expr< Sexplib.Conv_error.sum_unknown_tag _loc sexp >>
        )
      ]
    in
    let bad_matches =
      List.rev_append rev_bad_sum_matches (bad_lst_match :: bad_tag_matches) in
    `Match (List.rev_append rev_good_sum_matches bad_matches)


  (* Record conversions *)

  (* Unpack the value bound to [name] of type [tp]. *)
  let unpack_fun_match type_name _loc name tp =
    match type_of_sexp type_name _loc tp with
    | `Fun fun_expr -> <:expr< $fun_expr$ $lid:name$ >>
    | `Match bindings -> <:expr< match $lid:name$ with [ $list:bindings$ ] >>

  (* Generate code for extracting record fields *)
  let mk_extract_field type_name (_loc, nm, _mut, tp) =
    let lcnm = String.copy nm in
    let ucnm = String.copy nm in
    lcnm.[0] <- Char.lowercase nm.[0];
    ucnm.[0] <- Char.uppercase nm.[0];
    let expr = unpack_fun_match type_name _loc "field_sexp" tp in
    let fvalue = <:patt< fvalue >>, expr in
    (
      <:patt< $str:lcnm$ | $str:ucnm$ >>,
      None,
      <:expr<
        let $list:[fvalue]$ in
        $lid:nm ^ "_field"$.contents := Some fvalue
      >>
    )

  (* Generate code for matching record fields *)
  let mk_match_fields type_name _loc flds =
    let rev_field_bindings = List.rev_map (mk_extract_field type_name) flds in
    let field_matches =
      let match_unknown = <:patt< _ >>, None, <:expr< () >> in
      List.rev_append rev_field_bindings [match_unknown] in
    let match_field_name =
      <:expr< match field_name with [ $list:field_matches$ ] >> in
    let match_tpl =
      [
        (
          <:patt<
            [
              Sexplib.Sexp.List
                [(Sexplib.Sexp.Atom field_name); field_sexp] ::
              tail
            ]
          >>,
          None,
          <:expr< do { $match_field_name$; iter tail } >>
        );
        (
          <:patt< [sexp :: _] >>,
          None,
          <:expr< Sexplib.Conv_error.record_only_pairs_expected _loc sexp >>
        );
        (
          <:patt< [] >>,
          None,
          <:expr< () >>
        );
      ] in
    let match_expr = <:expr< fun [ $list:match_tpl$ ] >> in
    <:patt< iter >>, match_expr

  (* Generate code for handling the result of matching record fields *)
  let mk_handle_record_match_result _loc flds =
    let res_tuples, bindings_lst =
      List.split (
        List.map (fun (_loc, nm, _mut, _tp) ->
          let fld = <:expr< $lid:nm ^ "_field"$.contents >> in
          (
            <:expr< $fld$ >>,
            <:expr< ($fld$ = None, $str:nm$) >>
          )
        ) flds) in
    let good_pats =
      List.map (fun (_loc, nm, _mut, _tp) ->
        <:patt< Some $lid:nm ^ "_value"$ >>) flds in
    let return_good =
      List.map (fun (_loc, nm, _mut, _tp) ->
        <:patt< $lid:nm$ >>, <:expr< $lid:nm ^ "_value"$ >>) flds in
    let match_good_res_patt =
      match good_pats with
      | [good_pat] -> <:patt< $good_pat$ >>
      | _ -> <:patt< ( $list:good_pats$ ) >>
    in
    let match_good_res =
      (
        match_good_res_patt,
        None,
        <:expr< { $list:return_good$ } >>
      )
    in
    let match_bad_res =
      <:patt< _ >>,
      None,
      let bindings = Gen.mk_expr_lst _loc bindings_lst in
      <:expr<
        Sexplib.Conv_error.record_undefined_elements _loc sexp $bindings$
      >>
    in
    let match_expr =
      match res_tuples with
      | [res_tuple] -> <:expr< $res_tuple$ >>
      | _ -> <:expr< ( $list:res_tuples$ ) >>
    in
    <:expr<
      match $match_expr$ with
      [ $list:[match_good_res; match_bad_res]$ ]
    >>

  (* Generate code for converting record fields *)
  let mk_cnv_fields type_name _loc flds =
    let field_refs =
      let mk_field (_loc, nm, _mut, _tp) =
        (<:patt< $lid:nm ^ "_field"$ >>, <:expr< ref None >>) in
      List.map mk_field flds in
    let match_patt_expr = mk_match_fields type_name _loc flds in
    let match_result = mk_handle_record_match_result _loc flds in
    <:expr<
      let $list:field_refs$ in
      let $opt:true$ $list:[match_patt_expr]$ in
      do { iter field_sexps; $match_result$ }
    >>

  (* Generate matching code for records *)
  let record_of_sexp type_name _loc flds =
    let cnv_fields = mk_cnv_fields type_name _loc flds in
    let handle_bad_atom =
      <:expr< Sexplib.Conv_error.record_list_instead_atom _loc sexp >> in
    let top_match =
      [
        (
          <:patt< (Sexplib.Sexp.List field_sexps as sexp) >>,
          None,
          cnv_fields
        );
        (
          <:patt< (Sexplib.Sexp.Atom _ as sexp) >>,
          None,
          handle_bad_atom
        );
      ] in
    `Match top_match

  (* Generate code from type definitions *)
  let generate =
    Gen.switch_tp_def (fun _loc type_name tps case ->
      let rec coll_args tp (arg, _) = <:ctyp< $tp$ '$lid:arg$ >> in
      let full_type = List.fold_left coll_args <:ctyp< $lid:type_name$ >> tps in
      let is_variant_ref = ref false in
      let handle_variant row_fields =
        is_variant_ref := true;
        variant_of_sexp type_name ~full_type row_fields in
      let is_alias_ref = ref false in
      let handle_alias _loc tp =
        is_alias_ref := true;
        type_of_sexp type_name _loc tp in
      let body_func =
        match
          case
            ~alias:handle_alias
            ~sum:(sum_of_sexp type_name)
            ~variants:handle_variant
            ~record:(record_of_sexp type_name)
        with
        | `Fun fun_expr ->
            (* Prevent violation of value restriction *)
            if !is_alias_ref && tps = [] then
              let matching =
                (
                  <:patt< sexp >>,
                  None,
                  <:expr< $fun_expr$ sexp >>
                )
              in
              <:expr< fun [ $list:[matching]$ ] >>
            else <:expr< $fun_expr$ >>
        | `Match matchings -> <:expr< fun [ $list:matchings$ ] >>
      in
      let func' =
        let patts = List.map (Gen.idp _loc *** (^) "of__" *** fst) tps in
        let body_func = Gen.abstract _loc patts body_func in
        let path = get_conv_path () in
        let bindings =
          (
            <:patt< _loc >>,
            <:expr< $str:sprintf "%s.%s" path type_name$ >>
          )
        in
        <:expr< let $list:[bindings]$ in $body_func$ >>
      in
      if !is_variant_ref then
        let name__ = type_name ^ "_of_sexp__" in
        let fun_name__ = Gen.ide _loc name__ in
        let f' = Gen.idp _loc name__, func' in
        let func =
          let handle_exc =
            (
              <:patt< Sexplib.Conv_error.No_variant_match (msg, sexp) >>,
              None,
              <:expr< Sexplib.Conv.of_sexp_error msg sexp >>
            )
          in
          let matching =
            (
              <:patt< sexp >>,
              None,
              <:expr< try $fun_name__$ sexp with [ $list:[handle_exc]$ ] >>
            )
          in
          <:expr< fun [ $list:[matching]$ ] >>
        in
        let f = Gen.idp _loc (type_name ^ "_of_sexp"), func in
        [f'; f]
      else [(Gen.idp _loc (type_name ^ "_of_sexp"), func')]
    )

  (* Add code generator to the set of known generators *)
  let () = add_generator "of_sexp" (fun td -> (generate td, []))
end

(* Add "of_sexp" and "sexp_of" as "sexp" to the set of generators *)
let () =
  add_generator
    "sexp"
    (fun td -> Generate_sexp_of.generate td :: Generate_of_sexp.generate td, [])


(* Functions for interpreting derivation types *)

(* Generates a tuple of lists of functions and types. *)
let generate _loc tp id =
  try GeneratorMap.find id !generators tp
  with Not_found ->
    failwith ("Pa_sexp_conv: '"^ id ^"' is not a supported generator.")

let gen_derived_defs _loc (ty, drvs) =
  let derived = List.map (generate _loc ty) drvs in
  both List.concat List.concat (List.split derived)


(* Syntax extension *)

DELETE_RULE
  Pcaml.str_item: "type"; LIST1 Pcaml.type_declaration SEP "and"
END

EXTEND
  GLOBAL: Pcaml.str_item;
  with_decl:
  [[ td = Pcaml.type_declaration; "with"; idl = LIST1 LIDENT SEP "," -> td, idl
  |  td = Pcaml.type_declaration -> td, [] ]];

  Pcaml.str_item:
  [[ "type"; twdl = LIST1 with_decl SEP "and" ->
    let tys = List.map fst twdl in
    let (der_vals, der_tys) =
      both List.concat List.concat
        (List.split (List.map (gen_derived_defs _loc) twdl)) in
    let tydef = <:str_item< type $list:tys@der_tys$ >> in
    let valdef = <:str_item< value rec $list:der_vals$ >> in
    match der_vals with
    | [] -> tydef
    | _ -> <:str_item< declare $tydef$; $valdef$; end >> ]];

  Pcaml.str_item:
  [[ "SEXP_CONV_PATH"; conv_path = STRING ->
    set_conv_path conv_path;
    StDcl (_loc, [])
  ]];
END