Represent the arity of uncurried function definitions directly in the…

… AST.

analysis/src/CompletionFrontEnd.ml

@@ -1318,7 +1318,7 @@ let completionWithParser1 ~currentFile ~debug ~offset ~path ~posCursor
             match exprToContextPath lhs with
             | Some contextPath -> setResult (Cpath (CPObj (contextPath, label)))
             | None -> ())
-        | Pexp_fun (lbl, defaultExpOpt, pat, e) ->
+        | Pexp_fun (lbl, defaultExpOpt, pat, e, _) ->
           let oldScope = !scope in
           (match (!processingFun, !currentCtxPath) with
           | None, Some ctxPath -> processingFun := Some (ctxPath, 0)

analysis/src/DumpAst.ml

@@ -213,7 +213,7 @@ and printExprItem expr ~pos ~indentation =
       | None -> ""
       | Some expr -> "," ^ printExprItem expr ~pos ~indentation)
     ^ ")"
-  | Pexp_fun (arg, _maybeDefaultArgExpr, pattern, nextExpr) ->
+  | Pexp_fun (arg, _maybeDefaultArgExpr, pattern, nextExpr, _) ->
     "Pexp_fun(\n"
     ^ addIndentation (indentation + 1)
     ^ "arg: "

analysis/src/Xform.ml

@@ -261,7 +261,7 @@ module AddBracesToFn = struct
         | _ -> false
       in
       (match e.pexp_desc with
-      | Pexp_fun (_, _, _, bodyExpr)
+      | Pexp_fun (_, _, _, bodyExpr, _)
         when Loc.hasPos ~pos bodyExpr.pexp_loc
              && isBracedExpr bodyExpr = false
              && isFunction bodyExpr = false ->
@@ -303,10 +303,10 @@ module AddTypeAnnotation = struct
     in
     let rec processFunction ~argNum (e : Parsetree.expression) =
       match e.pexp_desc with
-      | Pexp_fun (argLabel, _, pat, e)
+      | Pexp_fun (argLabel, _, pat, e, _)
       | Pexp_construct
           ( {txt = Lident "Function$"},
-            Some {pexp_desc = Pexp_fun (argLabel, _, pat, e)} ) ->
+            Some {pexp_desc = Pexp_fun (argLabel, _, pat, e, _)} ) ->
         let isUnlabeledOnlyArg =
           argNum = 1 && argLabel = Nolabel
           &&

compiler/frontend/ast_compatible.ml

@@ -63,11 +63,11 @@ let app3 ?(loc = default_loc) ?(attrs = []) fn arg1 arg2 arg3 : expression =
       Pexp_apply (fn, [(Nolabel, arg1); (Nolabel, arg2); (Nolabel, arg3)]);
   }
 
-let fun_ ?(loc = default_loc) ?(attrs = []) pat exp =
+let fun_ ?(loc = default_loc) ?(attrs = []) ~arity pat exp =
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
-    pexp_desc = Pexp_fun (Nolabel, None, pat, exp);
+    pexp_desc = Pexp_fun (Nolabel, None, pat, exp, arity);
   }
 
 let const_exp_string ?(loc = default_loc) ?(attrs = []) ?delimiter (s : string)

compiler/frontend/ast_compatible.mli

@@ -72,7 +72,12 @@ val apply_labels :
 *)
 
 val fun_ :
-  ?loc:Location.t -> ?attrs:attrs -> pattern -> expression -> expression
+  ?loc:Location.t ->
+  ?attrs:attrs ->
+  arity:int option ->
+  pattern ->
+  expression ->
+  expression
 
 (* val opt_label : string -> Asttypes.arg_label *)
 

compiler/frontend/ast_derive_js_mapper.ml

@@ -167,7 +167,7 @@ let init () =
               in
               let to_js_body body =
                 Ast_comb.single_non_rec_value pat_to_js
