Distinguish fun_param for expr/class, unify parsing of val/newtype fun_param

CHANGES.md

@@ -33,6 +33,7 @@ profile. This started with version 0.26.0.
 - Fix unstable comment around docked functor argument (#2506, @Julow)
 - \* Fix unwanted alignment after comment (#2507, @Julow)
 - \* Fix unwanted alignment in if-then-else (#2511, @Julow)
+- Fix position of comments around and within `(type ...)` function arguments (#2503, @gpetiot)
 
 ## 0.26.1 (2023-09-15)
 

lib/Ast.ml

@@ -640,7 +640,8 @@ module T = struct
     | Ctd of class_type_declaration
     | Pat of pattern
     | Exp of expression
-    | Fp of function_param
+    | Fpe of expr_function_param
+    | Fpc of class_function_param
     | Vc of value_constraint
     | Lb of value_binding
     | Bo of binding_op
@@ -663,7 +664,8 @@ module T = struct
     | Td t -> Format.fprintf fs "Td:@\n%a" Printast.type_declaration t
     | Pat p -> Format.fprintf fs "Pat:@\n%a" Printast.pattern p
     | Exp e -> Format.fprintf fs "Exp:@\n%a" Printast.expression e
-    | Fp p -> Format.fprintf fs "Fp:@\n%a" Printast.function_param p
+    | Fpe p -> Format.fprintf fs "Fpe:@\n%a" Printast.expr_function_param p
+    | Fpc p -> Format.fprintf fs "Fpc:@\n%a" Printast.class_function_param p
     | Vc c -> Format.fprintf fs "Vc:@\n%a" Printast.value_constraint c
     | Lb b -> Format.fprintf fs "Lb:@\n%a" Printast.value_binding b
     | Bo b -> Format.fprintf fs "Bo:@\n%a" Printast.binding_op b
@@ -701,7 +703,7 @@ let attributes = function
   | Cty x -> x.pcty_attributes
   | Pat x -> x.ppat_attributes
   | Exp x -> x.pexp_attributes
-  | Fp _ -> []
+  | Fpe _ | Fpc _ -> []
   | Vc _ -> []
   | Lb x -> x.pvb_attributes
   | Bo _ -> []
@@ -727,7 +729,8 @@ let location = function
   | Cty x -> x.pcty_loc
   | Pat x -> x.ppat_loc
   | Exp x -> x.pexp_loc
-  | Fp x -> x.pparam_loc
+  | Fpe x -> x.pparam_loc
+  | Fpc x -> x.pparam_loc
   | Vc _ -> Location.none
   | Lb x -> x.pvb_loc
   | Bo x -> x.pbop_loc
@@ -1014,7 +1017,7 @@ end = struct
                   | Pcoerce (t1, t2) -> Option.exists t1 ~f || f t2 ) ) )
       | Pexp_let (lbs, _) -> assert (check_let_bindings lbs)
       | _ -> assert false )
-    | Fp _ -> assert false
+    | Fpe _ | Fpc _ -> assert false
     | Vc c -> assert (check_value_constraint c)
     | Lb _ -> assert false
     | Bo _ -> assert false
@@ -1106,7 +1109,7 @@ end = struct
   let check_cty {ctx; ast= cty} =
     match (ctx : t) with
     | Exp _ -> assert false
-    | Fp _ -> assert false
+    | Fpe _ | Fpc _ -> assert false
     | Vc _ -> assert false
     | Lb _ -> assert false
     | Bo _ -> assert false
@@ -1165,7 +1168,7 @@ end = struct
   let check_cl {ctx; ast= cl} =
     match (ctx : t) with
     | Exp _ -> assert false
-    | Fp _ -> assert false
+    | Fpe _ | Fpc _ -> assert false
     | Vc _ -> assert false
     | Lb _ -> assert false
     | Bo _ -> assert false
@@ -1217,12 +1220,18 @@ end = struct
     let check_bindings l =
       List.exists l ~f:(fun {pvb_pat; _} -> check_subpat pvb_pat)
     in
-    let check_function_param param =
+    let check_param_val (_, _, p) = p == pat in
+    let check_expr_function_param param =
       match param.pparam_desc with
