AST: store arity in function type

CHANGELOG.md

@@ -30,7 +30,7 @@
 - Remove unused code from Location and Rescript_cpp modules. https://github.com/rescript-lang/rescript/pull/7150
 - Build with OCaml 5.2.1. https://github.com/rescript-lang/rescript-compiler/pull/7201
 - AST cleanup: Remove `Function$` entirely for function definitions. https://github.com/rescript-lang/rescript/pull/7200
-
+- AST cleanup: store arity in function type https://github.com/rescript-lang/rescript/pull/7195
 
 # 12.0.0-alpha.5
 

analysis/reanalyze/src/DeadOptionalArgs.ml

@@ -31,17 +31,17 @@ let addFunctionReference ~(locFrom : Location.t) ~(locTo : Location.t) =
 let rec hasOptionalArgs (texpr : Types.type_expr) =
   match texpr.desc with
   | _ when not (active ()) -> false
-  | Tarrow (Optional _, _tFrom, _tTo, _) -> true
-  | Tarrow (_, _tFrom, tTo, _) -> hasOptionalArgs tTo
+  | Tarrow (Optional _, _tFrom, _tTo, _, _) -> true
+  | Tarrow (_, _tFrom, tTo, _, _) -> hasOptionalArgs tTo
   | Tlink t -> hasOptionalArgs t
   | Tsubst t -> hasOptionalArgs t
   | _ -> false
 
 let rec fromTypeExpr (texpr : Types.type_expr) =
   match texpr.desc with
   | _ when not (active ()) -> []
-  | Tarrow (Optional s, _tFrom, tTo, _) -> s :: fromTypeExpr tTo
-  | Tarrow (_, _tFrom, tTo, _) -> fromTypeExpr tTo
+  | Tarrow (Optional s, _tFrom, tTo, _, _) -> s :: fromTypeExpr tTo
+  | Tarrow (_, _tFrom, tTo, _, _) -> fromTypeExpr tTo
   | Tlink t -> fromTypeExpr t
   | Tsubst t -> fromTypeExpr t
   | _ -> []

analysis/src/CompletionBackEnd.ml

@@ -898,7 +898,7 @@ and getCompletionsForContextPath ~debug ~full ~opens ~rawOpens ~pos ~env ~exact
         | [] -> tRet
         | (label, tArg) :: rest ->
           let restType = reconstructFunctionType rest tRet in
-          {typ with desc = Tarrow (label, tArg, restType, Cok)}
+          {typ with desc = Tarrow (label, tArg, restType, Cok, None)}
       in
       let rec processApply args labels =
         match (args, labels) with
@@ -1362,7 +1362,7 @@ let rec completeTypedValue ?(typeArgContext : typeArgContext option) ~rawOpens
       | Tlink t1
       | Tsubst t1
       | Tpoly (t1, [])