-                  (Ast_compatible.fun_
+                  (Ast_compatible.fun_ ~arity:None
                      (Pat.constraint_ (Pat.var pat_param) core_type)
                      body)
               in
@@ -211,7 +211,7 @@ let init () =
                 in
                 let from_js =
                   Ast_comb.single_non_rec_value pat_from_js
-                    (Ast_compatible.fun_ (Pat.var pat_param)
+                    (Ast_compatible.fun_ ~arity:None (Pat.var pat_param)
                        (if create_type then
                           Exp.let_ Nonrecursive
                             [Vb.mk (Pat.var pat_param) (exp_param +: new_type)]
@@ -253,7 +253,7 @@ let init () =
                            app2 unsafe_index_get_exp exp_map exp_param
                          else app1 erase_type_exp exp_param);
                       Ast_comb.single_non_rec_value pat_from_js
-                        (Ast_compatible.fun_ (Pat.var pat_param)
+                        (Ast_compatible.fun_ ~arity:None (Pat.var pat_param)
                            (let result =
                               app2 unsafe_index_get_exp rev_exp_map exp_param
                             in

compiler/frontend/ast_derive_projector.ml

@@ -46,7 +46,7 @@ let init () =
                   ->
                     let txt = "param" in
                     Ast_comb.single_non_rec_value ?attrs:gentype_attrs pld_name
-                      (Ast_compatible.fun_
+                      (Ast_compatible.fun_ ~arity:None
                          (Pat.constraint_ (Pat.var {txt; loc}) core_type)
                          (Exp.field
                             (Exp.ident {txt = Lident txt; loc})
@@ -108,7 +108,9 @@ let init () =
                              annotate_type
                          in
                          Ext_list.fold_right vars exp (fun var b ->
-                             Ast_compatible.fun_ (Pat.var {loc; txt = var}) b)
+                             Ast_compatible.fun_ ~arity:None
+                               (Pat.var {loc; txt = var})
+                               b)
                          |> handle_uncurried_accessor_tranform ~loc ~arity))
               | Ptype_abstract | Ptype_open ->
                 Ast_derive_util.not_applicable tdcl.ptype_loc deriving_name;

compiler/frontend/ast_pat.ml

@@ -35,7 +35,7 @@ let is_unit_cont ~yes ~no (p : t) =
 let arity_of_fun (pat : Parsetree.pattern) (e : Parsetree.expression) =
   let rec aux (e : Parsetree.expression) =
     match e.pexp_desc with
-    | Pexp_fun (_, _, _, e) -> 1 + aux e (*FIXME error on optional*)
+    | Pexp_fun (_, _, _, e, _) -> 1 + aux e (*FIXME error on optional*)
     (* | Pexp_fun _
        -> Location.raise_errorf
            ~loc:e.pexp_loc "Label is not allowed in JS object" *)
@@ -45,7 +45,7 @@ let arity_of_fun (pat : Parsetree.pattern) (e : Parsetree.expression) =
 
 let rec labels_of_fun (e : Parsetree.expression) =
   match e.pexp_desc with
-  | Pexp_fun (l, _, _, e) -> l :: labels_of_fun e
+  | Pexp_fun (l, _, _, e, _) -> l :: labels_of_fun e
   | _ -> []
 
 let rec is_single_variable_pattern_conservative (p : t) =

compiler/frontend/ast_uncurry_gen.ml

@@ -36,7 +36,7 @@ let to_method_callback loc (self : Bs_ast_mapper.mapper) label
     match Ast_attributes.process_attributes_rev body.pexp_attributes with
     | Nothing, attrs -> (
       match body.pexp_desc with
-      | Pexp_fun (arg_label, _, arg, body) ->
+      | Pexp_fun (arg_label, _, arg, body, _) ->
         Bs_syntaxerr.optional_err loc arg_label;
         aux ((arg_label, self.pat self arg, attrs) :: acc) body
       | _ -> (self.expr self body, acc))
@@ -45,7 +45,7 @@ let to_method_callback loc (self : Bs_ast_mapper.mapper) label
   let result, rev_extra_args = aux [(label, self_pat, [])] body in
   let body =
     Ext_list.fold_left rev_extra_args result (fun e (label, p, attrs) ->
-        Ast_helper.Exp.fun_ ~loc ~attrs label None p e)
+        Ast_helper.Exp.fun_ ~loc ~attrs ~arity:None label None p e)
   in
   let arity = List.length rev_extra_args in
   let arity_s = string_of_int arity in

compiler/frontend/bs_ast_mapper.ml

@@ -315,8 +315,8 @@ module E = struct
            sub vbs)
         (sub.expr sub e)
     (* #end *)
-    | Pexp_fun (lab, def, p, e) ->
-      fun_ ~loc ~attrs lab
+    | Pexp_fun (lab, def, p, e, arity) ->
+      fun_ ~loc ~attrs ~arity lab
         (map_opt (sub.expr sub) def)
         (sub.pat sub p) (sub.expr sub e)
     | Pexp_function pel -> function_ ~loc ~attrs (sub.cases sub pel)

compiler/frontend/bs_builtin_ppx.ml

@@ -130,7 +130,7 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
     let body = Ast_async.add_async_attribute ~async body in
     let res = self.expr self body in
     {e with pexp_desc = Pexp_newtype (s, res)}
-  | Pexp_fun (label, _, pat, body) -> (
+  | Pexp_fun (label, _, pat, body, _arity) -> (
     let async = Ast_attributes.has_async_payload e.pexp_attributes <> None in
     match Ast_attributes.process_attributes_rev e.pexp_attributes with
     | Nothing, _ ->
@@ -594,7 +594,7 @@ let rec structure_mapper ~await_context (self : mapper) (stru : Ast_structure.t)
             | Pexp_ifthenelse (_, then_expr, Some else_expr) ->
               aux then_expr @ aux else_expr
             | Pexp_construct (_, Some expr) -> aux expr
-            | Pexp_fun (_, _, _, expr) | Pexp_newtype (_, expr) -> aux expr
+            | Pexp_fun (_, _, _, expr, _) | Pexp_newtype (_, expr) -> aux expr
             | _ -> acc
           in
           aux pvb_expr @ spelunk_vbs acc tl

compiler/ml/ast_async.ml

@@ -40,14 +40,14 @@ let add_async_attribute ~async (body : Parsetree.expression) =
 
 let rec add_promise_to_result ~loc (e : Parsetree.expression) =
   match e.pexp_desc with
-  | Pexp_fun (label, eo, pat, body) ->
+  | Pexp_fun (label, eo, pat, body, arity) ->
     let body = add_promise_to_result ~loc body in
-    {e with pexp_desc = Pexp_fun (label, eo, pat, body)}
+    {e with pexp_desc = Pexp_fun (label, eo, pat, body, arity)}
   | _ -> add_promise_type ~loc ~async:true e
 
 let make_function_async ~async (e : Parsetree.expression) =
   if async then
     match e.pexp_desc with
-    | Pexp_fun (_, _, {ppat_loc}, _) -> add_promise_to_result ~loc:ppat_loc e
+    | Pexp_fun (_, _, {ppat_loc}, _, _) -> add_promise_to_result ~loc:ppat_loc e
     | _ -> assert false
   else e

compiler/ml/ast_helper.ml

@@ -150,7 +150,8 @@ module Exp = struct
   let ident ?loc ?attrs a = mk ?loc ?attrs (Pexp_ident a)
   let constant ?loc ?attrs a = mk ?loc ?attrs (Pexp_constant a)
   let let_ ?loc ?attrs a b c = mk ?loc ?attrs (Pexp_let (a, b, c))
-  let fun_ ?loc ?attrs a b c d = mk ?loc ?attrs (Pexp_fun (a, b, c, d))
+  let fun_ ?loc ?attrs ~arity a b c d =
+    mk ?loc ?attrs (Pexp_fun (a, b, c, d, arity))
   let function_ ?loc ?attrs a = mk ?loc ?attrs (Pexp_function a)
   let apply ?loc ?attrs a b = mk ?loc ?attrs (Pexp_apply (a, b))
   let match_ ?loc ?attrs a b = mk ?loc ?attrs (Pexp_match (a, b))

compiler/ml/ast_helper.mli

@@ -132,6 +132,7 @@ module Exp : sig
   val fun_ :
     ?loc:loc ->
     ?attrs:attrs ->
+    arity:int option ->
     arg_label ->
     expression option ->
     pattern ->

compiler/ml/ast_iterator.ml

@@ -282,7 +282,7 @@ module E = struct
     | Pexp_let (_r, vbs, e) ->
       List.iter (sub.value_binding sub) vbs;
       sub.expr sub e
-    | Pexp_fun (_lab, def, p, e) ->
+    | Pexp_fun (_lab, def, p, e, _) ->
       iter_opt (sub.expr sub) def;
       sub.pat sub p;
       sub.expr sub e

compiler/ml/ast_mapper.ml

@@ -278,8 +278,8 @@ module E = struct
     | Pexp_constant x -> constant ~loc ~attrs x
     | Pexp_let (r, vbs, e) ->
       let_ ~loc ~attrs r (List.map (sub.value_binding sub) vbs) (sub.expr sub e)
-    | Pexp_fun (lab, def, p, e) ->
-      fun_ ~loc ~attrs lab
+    | Pexp_fun (lab, def, p, e, arity) ->
+      fun_ ~loc ~attrs ~arity lab
         (map_opt (sub.expr sub) def)
         (sub.pat sub p) (sub.expr sub e)
     | Pexp_function pel -> function_ ~loc ~attrs (sub.cases sub pel)

compiler/ml/ast_mapper_from0.ml

@@ -285,7 +285,7 @@ module E = struct
     | Pexp_let (r, vbs, e) ->
       let_ ~loc ~attrs r (List.map (sub.value_binding sub) vbs) (sub.expr sub e)
     | Pexp_fun (lab, def, p, e) ->
-      fun_ ~loc ~attrs lab
+      fun_ ~loc ~attrs ~arity:None lab
         (map_opt (sub.expr sub) def)
         (sub.pat sub p) (sub.expr sub e)
     | Pexp_function pel -> function_ ~loc ~attrs (sub.cases sub pel)
@@ -295,8 +295,42 @@ module E = struct
       match_ ~loc ~attrs (sub.expr sub e) (sub.cases sub pel)
     | Pexp_try (e, pel) -> try_ ~loc ~attrs (sub.expr sub e) (sub.cases sub pel)
     | Pexp_tuple el -> tuple ~loc ~attrs (List.map (sub.expr sub) el)
-    | Pexp_construct (lid, arg) ->
-      construct ~loc ~attrs (map_loc sub lid) (map_opt (sub.expr sub) arg)
+    | Pexp_construct (lid, arg) -> (
+      let lid1 = map_loc sub lid in
+      let arg1 = map_opt (sub.expr sub) arg in
+      let exp1 = construct ~loc ~attrs lid1 arg1 in
+      match lid.txt with
+      | Lident "Function$" -> (
+        let rec attributes_to_arity (attrs : Parsetree.attributes) =
+          match attrs with
+          | ( {txt = "res.arity"},
+              PStr
+                [
+                  {
+                    pstr_desc =
+                      Pstr_eval
+                        ( {pexp_desc = Pexp_constant (Pconst_integer (arity, _))},
+                          _ );
+                  };
+                ] )
+            :: _ ->
+            int_of_string arity
+          | _ :: rest -> attributes_to_arity rest
+          | [] -> assert false
+        in
+        match arg1 with
+        | Some ({pexp_desc = Pexp_fun (l, eo, p, e, _)} as e1) ->
+          let arity = attributes_to_arity attrs in
+          {
+            e1 with
+            pexp_desc =
+              Pexp_construct
+                ( lid1,
+                  Some {e with pexp_desc = Pexp_fun (l, eo, p, e, Some arity)}
+                );
+          }
+        | _ -> exp1)
+      | _ -> exp1)
     | Pexp_variant (lab, eo) ->
       variant ~loc ~attrs lab (map_opt (sub.expr sub) eo)
     | Pexp_record (l, eo) ->

compiler/ml/ast_mapper_to0.ml

@@ -283,7 +283,7 @@ module E = struct
     | Pexp_constant x -> constant ~loc ~attrs (map_constant x)
     | Pexp_let (r, vbs, e) ->
       let_ ~loc ~attrs r (List.map (sub.value_binding sub) vbs) (sub.expr sub e)
-    | Pexp_fun (lab, def, p, e) ->
+    | Pexp_fun (lab, def, p, e, _) ->
       fun_ ~loc ~attrs lab
         (map_opt (sub.expr sub) def)
         (sub.pat sub p) (sub.expr sub e)
@@ -294,8 +294,31 @@ module E = struct
       match_ ~loc ~attrs (sub.expr sub e) (sub.cases sub pel)
     | Pexp_try (e, pel) -> try_ ~loc ~attrs (sub.expr sub e) (sub.cases sub pel)
     | Pexp_tuple el -> tuple ~loc ~attrs (List.map (sub.expr sub) el)
-    | Pexp_construct (lid, arg) ->
-      construct ~loc ~attrs (map_loc sub lid) (map_opt (sub.expr sub) arg)
+    | Pexp_construct (lid, arg) -> (
+      let exp0 =
+        construct ~loc ~attrs (map_loc sub lid) (map_opt (sub.expr sub) arg)
+      in
+      match lid.txt with
+      | Lident "Function$" -> (
+        match arg with
+        | Some {pexp_desc = Pexp_fun (_, _, _, _, Some arity)} ->
+          let arity_to_attributes arity =
+            [
+              ( Location.mknoloc "res.arity",
+                Parsetree0.PStr
+                  [
+                    Ast_helper0.Str.eval
+                      (Ast_helper0.Exp.constant
+                         (Pconst_integer (string_of_int arity, None)));
+                  ] );
+            ]
+          in
+          {
+            exp0 with
+            pexp_attributes = arity_to_attributes arity @ exp0.pexp_attributes;
+          }
+        | _ -> assert false)
+      | _ -> exp0)
     | Pexp_variant (lab, eo) ->
       variant ~loc ~attrs lab (map_opt (sub.expr sub) eo)
     | Pexp_record (l, eo) ->

compiler/ml/ast_uncurried.ml

@@ -19,36 +19,14 @@ let uncurried_type ~loc ~arity t_arg =
   let t_arity = arity_type ~loc arity in
   Ast_helper.Typ.constr ~loc {txt = Lident "function$"; loc} [t_arg; t_arity]
 
-let arity_to_attributes arity =
-  [
-    ( Location.mknoloc "res.arity",
-      Parsetree.PStr
-        [
-          Ast_helper.Str.eval
-            (Ast_helper.Exp.constant
-               (Pconst_integer (string_of_int arity, None)));
-        ] );
-  ]
-
-let rec attributes_to_arity (attrs : Parsetree.attributes) =
-  match attrs with
-  | ( {txt = "res.arity"},
-      PStr
-        [
-          {
-            pstr_desc =
-              Pstr_eval
-                ({pexp_desc = Pexp_constant (Pconst_integer (arity, _))}, _);
-          };
-        ] )
-    :: _ ->
-    int_of_string arity
-  | _ :: rest -> attributes_to_arity rest
-  | _ -> assert false
-
 let uncurried_fun ~loc ~arity fun_expr =
+  let fun_expr =
+    match fun_expr.Parsetree.pexp_desc with
+    | Pexp_fun (l, eo, p, e, _) ->
+      {fun_expr with pexp_desc = Pexp_fun (l, eo, p, e, Some arity)}
+    | _ -> assert false
+  in
   Ast_helper.Exp.construct ~loc
-    ~attrs:(arity_to_attributes arity)
     (Location.mknoloc (Longident.Lident "Function$"))
     (Some fun_expr)
 
@@ -59,7 +37,13 @@ let expr_is_uncurried_fun (expr : Parsetree.expression) =
 
 let expr_extract_uncurried_fun (expr : Parsetree.expression) =
   match expr.pexp_desc with
-  | Pexp_construct ({txt = Lident "Function$"}, Some e) -> e
+  | Pexp_construct ({txt = Lident "Function$"}, Some e) ->
+    let () =
+      match e.pexp_desc with
+      | Pexp_fun (_, _, _, _, Some _arity) -> ()
+      | _ -> assert false
+    in
+    e
   | _ -> assert false
 
 let core_type_is_uncurried_fun (typ : Parsetree.core_type) =

compiler/ml/depend.ml

@@ -218,7 +218,7 @@ let rec add_expr bv exp =
   | Pexp_let (rf, pel, e) ->
     let bv = add_bindings rf bv pel in
     add_expr bv e
-  | Pexp_fun (_, opte, p, e) ->
+  | Pexp_fun (_, opte, p, e, _) ->
     add_opt add_expr bv opte;
     add_expr (add_pattern bv p) e
   | Pexp_function pel -> add_cases bv pel