-      | Pparam_val (_, _, p) -> p == pat
-      | Pparam_newtype _ -> false
+      | Param_val x -> check_param_val x
+      | Param_newtype _ -> false
+    in
+    let check_class_function_param param =
+      check_param_val param.pparam_desc
+    in
+    let check_class_function_params =
+      List.exists ~f:check_class_function_param
     in
-    let check_function_params l = List.exists l ~f:check_function_param in
     match ctx with
     | Pld (PPat (p1, _)) -> assert (p1 == pat)
     | Pld _ -> assert false
@@ -1280,8 +1289,9 @@ end = struct
               | {pc_lhs; _} when pc_lhs == pat -> true
               | _ -> false ) )
       | Pexp_for (p, _, _, _, _) -> assert (p == pat)
-      | Pexp_fun (p, _) -> assert (check_function_param p) )
-    | Fp ctx -> assert (check_function_param ctx)
+      | Pexp_fun (p, _) -> assert (check_expr_function_param p) )
+    | Fpe ctx -> assert (check_expr_function_param ctx)
+    | Fpc ctx -> assert (check_class_function_param ctx)
     | Vc _ -> assert false
     | Lb x -> assert (x.pvb_pat == pat)
     | Bo x -> assert (x.pbop_pat == pat)
@@ -1290,7 +1300,7 @@ end = struct
     | Cl ctx ->
         assert (
           match ctx.pcl_desc with
-          | Pcl_fun (p, _) -> check_function_params p
+          | Pcl_fun (p, _) -> check_class_function_params p
           | Pcl_constr _ -> false
           | Pcl_structure {pcstr_self; _} ->
               Option.exists ~f:(fun self_ -> self_ == pat) pcstr_self
@@ -1332,12 +1342,18 @@ end = struct
       | PStr [{pstr_desc= Pstr_eval (e, _); _}] -> e == exp
       | _ -> false
     in
-    let check_function_param param =
+    let check_param_val (_, e, _) = Option.exists e ~f:(fun x -> x == exp) in
+    let check_expr_function_param param =
       match param.pparam_desc with
-      | Pparam_val (_, e, _) -> Option.exists e ~f:(fun x -> x == exp)
-      | Pparam_newtype _ -> false
+      | Param_val x -> check_param_val x
+      | Param_newtype _ -> false
+    in
+    let check_class_function_param param =
+      check_param_val param.pparam_desc
+    in
+    let check_class_function_params =
+      List.exists ~f:check_class_function_param
     in
-    let check_function_params l = List.exists l ~f:check_function_param in
     match ctx with
     | Pld (PPat (_, Some e1)) -> assert (e1 == exp)
     | Pld _ -> assert false
@@ -1367,7 +1383,7 @@ end = struct
                 | {pc_rhs; _} when pc_rhs == exp -> true
                 | _ -> false ) )
         | Pexp_fun (param, body) ->
-            assert (check_function_param param || body == exp)
+            assert (check_expr_function_param param || body == exp)
         | Pexp_indexop_access {pia_lhs; pia_kind= Builtin idx; pia_rhs; _} ->
             assert (
               pia_lhs == exp || idx == exp
@@ -1416,7 +1432,8 @@ end = struct
         | Pexp_for (_, e1, e2, _, e3) ->
             assert (e1 == exp || e2 == exp || e3 == exp)
         | Pexp_override e1N -> assert (List.exists e1N ~f:snd_f) )
-    | Fp ctx -> assert (check_function_param ctx)
+    | Fpe ctx -> assert (check_expr_function_param ctx)
+    | Fpc ctx -> assert (check_class_function_param ctx)
     | Vc _ -> assert false
     | Lb x -> assert (x.pvb_expr == exp)
     | Bo x -> assert (x.pbop_exp == exp)
@@ -1439,7 +1456,7 @@ end = struct
     | Cl ctx ->
         let rec loop ctx =
           match ctx.pcl_desc with
-          | Pcl_fun (param, e) -> check_function_params param || loop e
+          | Pcl_fun (param, e) -> check_class_function_params param || loop e
           | Pcl_constr _ -> false
           | Pcl_structure _ -> false
           | Pcl_apply (_, l) -> List.exists l ~f:(fun (_, e) -> e == exp)