-      | Tconstr (Pident {name = "function$"}, [t1; _], _) ->
+      | Tconstr (Pident {name = "function$"}, [t1], _) ->
         fnReturnsTypeT t1
       | Tarrow _ -> (
         match TypeUtils.extractFunctionType ~env ~package:full.package t with

analysis/src/CompletionJsx.ml

@@ -238,7 +238,7 @@ let getJsxLabels ~componentPath ~findTypeOfValue ~package =
       | Tlink t1
       | Tsubst t1
       | Tpoly (t1, [])
-      | Tconstr (Pident {name = "function$"}, [t1; _], _) ->
+      | Tconstr (Pident {name = "function$"}, [t1], _) ->
         getLabels t1
       | Tconstr (p, [propsType], _) when Path.name p = "React.component" -> (
         let rec getPropsType (t : Types.type_expr) =
@@ -251,15 +251,15 @@ let getJsxLabels ~componentPath ~findTypeOfValue ~package =
         match propsType |> getPropsType with
         | Some (path, typeArgs) -> getFields ~path ~typeArgs
         | None -> [])
-      | Tarrow (Nolabel, {desc = Tconstr (path, typeArgs, _)}, _, _)
+      | Tarrow (Nolabel, {desc = Tconstr (path, typeArgs, _)}, _, _, _)
         when Path.last path = "props" ->
         getFields ~path ~typeArgs
       | Tconstr (clPath, [{desc = Tconstr (path, typeArgs, _)}; _], _)
         when Path.name clPath = "React.componentLike"
              && Path.last path = "props" ->
         (* JSX V4 external or interface *)
         getFields ~path ~typeArgs
-      | Tarrow (Nolabel, typ, _, _) -> (
+      | Tarrow (Nolabel, typ, _, _, _) -> (
         (* Component without the JSX PPX, like a make fn taking a hand-written
            type props. *)
         let rec digToConstr typ =

analysis/src/CreateInterface.ml

@@ -123,8 +123,9 @@ let printSignature ~extractor ~signature =
       Ctype.newconstr (Pdot (Pident (Ident.create "React"), "element", 0)) []
     in
     match typ.desc with
-    | Tconstr (Pident {name = "function$"}, [typ; _], _) -> getComponentType typ
-    | Tarrow (_, {desc = Tconstr (Path.Pident propsId, typeArgs, _)}, retType, _)
+    | Tconstr (Pident {name = "function$"}, [typ], _) -> getComponentType typ
+    | Tarrow
+        (_, {desc = Tconstr (Path.Pident propsId, typeArgs, _)}, retType, _, _)
       when Ident.name propsId = "props" ->
       Some (typeArgs, retType)
     | Tconstr
@@ -173,14 +174,17 @@ let printSignature ~extractor ~signature =
               if labelDecl.ld_optional then Asttypes.Optional lblName
               else Labelled lblName
             in
-            {retType with desc = Tarrow (lbl, propType, mkFunType rest, Cok)}
+            {
+              retType with
+              desc = Tarrow (lbl, propType, mkFunType rest, Cok, None);
+            }
         in
         let funType =
           if List.length labelDecls = 0 (* No props *) then
             let tUnit =
               Ctype.newconstr (Path.Pident (Ident.create "unit")) []
             in
-            {retType with desc = Tarrow (Nolabel, tUnit, retType, Cok)}
+            {retType with desc = Tarrow (Nolabel, tUnit, retType, Cok, None)}
           else mkFunType labelDecls
         in
         sigItemToString

analysis/src/Shared.ml

@@ -52,7 +52,7 @@ let findTypeConstructors (tel : Types.type_expr list) =
     | Tconstr (path, args, _) ->
       addPath path;
       args |> List.iter loop
-    | Tarrow (_, te1, te2, _) ->
+    | Tarrow (_, te1, te2, _, _) ->
       loop te1;
       loop te2
     | Ttuple tel -> tel |> List.iter loop

analysis/src/SignatureHelp.ml

@@ -118,7 +118,7 @@ let extractParameters ~signature ~typeStrForParser ~labelPrefixLen =
                  ptyp_desc =
                    Ptyp_constr
                      ( {txt = Lident "function$"},
-                       [({ptyp_desc = Ptyp_arrow _} as expr); _] );
+                       [({ptyp_desc = Ptyp_arrow _} as expr)] );
                };
            };
      } );
@@ -128,7 +128,7 @@ let extractParameters ~signature ~typeStrForParser ~labelPrefixLen =
       | {
        (* Gotcha: functions with multiple arugments are modelled as a series of single argument functions. *)
        Parsetree.ptyp_desc =
-         Ptyp_arrow (argumentLabel, argumentTypeExpr, nextFunctionExpr);
+         Ptyp_arrow (argumentLabel, argumentTypeExpr, nextFunctionExpr, _);
        ptyp_loc;
       } ->
         let startOffset =

analysis/src/TypeUtils.ml

@@ -10,7 +10,7 @@ let debugLogTypeArgContext {env; typeArgs; typeParams} =
 let rec hasTvar (ty : Types.type_expr) : bool =
   match ty.desc with
   | Tvar _ -> true
-  | Tarrow (_, ty1, ty2, _) -> hasTvar ty1 || hasTvar ty2
+  | Tarrow (_, ty1, ty2, _, _) -> hasTvar ty1 || hasTvar ty2
   | Ttuple tyl -> List.exists hasTvar tyl
   | Tconstr (_, tyl, _) -> List.exists hasTvar tyl
   | Tobject (ty, _) -> hasTvar ty
