Define type function_param and replace Sugar.arg_kind (#2466)

* Define type 'function_param'
* Add Fp in the Ast context
* Replace Sugar.arg_kind with function_param

lib/Ast.ml

@@ -638,6 +638,7 @@ module T = struct
     | Cty of class_type
     | Pat of pattern
     | Exp of expression
+    | Fp of function_param
     | Lb of value_binding
     | Mb of module_binding
     | Md of module_declaration
@@ -658,6 +659,7 @@ 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
     | Lb b -> Format.fprintf fs "Lb:@\n%a" Printast.value_binding b
     | Mb m -> Format.fprintf fs "Mb:@\n%a" Printast.module_binding m
     | Md m -> Format.fprintf fs "Md:@\n%a" Printast.module_declaration m
@@ -690,6 +692,7 @@ let attributes = function
   | Cty x -> x.pcty_attributes
   | Pat x -> x.ppat_attributes
   | Exp x -> x.pexp_attributes
+  | Fp _ -> []
   | Lb x -> x.pvb_attributes
   | Mb x -> attrs_of_ext_attrs x.pmb_ext_attrs
   | Md x -> attrs_of_ext_attrs x.pmd_ext_attrs
@@ -711,6 +714,7 @@ let location = function
   | Cty x -> x.pcty_loc
   | Pat x -> x.ppat_loc
   | Exp x -> x.pexp_loc
+  | Fp x -> x.pparam_loc
   | Lb x -> x.pvb_loc
   | Mb x -> x.pmb_loc
   | Md x -> x.pmd_loc
@@ -987,6 +991,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
     | Lb _ -> assert false
     | Mb _ -> assert false
     | Md _ -> assert false
@@ -1102,6 +1107,7 @@ end = struct
     in
     match (ctx : t) with
     | Exp _ -> assert false
+    | Fp _ -> assert false
     | Lb _ -> assert false
     | Mb _ -> assert false
     | Md _ -> assert false
@@ -1169,6 +1175,7 @@ end = struct
   let check_cl {ctx; ast= cl} =
     match (ctx : t) with
     | Exp _ -> assert false
+    | Fp _ -> assert false
     | Lb _ -> assert false
     | Mb _ -> assert false
     | Md _ -> assert false
@@ -1288,6 +1295,11 @@ end = struct
               | _ -> false ) )
       | Pexp_for (p, _, _, _, _) | Pexp_fun (_, _, p, _) -> assert (p == pat)
       )
+    | Fp ctx ->
+        assert (
+          match ctx.pparam_desc with
+          | Pparam_val (_, _, p) -> p == pat
+          | Pparam_newtype _ -> false )
     | Lb x -> assert (x.pvb_pat == pat)
     | Mb _ -> assert false
     | Md _ -> assert false
@@ -1412,6 +1424,11 @@ 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 (
+          match ctx.pparam_desc with
+          | Pparam_val (_, e, _) -> Option.exists e ~f:(fun x -> x == exp)
+          | Pparam_newtype _ -> false )
     | Lb x -> assert (x.pvb_expr == exp)
     | Mb _ -> assert false
     | Md _ -> assert false
@@ -1662,6 +1679,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= Lb _; ast= _}
      |{ctx= _; ast= Lb _}
      |{ctx= Td _; ast= _}
@@ -1745,6 +1764,7 @@ end = struct
       | Pexp_field _ -> Some Dot
       | Pexp_send _ -> Some Dot
       | _ -> None )
+    | Fp _ -> None
     | Lb _ -> None
     | Cl c -> (
       match c.pcl_desc with
@@ -1916,7 +1936,10 @@ end = struct
         | Ppat_variant _ ) ) ->
         true
     | (Str _ | Exp _), Ppat_lazy _ -> true