@@ -1670,8 +1687,8 @@ end = struct
       ; ast=
           ( Pld _ | Top | Tli _ | Pat _ | Cl _ | Mty _ | Mod _ | Sig _
           | Str _ | Clf _ | Ctf _ | Rep | Mb _ | Md _ ) }
-     |{ctx= Fp _; ast= _}
-     |{ctx= _; ast= Fp _}
+     |{ctx= Fpe _ | Fpc _; ast= _}
+     |{ctx= _; ast= Fpe _ | Fpc _}
      |{ctx= Vc _; ast= _}
      |{ctx= _; ast= Vc _}
      |{ctx= Lb _; ast= _}
@@ -1763,7 +1780,7 @@ end = struct
       | Pexp_field _ -> Some Dot
       | Pexp_send _ -> Some Dot
       | _ -> None )
-    | Fp _ -> None
+    | Fpe _ | Fpc _ -> None
     | Vc _ -> None
     | Lb _ -> None
     | Bo _ -> None
@@ -1884,7 +1901,8 @@ end = struct
       | Ppat_cons _ -> true
       | Ppat_construct _ | Ppat_record _ | Ppat_variant _ -> false
       | _ -> true )
-    | Fp {pparam_desc= Pparam_val (_, _, _); _}, Ppat_cons _ -> true
+    | Fpe {pparam_desc= Param_val (_, _, _); _}, Ppat_cons _ -> true
+    | Fpc {pparam_desc= _; _}, Ppat_cons _ -> true
     | Pat {ppat_desc= Ppat_construct _; _}, Ppat_cons _ -> true
     | _, Ppat_constraint (_, {ptyp_desc= Ptyp_poly _; _}) -> false
     | ( Exp {pexp_desc= Pexp_letop _; _}
@@ -1945,7 +1963,7 @@ end = struct
         | Ppat_variant _ ) ) ->
         true
     | (Str _ | Exp _), Ppat_lazy _ -> true
-    | ( Fp _
+    | ( (Fpe _ | Fpc _)
       , ( Ppat_tuple _ | Ppat_construct _ | Ppat_alias _ | Ppat_variant _
         | Ppat_lazy _ | Ppat_exception _ | Ppat_or _ ) )
      |( Pat {ppat_desc= Ppat_construct _ | Ppat_variant _; _}

lib/Ast.mli

@@ -112,7 +112,8 @@ type t =
   | Ctd of class_type_declaration
   | Pat of pattern
   | Exp of expression
-  | Fp of function_param
+  | Fpe of expr_function_param
+  | Fpc of class_function_param
   | Vc of value_constraint
   | Lb of value_binding
   | Bo of binding_op

lib/Extended_ast.ml

@@ -183,7 +183,7 @@ module Parse = struct
           {p with ppat_desc= Ppat_unpack (name, Some pt)}
       | p -> Ast_mapper.default_mapper.pat m p
     in
-    let expr (m : Ast_mapper.mapper) = function
+    let rec expr (m : Ast_mapper.mapper) = function
       | {pexp_desc= Pexp_cons (_ :: _ :: _ :: _ as l); _} as e
         when match List.last_exn l with
              (* Empty lists are always represented as Lident [] *)
@@ -233,6 +233,25 @@ module Parse = struct
              (module S) = (module M)] - [let _ = ((module M) : (module
              S))] *)
           {p with pexp_desc= Pexp_pack (name, Some pt)}
+      | { pexp_desc=
+            Pexp_fun
+              ({pparam_desc= Param_newtype types1; pparam_loc= loc1}, e1)
+        ; pexp_attributes= []
+        ; _ } as e ->
+          let e =
+            match (expr m e1).pexp_desc with
+            | Pexp_fun
+                ({pparam_desc= Param_newtype types2; pparam_loc= loc2}, e2)
+              ->
+                { e with
+                  pexp_desc=
+                    Pexp_fun
+                      ( { pparam_desc= Param_newtype (types1 @ types2)
+                        ; pparam_loc= {loc1 with loc_end= loc2.loc_end} }
+                      , e2 ) }
+            | _ -> e
+          in
+          Ast_mapper.default_mapper.expr m e
       | e -> Ast_mapper.default_mapper.expr m e
     in
     Ast_mapper.{default_mapper with expr; pat; binding_op}