@@ -36,7 +36,7 @@ let findTypeViaLoc ~full ~debug (loc : Location.t) =
 
 let rec pathFromTypeExpr (t : Types.type_expr) =
   match t.desc with
-  | Tconstr (Pident {name = "function$"}, [t; _], _) -> pathFromTypeExpr t
+  | Tconstr (Pident {name = "function$"}, [t], _) -> pathFromTypeExpr t
   | Tconstr (path, _typeArgs, _)
   | Tlink {desc = Tconstr (path, _typeArgs, _)}
   | Tsubst {desc = Tconstr (path, _typeArgs, _)}
@@ -116,8 +116,8 @@ let instantiateType ~typeParams ~typeArgs (t : Types.type_expr) =
       | Tsubst t -> loop t
       | Tvariant rd -> {t with desc = Tvariant (rowDesc rd)}
       | Tnil -> t
-      | Tarrow (lbl, t1, t2, c) ->
-        {t with desc = Tarrow (lbl, loop t1, loop t2, c)}
+      | Tarrow (lbl, t1, t2, c, arity) ->
+        {t with desc = Tarrow (lbl, loop t1, loop t2, c, arity)}
       | Ttuple tl -> {t with desc = Ttuple (tl |> List.map loop)}
       | Tobject (t, r) -> {t with desc = Tobject (loop t, r)}
       | Tfield (n, k, t1, t2) -> {t with desc = Tfield (n, k, loop t1, loop t2)}
@@ -169,8 +169,8 @@ let instantiateType2 ?(typeArgContext : typeArgContext option)
       | Tsubst t -> loop t
       | Tvariant rd -> {t with desc = Tvariant (rowDesc rd)}
       | Tnil -> t
-      | Tarrow (lbl, t1, t2, c) ->
-        {t with desc = Tarrow (lbl, loop t1, loop t2, c)}
+      | Tarrow (lbl, t1, t2, c, arity) ->
+        {t with desc = Tarrow (lbl, loop t1, loop t2, c, arity)}
       | Ttuple tl -> {t with desc = Ttuple (tl |> List.map loop)}
       | Tobject (t, r) -> {t with desc = Tobject (loop t, r)}
       | Tfield (n, k, t1, t2) -> {t with desc = Tfield (n, k, loop t1, loop t2)}
@@ -242,8 +242,8 @@ let rec extractFunctionType ~env ~package typ =
   let rec loop ~env acc (t : Types.type_expr) =
     match t.desc with
     | Tlink t1 | Tsubst t1 | Tpoly (t1, []) -> loop ~env acc t1
