Extraction of polymorphic recursion

src/mlw/compile.ml

@@ -29,6 +29,7 @@ open Ident
 open Ity
 open Ty
 open Term
+open Stdlib
 
 let clean_name fname =
   (* TODO: replace with Filename.remove_extension
@@ -53,6 +54,13 @@ module ML = struct
     | Tapp    of ident * ty list
     | Ttuple  of ty list
 
+  module Vvar = MakeMSHW (struct
+    type t = tvsymbol
+    let tag tvs = tvs.tv_name.id_tag
+  end)
+
+  module Svar = Vvar.S
+
   type is_ghost = bool
 
   type var = ident * ty * is_ghost
@@ -109,6 +117,7 @@ module ML = struct
     rec_args : var list;
     rec_exp  : expr;
     rec_res  : ty;
+    rec_svar : Svar.t;
   }
 
   type is_mutable = bool
@@ -514,11 +523,11 @@ module Translate = struct
       let for_call      = ML.Eapp (for_rs, [from_expr]) in
       ML.mk_expr for_call ML.ity_unit eff in
     let pv_name pv = pv.pv_vs.vs_name in
-    let args = [ pv_name  i_pv, ty_int, false ] in
+    let args = [ pv_name i_pv, ty_int, false ] in
     let for_rec_def = {
-      ML.rec_sym  = for_rs; ML.rec_args = args;
-      ML.rec_rsym = for_rs; ML.rec_exp  = for_expr;
-      ML.rec_res  = ML.tunit;
+      ML.rec_sym  = for_rs;   ML.rec_args = args;
+      ML.rec_rsym = for_rs;   ML.rec_exp  = for_expr;
+      ML.rec_res  = ML.tunit; ML.rec_svar = ML.Svar.empty;
     } in
     let for_let = ML.Elet (ML.Lrec [for_rec_def], for_call_expr) in
     ML.mk_expr for_let ML.ity_unit eff
@@ -537,6 +546,12 @@ module Translate = struct
 
   exception ExtractionAny
 
+  (* build the set of type variables from functions arguments *)
+  let rec add_tvar acc = function
+    | ML.Tvar tv -> ML.Svar.add tv acc
+    | ML.Tapp (_, tyl) | ML.Ttuple tyl ->
+      List.fold_left add_tvar acc tyl
+
   (* expressions *)
   let rec expr info ({e_effect = eff} as e) =
     assert (not eff.eff_ghost);
@@ -592,10 +607,16 @@ module Translate = struct
       let def = function
         | {rec_sym = rs1; rec_rsym = rs2;
            rec_fun = {c_node = Cfun ef; c_cty = cty}} ->
-            let res  = ity rs1.rs_cty.cty_result in
-            let args = params cty.cty_args in let ef = expr info ef in
-            { ML.rec_sym  = rs1;  ML.rec_rsym = rs2;
-              ML.rec_args = args; ML.rec_exp  = ef ; ML.rec_res  = res }
+          let res  = ity rs1.rs_cty.cty_result in
+          let args = params cty.cty_args in
+          let svar =
+            let args' = List.map (fun (_, ty, _) -> ty) args in
+            let svar  = List.fold_left add_tvar ML.Svar.empty args' in
+            add_tvar svar res in
+          let ef = expr info ef in
+          { ML.rec_sym  = rs1;  ML.rec_rsym = rs2;
+            ML.rec_args = args; ML.rec_exp  = ef ;
+            ML.rec_res  = res;  ML.rec_svar = svar; }
         | _ -> assert false in
       let rdefl = List.map def rdefl in
       if rdefl <> [] then
@@ -749,9 +770,13 @@ module Translate = struct
         let e = match e.c_node with Cfun e -> e | _ -> assert false in
         let args = params rs1.rs_cty.cty_args in
         let res  = ity rs1.rs_cty.cty_result in
+        let svar =
+          let args' = List.map (fun (_, ty, _) -> ty) args in
+          let svar  = List.fold_left add_tvar ML.Svar.empty args' in
+          add_tvar svar res in
         { ML.rec_sym  = rs1;  ML.rec_rsym = rs2;
           ML.rec_args = args; ML.rec_exp  = expr info e;
-          ML.rec_res  = res } in
+          ML.rec_res  = res;  ML.rec_svar = svar; } in
       if rl = [] then [] else [ML.Dlet (ML.Lrec (List.map def rl))]
     | PDlet (LDsym _) | PDpure | PDlet (LDvar _) ->
       []

src/mlw/ocaml_printer.ml

@@ -171,7 +171,7 @@ module Print = struct
     if is_optional ~labels then print_vsty_opt info fmt id ty
     else if is_named ~labels then print_vsty_named info fmt id ty
     else fprintf fmt "(%a:@ %a)" (print_lident info) id
-        (print_ty ~paren:false info) ty
+      (print_ty ~paren:false info) ty
 
   let print_tv_arg = print_tv
   let print_tv_args fmt = function
@@ -198,8 +198,8 @@ module Print = struct
   let get_record info rs =
     match Mid.find_opt rs.rs_name info.info_mo_known_map with
     | Some {pd_node = PDtype itdl} ->
-      let itd = List.find (fun {itd_constructors=constr} ->
-          List.exists (rs_equal rs) constr) itdl in
+      let eq_rs {itd_constructors=constr} = List.exists (rs_equal rs) constr in
+      let itd = List.find eq_rs itdl in
       List.filter (fun e -> not (rs_ghost e)) itd.itd_fields
     | _ -> []
 
@@ -314,6 +314,23 @@ module Print = struct
         (print_apply_args info) (tl, rs.rs_cty.cty_args)
         (* (print_list space (print_expr ~paren:true info)) tl *)
 
+  and print_svar fmt s =
+    Svar.iter (fun tv -> fprintf fmt "%a" print_tv tv) s
+
+  and print_fun_type_args info fmt (args, s, res) =
+    if Svar.is_empty s then
+      fprintf fmt "@[%a@] :@ %a@ =@"
+        (print_list space (print_vs_arg info)) args
+        (print_ty info) res
+    else
+      let ty_args = List.map (fun (_, ty, _) -> ty) args in
+      let id_args = List.map (fun (id, _, _) -> id) args in
+      fprintf fmt ": @[%a@]. @[%a@] ->@ %a@ =@ fun @[%a@]@ ->@\n"
+        print_svar s
+        (print_list arrow (print_ty info)) ty_args
+        (print_ty info) res
+        (print_list space (print_lident info)) id_args
+
   and print_let_def info fmt = function
     | Lvar (pv, e) ->
       fprintf fmt "@[<hov 2>let %a =@ %a@]"
@@ -326,12 +343,13 @@ module Print = struct
       forget_vars args
     | Lrec rdef ->
       let print_one fst fmt = function
-        | { rec_sym = rs1; rec_args = args; rec_exp = e; rec_res = res } ->
-          fprintf fmt "@[<hov 2>%s %a @[%a@] :@ %a@ =@ %a@]"
+        | { rec_sym = rs1; rec_args = args; rec_exp = e;
+            rec_res = res; rec_svar = s } ->
+          fprintf fmt "@[<hov 2>%s %a @[%a@] %a@]"
             (if fst then "let rec" else "and")
             (print_lident info) rs1.rs_name
-            (print_list space (print_vs_arg info)) args
-            (print_ty info) res (print_expr info) e;
+            (print_fun_type_args info) (args, s, res)
+            (print_expr info) e;
           forget_vars args
       in
       List.iter (fun fd -> Hrs.replace ht_rs fd.rec_rsym fd.rec_sym) rdef;