-    | ( Pat {ppat_desc= Ppat_construct _ | Ppat_variant _; _}
+    | ( Fp _
+      , ( Ppat_tuple _ | Ppat_construct _ | Ppat_alias _ | Ppat_variant _
+        | Ppat_lazy _ | Ppat_exception _ | Ppat_or _ ) )
+     |( Pat {ppat_desc= Ppat_construct _ | Ppat_variant _; _}
       , (Ppat_construct (_, Some _) | Ppat_cons _ | Ppat_variant (_, Some _))
       ) ->
         true

lib/Ast.mli

@@ -110,6 +110,7 @@ type t =
   | Cty of class_type
   | Pat of pattern
   | Exp of expression
+  | Fp of function_param
   | Lb of value_binding
   | Mb of module_binding
   | Md of module_declaration

lib/Fmt_ast.ml

@@ -1234,29 +1234,30 @@ and fmt_pattern ?ext c ?pro ?parens ?(box = false)
             (fmt "@;<0 2>" $ fmt_pattern c (sub_pat ~ctx pat)) )
 
 and fmt_fun_args c args =
-  let fmt_fun_arg (a : Sugar.arg_kind) =
-    match a with
-    | Val
+  let fmt_fun_arg (a : function_param) =
+    let ctx = Fp a in
+    match a.pparam_desc with
+    | Pparam_val
         ( ((Labelled l | Optional l) as lbl)
-        , ( { ast=
-                { ppat_desc=
-                    ( Ppat_var {txt; loc= _}
-                    | Ppat_constraint
-                        ( { ppat_desc= Ppat_var {txt; loc= _}
-                          ; ppat_attributes= []
-                          ; _ }
-                        , _ ) )
-                ; ppat_attributes= []
-                ; _ }
-            ; _ } as xpat )
-        , None )
+        , None
+        , ( { ppat_desc=
+                ( Ppat_var {txt; loc= _}
+                | Ppat_constraint
+                    ( { ppat_desc= Ppat_var {txt; loc= _}
+                      ; ppat_attributes= []
+                      ; _ }
+                    , _ ) )
+            ; ppat_attributes= []
+            ; _ } as pat ) )
       when String.equal l.txt txt ->
         let symbol = match lbl with Labelled _ -> "~" | _ -> "?" in
+        let xpat = sub_pat ~ctx pat in
         cbox 0 (str symbol $ fmt_pattern ~box:true c xpat)
-    | Val ((Optional _ as lbl), xpat, None) ->
-        let has_attr = not (List.is_empty xpat.ast.ppat_attributes) in
+    | Pparam_val ((Optional _ as lbl), None, pat) ->
+        let xpat = sub_pat ~ctx pat in
+        let has_attr = not (List.is_empty pat.ppat_attributes) in
         let outer_parens, inner_parens =
-          match xpat.ast.ppat_desc with
+          match pat.ppat_desc with
           | Ppat_any | Ppat_var _ -> (false, false)
           | Ppat_unpack _ -> (not has_attr, true)
           | Ppat_tuple _ -> (false, true)
@@ -1268,35 +1269,39 @@ and fmt_fun_args c args =
           $ hovbox 0
             @@ Params.parens_if outer_parens c.conf
                  (fmt_pattern ~parens:inner_parens c xpat) )
-    | Val (((Labelled _ | Nolabel) as lbl), xpat, None) ->
+    | Pparam_val (((Labelled _ | Nolabel) as lbl), None, pat) ->
+        let xpat = sub_pat ~ctx pat in
         cbox 2 (fmt_label lbl ":@," $ fmt_pattern c xpat)
-    | Val
+    | Pparam_val
         ( Optional l
-        , ( { ast= {ppat_desc= Ppat_var {txt; loc= _}; ppat_attributes= []; _}
-            ; _ } as xpat )
-        , Some xexp )
+        , Some exp
+        , ({ppat_desc= Ppat_var {txt; loc= _}; ppat_attributes= []; _} as pat)
+        )
       when String.equal l.txt txt ->
+        let xexp = sub_exp ~ctx exp in
+        let xpat = sub_pat ~ctx pat in
         cbox 0
           (wrap "?(" ")"
              ( fmt_pattern c ~box:true xpat
              $ fmt " =@;<1 2>"
              $ hovbox 2 (fmt_expression c xexp) ) )
