Add args and constraint to class_infos (#2502)

Author: gpetiot

lib/Ast.ml

@@ -1138,8 +1138,12 @@ end = struct
           | Pcl_constraint (_, x) -> x == cty
           | Pcl_extension _ -> false
           | Pcl_open _ -> false )
-    | Cd _ -> assert false
-    | Ctd ctx -> assert (ctx.pci_expr == cty)
+    | Cd ctx ->
+        assert (Option.exists ctx.pci_constraint ~f:(fun x -> x == cty))
+    | Ctd ctx ->
+        assert (
+          Option.exists ctx.pci_constraint ~f:(fun x -> x == cty)
+          || ctx.pci_expr == cty )
     | Clf _ -> assert false
     | Ctf {pctf_desc; _} ->
         assert (

lib/Fmt_ast.ml

@@ -3741,6 +3741,7 @@ and fmt_signature_item c ?ext {ast= si; _} =
 
 and fmt_class_types ?ext c ~pre ~sep cls =
   list_fl cls (fun ~first ~last:_ cl ->
+      (* [pci_args] and [pci_constraint] are not used for class types. *)
       update_config_maybe_disabled c cl.pci_loc cl.pci_attributes
       @@ fun c ->
       let ctx = Ctd cl in
@@ -3773,19 +3774,12 @@ and fmt_class_exprs ?ext c cls =
          update_config_maybe_disabled c cl.pci_loc cl.pci_attributes
          @@ fun c ->
          let ctx = Cd cl in
-         let xargs, xbody = Sugar.cl_fun c.cmts (sub_cl ~ctx cl.pci_expr) in
-         let ty, e =
-           match xbody.ast with
-           | {pcl_desc= Pcl_constraint (e, t); _} as ce ->
-               let ctx = Cl ce in
-               (Some (sub_cty ~ctx t), sub_cl ~ctx e)
-           | _ -> (None, xbody)
-         in
+         let xargs = cl.pci_args in
          let doc_before, doc_after, atrs =
            let force_before = not (Cl.is_simple cl.pci_expr) in
            fmt_docstring_around_item ~force_before c cl.pci_attributes
          in
-         let class_exprs =
+         let class_expr =
            let pro =
              box_fun_decl_args c 2
                ( hovbox 2
@@ -3799,19 +3793,20 @@ and fmt_class_exprs ?ext c cls =
                $ wrap_fun_decl_args c (fmt_fun_args c xargs) )
            in
            let intro =
-             match ty with
+             match cl.pci_constraint with
              | Some ty ->
-                 let pro = pro $ fmt " :@ " in
-                 fmt_class_type c ~pro ty
+                 fmt_class_type c ~pro:(pro $ fmt " :@ ") (sub_cty ~ctx ty)
              | None -> pro
            in
            hovbox 2
-             (hovbox 2 (intro $ fmt "@ =") $ fmt "@;" $ fmt_class_expr c e)
+             ( hovbox 2 (intro $ fmt "@ =")
+             $ fmt "@;"
+             $ fmt_class_expr c (sub_cl ~ctx cl.pci_expr) )
            $ fmt_item_attributes c ~pre:(Break (1, 0)) atrs
          in
          fmt_if (not first) "\n@;<1000 0>"
          $ hovbox 0
-           @@ Cmts.fmt c cl.pci_loc (doc_before $ class_exprs $ doc_after) )
+           @@ Cmts.fmt c cl.pci_loc (doc_before $ class_expr $ doc_after) )
 
 and fmt_module c ctx ?rec_ ?epi ?(can_sparse = false) keyword ?(eqty = "=")
     name xargs xbody xmty ~attrs ~rec_flag =

lib/Sugar.ml

@@ -50,16 +50,6 @@ let fun_ cmts ?(will_keep_first_ast_node = true) xexp =
   in
   fun_ ~will_keep_first_ast_node xexp
 
-let cl_fun cmts ({ast= exp; _} as xexp) =
-  let ctx = Cl exp in
-  match (exp.pcl_attributes, exp.pcl_desc) with
-  | [], Pcl_fun (p, body) ->
-      let before = (List.hd_exn p).pparam_loc in
-      let after = body.pcl_loc in
-      Cmts.relocate cmts ~src:exp.pcl_loc ~before ~after ;
-      (p, sub_cl ~ctx body)
-  | _ -> ([], xexp)
-
 module Exp = struct
   let infix cmts prec xexp =
     let assoc = Option.value_map prec ~default:Assoc.Non ~f:Assoc.of_prec in

lib/Sugar.mli

@@ -22,10 +22,6 @@ val fun_ :
     and the body of the function [exp]. [will_keep_first_ast_node] is set by
     default, otherwise the [exp] is returned without modification. *)
 