-    | Tarrow (label, tArg, tRet, _) -> loop ~env ((label, tArg) :: acc) tRet
-    | Tconstr (Pident {name = "function$"}, [t; _], _) ->
+    | Tarrow (label, tArg, tRet, _, _) -> loop ~env ((label, tArg) :: acc) tRet
+    | Tconstr (Pident {name = "function$"}, [t], _) ->
       extractFunctionType ~env ~package t
     | Tconstr (path, typeArgs, _) -> (
       match References.digConstructor ~env ~package path with
@@ -281,9 +281,9 @@ let rec extractFunctionType2 ?typeArgContext ~env ~package typ =
   let rec loop ?typeArgContext ~env acc (t : Types.type_expr) =
     match t.desc with
     | Tlink t1 | Tsubst t1 | Tpoly (t1, []) -> loop ?typeArgContext ~env acc t1
-    | Tarrow (label, tArg, tRet, _) ->
+    | Tarrow (label, tArg, tRet, _, _) ->
       loop ?typeArgContext ~env ((label, tArg) :: acc) tRet
-    | Tconstr (Pident {name = "function$"}, [t; _], _) ->
+    | Tconstr (Pident {name = "function$"}, [t], _) ->
       extractFunctionType2 ?typeArgContext ~env ~package t
     | Tconstr (path, typeArgs, _) -> (
       match References.digConstructor ~env ~package path with
@@ -334,7 +334,7 @@ let rec extractType ?(printOpeningDebug = true)
     Some (Tstring env, typeArgContext)
   | Tconstr (Path.Pident {name = "exn"}, [], _) ->
     Some (Texn env, typeArgContext)
-  | Tconstr (Pident {name = "function$"}, [t; _], _) -> (
+  | Tconstr (Pident {name = "function$"}, [t], _) -> (
     match extractFunctionType2 ?typeArgContext t ~env ~package with
     | args, tRet, typeArgContext when args <> [] ->
       Some
@@ -910,14 +910,14 @@ let getArgs ~env (t : Types.type_expr) ~full =
     | Tlink t1
     | Tsubst t1
     | Tpoly (t1, [])
-    | Tconstr (Pident {name = "function$"}, [t1; _], _) ->
+    | Tconstr (Pident {name = "function$"}, [t1], _) ->
       getArgsLoop ~full ~env ~currentArgumentPosition t1
-    | Tarrow (Labelled l, tArg, tRet, _) ->
+    | Tarrow (Labelled l, tArg, tRet, _, _) ->
       (SharedTypes.Completable.Labelled l, tArg)
       :: getArgsLoop ~full ~env ~currentArgumentPosition tRet
-    | Tarrow (Optional l, tArg, tRet, _) ->
+    | Tarrow (Optional l, tArg, tRet, _, _) ->
       (Optional l, tArg) :: getArgsLoop ~full ~env ~currentArgumentPosition tRet
-    | Tarrow (Nolabel, tArg, tRet, _) ->
+    | Tarrow (Nolabel, tArg, tRet, _, _) ->
       (Unlabelled {argumentPosition = currentArgumentPosition}, tArg)
       :: getArgsLoop ~full ~env
            ~currentArgumentPosition:(currentArgumentPosition + 1)

compiler/frontend/ast_comb.ml

@@ -40,7 +40,9 @@ let tuple_type_pair ?loc kind arity =
     match kind with
     | `Run -> (ty, [], ty)
     | `Make ->
-      (Ast_compatible.arrow ?loc (Ast_literal.type_unit ?loc ()) ty, [], ty)
+      ( Ast_compatible.arrow ?loc ~arity:None (Ast_literal.type_unit ?loc ()) ty,
+        [],
+        ty )
   else
     let number = arity + 1 in
     let tys =
@@ -50,7 +52,7 @@ let tuple_type_pair ?loc kind arity =
     match tys with
     | result :: rest ->
       ( Ext_list.reduce_from_left tys (fun r arg ->
-            Ast_compatible.arrow ?loc arg r),
+            Ast_compatible.arrow ?loc ~arity:None arg r),
         List.rev rest,
         result )
     | [] -> assert false

compiler/frontend/ast_compatible.ml

@@ -30,7 +30,8 @@ open Parsetree
 
 let default_loc = Location.none
 
-let arrow ?loc ?attrs a b = Ast_helper.Typ.arrow ?loc ?attrs Nolabel a b
+let arrow ?loc ?attrs ~arity a b =
+  Ast_helper.Typ.arrow ?loc ?attrs ~arity Nolabel a b
 
 let apply_simple ?(loc = default_loc) ?(attrs = []) (fn : expression)
     (args : expression list) : expression =
@@ -94,16 +95,16 @@ let apply_labels ?(loc = default_loc) ?(attrs = []) fn
       Pexp_apply (fn, Ext_list.map args (fun (l, a) -> (Asttypes.Labelled l, a)));
   }
 
-let label_arrow ?(loc = default_loc) ?(attrs = []) s a b : core_type =
+let label_arrow ?(loc = default_loc) ?(attrs = []) ~arity s a b : core_type =
   {
-    ptyp_desc = Ptyp_arrow (Asttypes.Labelled s, a, b);
+    ptyp_desc = Ptyp_arrow (Asttypes.Labelled s, a, b, arity);
     ptyp_loc = loc;
     ptyp_attributes = attrs;
   }
 
-let opt_arrow ?(loc = default_loc) ?(attrs = []) s a b : core_type =
+let opt_arrow ?(loc = default_loc) ?(attrs = []) ~arity s a b : core_type =
   {
-    ptyp_desc = Ptyp_arrow (Asttypes.Optional s, a, b);
+    ptyp_desc = Ptyp_arrow (Asttypes.Optional s, a, b, arity);
     ptyp_loc = loc;
     ptyp_attributes = attrs;
   }

compiler/frontend/ast_compatible.mli

@@ -90,11 +90,17 @@ val fun_ :
    expression *)
 
 val arrow :
-  ?loc:Location.t -> ?attrs:attrs -> core_type -> core_type -> core_type
+  ?loc:Location.t ->
+  ?attrs:attrs ->
+  arity:Asttypes.arity ->
+  core_type ->
+  core_type ->
+  core_type
 
 val label_arrow :
   ?loc:Location.t ->
   ?attrs:attrs ->
+  arity:Asttypes.arity ->
   string ->
   core_type ->
   core_type ->
@@ -103,6 +109,7 @@ val label_arrow :
 val opt_arrow :
   ?loc:Location.t ->
   ?attrs:attrs ->
+  arity:Asttypes.arity ->
   string ->
   core_type ->
   core_type ->

compiler/frontend/ast_core_type.ml

@@ -95,7 +95,8 @@ let from_labels ~loc arity labels : t =
   in
   Ext_list.fold_right2 labels tyvars result_type
     (fun label (* {loc ; txt = label }*) tyvar acc ->
-      Ast_compatible.label_arrow ~loc:label.loc label.txt tyvar acc)
+      Ast_compatible.label_arrow ~loc:label.loc ~arity:(Some arity) label.txt
+        tyvar acc)
 
 let make_obj ~loc xs = Typ.object_ ~loc xs Closed
 