-    | Val
+    | Pparam_val
         ( Optional l
-        , ( { ast=
-                { ppat_desc=
-                    Ppat_constraint
-                      ({ppat_desc= Ppat_var {txt; loc= _}; _}, _)
-                ; ppat_attributes= []
-                ; _ }
-            ; _ } as xpat )
-        , Some xexp )
+        , Some exp
+        , ( { ppat_desc=
+                Ppat_constraint ({ppat_desc= Ppat_var {txt; loc= _}; _}, _)
+            ; ppat_attributes= []
+            ; _ } as pat ) )
       when String.equal l.txt txt ->
+        let xexp = sub_exp ~ctx exp in
+        let xpat = sub_pat ~ctx pat in
         cbox 0
           (wrap "?(" ")"
              ( fmt_pattern c ~parens:false ~box:true xpat
              $ fmt " =@;<1 2>" $ fmt_expression c xexp ) )
-    | Val (Optional l, xpat, Some xexp) ->
+    | Pparam_val (Optional l, Some exp, pat) ->
+        let xexp = sub_exp ~ctx exp in
+        let xpat = sub_pat ~ctx pat in
         let parens =
           match xpat.ast.ppat_desc with
           | Ppat_unpack _ -> None
@@ -1307,10 +1312,10 @@ and fmt_fun_args c args =
           $ wrap_k (fmt ":@,(") (str ")")
               ( fmt_pattern c ?parens ~box:true xpat
               $ fmt " =@;<1 2>" $ fmt_expression c xexp ) )
-    | Val ((Labelled _ | Nolabel), _, Some _) ->
+    | Pparam_val ((Labelled _ | Nolabel), Some _, _) ->
         impossible "not accepted by parser"
-    | Newtypes [] -> impossible "not accepted by parser"
-    | Newtypes names ->
+    | Pparam_newtype [] -> impossible "not accepted by parser"
+    | Pparam_newtype names ->
         cbox 0
           (Params.parens c.conf
              (str "type " $ list names "@ " (fmt_str_loc c)) )

lib/Sugar.ml

@@ -14,9 +14,37 @@ open Asttypes
 open Ast
 open Extended_ast
 
-type arg_kind =
-  | Val of arg_label * pattern xt * expression xt option
-  | Newtypes of string loc list
+(* Temporary. Won't be necessary once the type [function_param] is used in
+   [Pexp_fun] and [Pcl_fun]. *)
+let mk_function_param pparam_desc =
+  let pparam_loc =
+    let init, locs =
+      match pparam_desc with
+      | Pparam_val (lbl, e, p) ->
+          let locs =
+            match lbl with
+            | Nolabel -> []
+            | Labelled x -> [x.loc]
+            | Optional x -> [x.loc]
+          in
+          let locs =
+            match e with Some e -> e.pexp_loc :: locs | None -> locs
+          in
+          (p.ppat_loc, locs)
+      | Pparam_newtype types -> (
+        match types with
+        | [] -> failwith "Pparam_newtype always contains at least one type"
+        | hd :: tl ->
+            let locs = List.map tl ~f:(fun x -> x.loc) in
+            (hd.loc, locs) )
+    in
+    let min acc x = if Location.compare_start acc x < 0 then acc else x in
+    let max acc x = if Location.compare_end acc x > 0 then acc else x in
+    let loc_start = (List.fold_left locs ~init ~f:min).loc_start in
+    let loc_end = (List.fold_left locs ~init ~f:max).loc_end in
+    {Location.loc_start; loc_end; loc_ghost= true}
+  in
+  {pparam_desc; pparam_loc}
 
 let fun_ cmts ?(will_keep_first_ast_node = true) xexp =
   let rec fun_ ?(will_keep_first_ast_node = false) ({ast= exp; _} as xexp) =
@@ -29,11 +57,7 @@ let fun_ cmts ?(will_keep_first_ast_node = true) xexp =
             Cmts.relocate cmts ~src:pexp_loc ~before:pattern.ppat_loc
               ~after:body.pexp_loc ;
           let xargs, xbody = fun_ (sub_exp ~ctx body) in