-val cl_fun :
-  Cmts.t -> class_expr Ast.xt -> function_param list * class_expr Ast.xt
-(** [cl_fun cmts exp] returns the list of arguments and the body of the function [exp]. *)
-
 module Exp : sig
   val infix :
        Cmts.t

vendor/parser-extended/ast_helper.ml

@@ -454,11 +454,14 @@ end
 module Ci = struct
   let mk ?(loc = !default_loc) ?(attrs = [])
         ?(docs = empty_docs) ?(text = [])
+        ?(args = []) ?constraint_
         ?(virt = Concrete) ?(params = []) name expr =
     {
      pci_virt = virt;
      pci_params = params;
      pci_name = name;
+     pci_args = args;
+     pci_constraint = constraint_;
      pci_expr = expr;
      pci_attributes =
        add_text_attrs text (add_docs_attrs docs attrs);

vendor/parser-extended/ast_mapper.ml

@@ -739,12 +739,14 @@ module CE = struct
     }
 
   let class_infos sub f {pci_virt; pci_params = pl; pci_name; pci_expr;
-                         pci_loc; pci_attributes} =
+                         pci_loc; pci_attributes; pci_args; pci_constraint} =
     let loc = sub.location sub pci_loc in
     let attrs = sub.attributes sub pci_attributes in
     Ci.mk ~loc ~attrs
      ~virt:(Flag.map_virtual sub pci_virt)
      ~params:(List.map (map_fst (sub.typ sub)) pl)
+     ~args:(List.map (map_function_param sub) pci_args)
+     ?constraint_:(map_opt (sub.class_type sub) pci_constraint)
       (map_loc sub pci_name)
       (f pci_expr)
 end

vendor/parser-extended/parser.mly

@@ -1773,14 +1773,15 @@ module_type_subst:
   virt = virtual_flag
   params = formal_class_parameters
   id = mkrhs(LIDENT)
-  body = class_fun_binding
+  cfb = class_fun_binding
   attrs2 = post_item_attributes
   {
     let attrs = attrs1 @ attrs2 in
     let loc = make_loc $sloc in
     let docs = symbol_docs $sloc in
+    let (args, constraint_, body) = cfb in
     ext,
-    Ci.mk id body ~virt ~params ~attrs ~loc ~docs
+    Ci.mk id body ~virt ~params ~attrs ~loc ~docs ~args ?constraint_
   }
 ;
 %inline and_class_declaration:
@@ -1789,14 +1790,15 @@ module_type_subst:
   virt = virtual_flag
   params = formal_class_parameters
   id = mkrhs(LIDENT)
-  body = class_fun_binding
+  cfb = class_fun_binding
   attrs2 = post_item_attributes
   {
     let attrs = attrs1 @ attrs2 in
     let loc = make_loc $sloc in
     let docs = symbol_docs $sloc in
     let text = symbol_text $symbolstartpos in
-    Ci.mk id body ~virt ~params ~attrs ~loc ~text ~docs
+    let (args, constraint_, body) = cfb in
+    Ci.mk id body ~virt ~params ~attrs ~loc ~text ~docs ~args ?constraint_
   }
 ;
 
@@ -1805,18 +1807,7 @@ class_fun_binding:
   ct = ioption(COLON class_type { $2 })
   EQUAL
   ce = class_expr
-  {
-    let ce =
-      match ct with
-      | Some ct ->
-        let loc = ($startpos(ct), $endpos(ce)) in
-        mkclass ~loc (Pcl_constraint (ce, ct))
-      | None -> ce
-    in
-    match params with
-    | [] -> ce
-    | _ :: _ -> mkclass ~loc:$sloc (Pcl_fun (params, ce))
-  }
+    { params, ct, ce }
 ;
 
 formal_class_parameters:

vendor/parser-extended/parsetree.mli

@@ -750,6 +750,8 @@ and 'a class_infos =
      pci_virt: virtual_flag;
      pci_params: (core_type * variance_and_injectivity) list;
      pci_name: string loc;
+     pci_args: function_param list;
+     pci_constraint: class_type option;
      pci_expr: 'a;
      pci_loc: Location.t;
      pci_attributes: attributes;  (** [... [\@\@id1] [\@\@id2]] *)

vendor/parser-extended/printast.ml

@@ -703,6 +703,9 @@ and class_infos : 'a. _ -> (_ -> _ -> 'a -> _) -> _ -> _ -> 'a class_infos -> _
   line i ppf "pci_params =\n";
   list (i+1) type_parameter ppf x.pci_params;
   line i ppf "pci_name = %a\n" fmt_string_loc x.pci_name;
+  line i ppf "pci_args =\n";
+  list (i+1) function_param ppf x.pci_args;
+  line i ppf "pci_constraint = %a\n" (fmt_opt (class_type i)) x.pci_constraint;
   line i ppf "pci_expr =\n";
   f (i+1) ppf x.pci_expr