@@ -108,7 +109,7 @@ let make_obj ~loc xs = Typ.object_ ~loc xs Closed
 *)
 let rec get_uncurry_arity_aux (ty : t) acc =
   match ty.ptyp_desc with
-  | Ptyp_arrow (_, _, new_ty) -> get_uncurry_arity_aux new_ty (succ acc)
+  | Ptyp_arrow (_, _, new_ty, _) -> get_uncurry_arity_aux new_ty (succ acc)
   | Ptyp_poly (_, ty) -> get_uncurry_arity_aux ty acc
   | _ -> acc
 
@@ -119,7 +120,7 @@ let rec get_uncurry_arity_aux (ty : t) acc =
 *)
 let get_uncurry_arity (ty : t) =
   match ty.ptyp_desc with
-  | Ptyp_arrow (_, _, rest) -> Some (get_uncurry_arity_aux rest 1)
+  | Ptyp_arrow (_, _, rest, _) -> Some (get_uncurry_arity_aux rest 1)
   | _ -> None
 
 let get_curry_arity (ty : t) =
@@ -139,15 +140,15 @@ type param_type = {
 let mk_fn_type (new_arg_types_ty : param_type list) (result : t) : t =
   Ext_list.fold_right new_arg_types_ty result (fun {label; ty; attr; loc} acc ->
       {
-        ptyp_desc = Ptyp_arrow (label, ty, acc);
+        ptyp_desc = Ptyp_arrow (label, ty, acc, None);
         ptyp_loc = loc;
         ptyp_attributes = attr;
       })
 
 let list_of_arrow (ty : t) : t * param_type list =
   let rec aux (ty : t) acc =
     match ty.ptyp_desc with
-    | Ptyp_arrow (label, t1, t2) ->
+    | Ptyp_arrow (label, t1, t2, _) ->
       aux t2
         (({label; ty = t1; attr = ty.ptyp_attributes; loc = ty.ptyp_loc}
            : param_type)

compiler/frontend/ast_core_type_class_type.ml

@@ -69,11 +69,11 @@ let typ_mapper (self : Bs_ast_mapper.mapper) (ty : Parsetree.core_type) =
   | {
    ptyp_attributes;
    ptyp_desc =
-     ( Ptyp_arrow (label, args, body)
+     ( Ptyp_arrow (label, args, body, _)
      | Ptyp_constr
          (* function$<...> is re-wrapped around only in case Nothing below *)
          ( {txt = Lident "function$"},
-           [{ptyp_desc = Ptyp_arrow (label, args, body)}; _] ) );
+           [{ptyp_desc = Ptyp_arrow (label, args, body, _)}] ) );
    (* let it go without regard label names,
       it will report error later when the label is not empty
    *)