-          ( Val
-              ( label
-              , sub_pat ~ctx pattern
-              , Option.map default ~f:(sub_exp ~ctx) )
-            :: xargs
+          ( mk_function_param (Pparam_val (label, default, pattern)) :: xargs
           , xbody )
       | Pexp_newtype (name, body) ->
           if not will_keep_first_ast_node then
@@ -42,8 +66,9 @@ let fun_ cmts ?(will_keep_first_ast_node = true) xexp =
           let xargs, xbody = fun_ (sub_exp ~ctx body) in
           let xargs =
             match xargs with
-            | Newtypes names :: xargs -> Newtypes (name :: names) :: xargs
-            | xargs -> Newtypes [name] :: xargs
+            | {pparam_desc= Pparam_newtype names; _} :: xargs ->
+                mk_function_param (Pparam_newtype (name :: names)) :: xargs
+            | xargs -> mk_function_param (Pparam_newtype [name]) :: xargs
           in
           (xargs, xbody)
       | _ -> ([], xexp)
@@ -62,11 +87,7 @@ let cl_fun ?(will_keep_first_ast_node = true) cmts xexp =
             Cmts.relocate cmts ~src:pcl_loc ~before:pattern.ppat_loc
               ~after:body.pcl_loc ;
           let xargs, xbody = fun_ (sub_cl ~ctx body) in
-          ( Val
-              ( label
-              , sub_pat ~ctx pattern
-              , Option.map default ~f:(sub_exp ~ctx) )
-            :: xargs
+          ( mk_function_param (Pparam_val (label, default, pattern)) :: xargs
           , xbody )
       | _ -> ([], xexp)
     else ([], xexp)
@@ -220,7 +241,7 @@ module Let_binding = struct
   type t =
     { lb_op: string loc
     ; lb_pat: pattern xt
-    ; lb_args: arg_kind list
+    ; lb_args: function_param list
     ; lb_typ:
         [ `Polynewtype of label loc list * core_type xt
         | `Coerce of core_type xt option * core_type xt

lib/Sugar.mli

@@ -13,15 +13,11 @@ open Migrate_ast
 open Asttypes
 open Extended_ast
 
-type arg_kind =
-  | Val of arg_label * pattern Ast.xt * expression Ast.xt option
-  | Newtypes of string loc list
-
 val fun_ :
      Cmts.t
   -> ?will_keep_first_ast_node:bool
   -> expression Ast.xt
-  -> arg_kind list * expression Ast.xt
+  -> function_param list * expression Ast.xt
 (** [fun_ cmts will_keep_first_ast_node exp] returns the list of arguments
     and the body of the function [exp]. [will_keep_first_ast_node] is set by
     default, otherwise the [exp] is returned without modification. *)
@@ -30,7 +26,7 @@ val cl_fun :
      ?will_keep_first_ast_node:bool
   -> Cmts.t
   -> class_expr Ast.xt
-  -> arg_kind list * class_expr Ast.xt
+  -> function_param list * class_expr Ast.xt
 (** [cl_fun will_keep_first_ast_node cmts exp] returns the list of arguments
     and the body of the function [exp]. [will_keep_first_ast_node] is set by
     default, otherwise the [exp] is returned without modification. *)
@@ -62,7 +58,7 @@ module Let_binding : sig
   type t =
     { lb_op: string loc
     ; lb_pat: pattern Ast.xt
-    ; lb_args: arg_kind list
+    ; lb_args: function_param list
     ; lb_typ:
         [ `Polynewtype of label loc list * core_type Ast.xt
         | `Coerce of core_type Ast.xt option * core_type Ast.xt

vendor/parser-extended/ast_mapper.ml

@@ -467,6 +467,20 @@ end
 module E = struct
   (* Value expressions for the core language *)
 
+  let map_function_param sub { pparam_loc = loc; pparam_desc = desc } =
+    let loc = sub.location sub loc in
+    let desc =
+      match desc with
+      | Pparam_val (lab, def, p) ->
+          Pparam_val
+            (lab,
+             map_opt (sub.expr sub) def,
+             sub.pat sub p)
+      | Pparam_newtype ty ->
+          Pparam_newtype (List.map (map_loc sub) ty)
+    in
+    { pparam_loc = loc; pparam_desc = desc }
+
   let map_constraint sub c =
     match c with
     | Pconstraint ty -> Pconstraint (sub.typ sub ty)