Make a function's locally abstract types part of the function node

CHANGELOG.md

@@ -29,6 +29,7 @@
 
 - Fix argument evaluation order when a function call is inlined: the beta reducer stacked argument bindings in reverse parameter order, so the last argument was evaluated first when arguments could not be substituted directly. https://github.com/rescript-lang/rescript/pull/8572
 - Preserve parentheses around multiplication, division, and modulo expressions used as exponents. https://github.com/rescript-lang/rescript/pull/8550
+- Make a function's locally abstract types (`(type t, x) => ...`) part of the function AST node instead of a chain of wrapper nodes. Fixes the formatter dropping the association of attributes with their `type` group (`(@attr type t, x, @attr2 type s, y)` used to print as `@attr @attr2` on the function) and comments written next to a type parameter migrating onto the following value parameter. https://github.com/rescript-lang/rescript/pull/8574
 - Enforce function arity in interface/module inclusion and type coercion. Previously a curried implementation (e.g. `int => int => int`) could satisfy an uncurried interface (`(int, int) => int`) or be coerced to it, which could miscompile calls made through the interface type. Such mismatches are now compile errors with an explanatory hint. https://github.com/rescript-lang/rescript/pull/8559
 - Fix termination-analysis false positives for functions whose progress flows through un-annotated helpers: collecting the callees of a function binding was accidentally disabled in 2024 (the collection guard required a node shape that uncurried code never produces), so helpers calling `@progress` functions were no longer added to the function table. https://github.com/rescript-lang/rescript/pull/8568
 - Fix default values of optional parameters being computed at the wrong time for curried functions: in `(~x=default, y) => (~z=default, w) => ...`, `x`'s default was only computed when the *inner* function was applied. Each default is now computed when its own parameter group is applied. https://github.com/rescript-lang/rescript/pull/8568

compiler/frontend/bs_ast_mapper.ml

@@ -325,8 +325,12 @@ module E = struct
            sub vbs)
         (sub.expr sub e)
     (* #end *)
-    | Pexp_fun {params; body; async} ->
+    | Pexp_fun {newtypes; params; body; async} ->
       fun_ ~loc ~attrs ~async
+        ~newtypes:
+          (List.map
+             (fun (name, attrs) -> (map_loc sub name, sub.attributes sub attrs))
+             newtypes)
         (List.map
            (fun (param : Parsetree.fun_param) ->
              {

compiler/frontend/bs_builtin_ppx.ml

@@ -95,7 +95,7 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
   | Pexp_newtype (s, body) ->
     let res = self.expr self body in
     {e with pexp_desc = Pexp_newtype (s, res)}
-  | Pexp_fun {params; body; async} -> (
+  | Pexp_fun {newtypes; params; body; async} -> (
     match Ast_attributes.process_attributes_rev e.pexp_attributes with
     | Nothing, _ ->
       (* Handle @async x => y => ... is in async context *)
@@ -116,19 +116,29 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
          mapper did (GH #7974). *)
       let body = self.expr self body in
       in_function_def := saved_in_function_def;
+      let newtypes =
+        Ext_list.map newtypes (fun (name, nt_attrs) ->
+            (name, self.attributes self nt_attrs))
+      in
       let mapped =
-        Ast_helper.Exp.fun_ ~loc:e.pexp_loc ~attrs ~async params body
+        Ast_helper.Exp.fun_ ~loc:e.pexp_loc ~attrs ~async ~newtypes params body
       in
       Ast_async.make_function_async ~async mapped
     | Meth_callback _, pexp_attributes ->
       (* FIXME: does it make sense to have a label for [this] ? *)
       async_context := false;
-      {
-        e with
-        pexp_desc =
-          Ast_uncurry_gen.to_method_callback ~async e.pexp_loc self params body;
-        pexp_attributes;
-      })
+      let callback =
+        {
+          e with
+          pexp_desc =
+            Ast_uncurry_gen.to_method_callback ~async e.pexp_loc self params
+              body;
+          pexp_attributes;
+        }
+      in
+      (* Keep the locally abstract types in scope around the callback. *)
+      Ext_list.fold_right newtypes callback (fun (name, nt_attrs) acc ->
+          Ast_helper.Exp.newtype ~loc:e.pexp_loc ~attrs:nt_attrs name acc))
   | Pexp_apply _ -> Ast_exp_apply.app_exp_mapper e self
   | Pexp_match
       ( b,

compiler/ml/ast_helper.ml

@@ -158,9 +158,9 @@ 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 ?(async = false) params body =
+  let fun_ ?loc ?attrs ?(async = false) ?(newtypes = []) params body =
     assert (params <> []);
-    mk ?loc ?attrs (Pexp_fun {params; body; async})
+    mk ?loc ?attrs (Pexp_fun {newtypes; params; body; async})
 
   let fun_param ?(attrs = []) ?default lbl pat =
     {p_attrs = attrs; p_lbl = lbl; p_default = default; p_pat = pat}

compiler/ml/ast_helper.mli

@@ -134,6 +134,7 @@ module Exp : sig
     ?loc:loc ->
     ?attrs:attrs ->
     ?async:bool ->
+    ?newtypes:(str * attrs) list ->
     fun_param list ->
     expression ->
     expression

compiler/ml/ast_iterator.ml

@@ -289,7 +289,12 @@ module E = struct
     | Pexp_let (_r, vbs, e) ->
       List.iter (sub.value_binding sub) vbs;
       sub.expr sub e
-    | Pexp_fun {params; body} ->
+    | Pexp_fun {newtypes; params; body} ->
+      List.iter
+        (fun (name, attrs) ->
+          iter_loc sub name;
+          sub.attributes sub attrs)
+        newtypes;
       List.iter
         (fun {p_default; p_pat} ->
           iter_opt (sub.expr sub) p_default;

compiler/ml/ast_mapper.ml

@@ -288,8 +288,12 @@ 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 {params; body; async} ->
+    | Pexp_fun {newtypes; params; body; async} ->
       fun_ ~loc ~attrs ~async
+        ~newtypes:
+          (List.map
+             (fun (name, attrs) -> (map_loc sub name, sub.attributes sub attrs))
+             newtypes)
         (List.map
            (fun param ->
              {

compiler/ml/ast_mapper_from0.ml

@@ -706,7 +706,8 @@ module E = struct
             in
             {
               e1 with
-              pexp_desc = Pexp_fun {params; body; async = f.async};
+              pexp_desc =
+                Pexp_fun {newtypes = []; params; body; async = f.async};
               pexp_attributes = e1.pexp_attributes @ node_attrs;
             })
         | _ -> exp1)
@@ -779,8 +780,55 @@ module E = struct
     | Pexp_lazy _ -> failwith "Pexp_lazy is no longer present in ReScript"
     | Pexp_poly _ -> failwith "Pexp_poly is no longer present in ReScript"
     | Pexp_object () -> assert false
-    | Pexp_newtype (s, e) ->
-      newtype ~loc ~attrs (map_loc sub s) (sub.expr sub e)
+    | Pexp_newtype (s, e) -> (
+      (* Fuse a chain of newtype wrappers over a Function$ node into the
+         function's [newtypes] field. Each wrapper's attributes are its
+         newtype's attributes, except on this outermost wrapper:
+         attributes before the internal [_res.newtype_attrs] marker (or
+         all of them, when there is no marker) are function-node
+         attributes, the ones after the marker belong to the first
+         newtype. Chains over anything else (e.g. the
+         [let f: type t. ...] sugar) keep their [Pexp_newtype] nodes. *)
+      let node_attrs, first_nt_attrs =
+        let rec split acc = function
+          | ({txt = "_res.newtype_attrs"}, _) :: rest -> (List.rev acc, rest)
+          | a :: rest -> split (a :: acc) rest
+          | [] -> (List.rev acc, [])
+        in
+        split [] attrs
+      in
+      let rec gather acc (e0 : Parsetree0.expression) =
+        match e0.pexp_desc with
+        | Pexp_newtype (s1, body) ->
+          gather
+            ((map_loc sub s1, sub.attributes sub e0.pexp_attributes) :: acc)
+            body
+        | Pexp_construct ({txt = Longident.Lident "Function$"}, Some _) ->
+          Some (List.rev acc, e0)
+        | _ -> None
+      in
+      match gather [(map_loc sub s, first_nt_attrs)] e with
+      | Some (newtypes, base) -> (
+        let base1 = sub.expr sub base in
+        match base1.pexp_desc with
+        | Pexp_fun ({newtypes = []} as f) ->
+          {
+            Pt.pexp_desc = Pexp_fun {f with newtypes};
+            pexp_attributes = base1.pexp_attributes @ node_attrs;
+            pexp_loc = loc;
+          }
+        | _ -> (
+          (* PPX-mangled Function$: keep the wrapper chain as-is. *)
+          match newtypes with
+          | [] -> assert false
+          | (n0, _) :: rest ->
+            let inner =
+              List.fold_right
+                (fun (n, a) acc -> newtype ~loc ~attrs:a n acc)
+                rest base1
+            in
+            newtype ~loc ~attrs n0 inner))
+      | None -> newtype ~loc ~attrs (map_loc sub s) (sub.expr sub e))
     | Pexp_pack me -> pack ~loc ~attrs (sub.module_expr sub me)
     | Pexp_open (ovf, lid, e) ->
       open_ ~loc ~attrs ovf (map_loc sub lid) (sub.expr sub e)

compiler/ml/ast_mapper_to0.ml

@@ -405,11 +405,12 @@ 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 {params; body; async} ->
+    | Pexp_fun {newtypes; params; body; async} -> (
       (* Re-curry the n-ary function into the v0 chain of unary funs, and
          wrap it in Function$ carrying the arity as a res.arity attribute.
-         The head carries the function node's own attributes (and the
-         res.async marker), matching what the old parser produced.
+         Without newtypes the head carries the function node's own
+         attributes (and the res.async marker), matching what the old parser
+         produced; with newtypes they travel on the newtype wrapper instead.
 
          v0 fun nodes have a single attribute slot for what the current
          parsetree splits into node attributes and parameter attributes.
@@ -433,7 +434,10 @@ module E = struct
                 :: param_attrs
             in
             if is_head then
-              let base = attrs @ marked_param_attrs in
+              let base =
+                if newtypes = [] then attrs @ marked_param_attrs
+                else marked_param_attrs
+              in
               if async then
                 ({txt = "res.async"; loc = Location.none}, Pt.PStr []) :: base
               else base
@@ -457,9 +461,41 @@ module E = struct
                    (Pconst_integer (string_of_int arity, None)));
             ] )
       in
-      Ast_helper0.Exp.construct ~attrs:[arity_attr]
-        (Location.mkloc (Longident.Lident "Function$") e.pexp_loc)
-        (Some e)
+      let fn =
+        Ast_helper0.Exp.construct ~attrs:[arity_attr]
+          (Location.mkloc (Longident.Lident "Function$") e.pexp_loc)
+          (Some e)
+      in
+      (* Expand the newtypes back into the v0 wrapper chain around the
+         Function$ node. Each wrapper carries its own newtype's attributes.
+         The outermost wrapper is the whole expression in v0, so it also
+         carries the function node's attributes: when the first newtype has
+         attributes of its own, an internal [_res.newtype_attrs] marker
+         separates node attributes (before) from the first newtype's
+         attributes (after); without a marker every attribute on the
+         outermost wrapper is a function-node attribute, which is also how
+         wrapper attributes behaved before newtypes became a field. *)
+      match newtypes with
+      | [] -> fn
+      | (first_name, first_attrs) :: rest_newtypes ->
+        let inner =
+          List.fold_right
+            (fun (name, nt_attrs) acc ->
+              Ast_helper0.Exp.newtype ~loc
+                ~attrs:(sub.attributes sub nt_attrs)
+                (map_loc sub name) acc)
+            rest_newtypes fn
+        in
+        let first_attrs = sub.attributes sub first_attrs in
+        let outer_attrs =
+          if first_attrs = [] then attrs
+          else
+            attrs
+            @ ({txt = "_res.newtype_attrs"; loc = Location.none}, Pt.PStr [])
+              :: first_attrs
+        in
+        Ast_helper0.Exp.newtype ~loc ~attrs:outer_attrs (map_loc sub first_name)
+          inner)
     | Pexp_apply {funct = e; args; partial} ->
       let e =
         match (e.pexp_desc, args) with

compiler/ml/parsetree.ml

@@ -232,10 +232,18 @@ and expression_desc =
     (* let P1 = E1 and ... and Pn = EN in E       (flag = Nonrecursive)
        let rec P1 = E1 and ... and Pn = EN in E   (flag = Recursive)
     *)
-  | Pexp_fun of {params: fun_param list; body: expression; async: bool}
-    (* (P1, ~l:P2, ?l:P3=E0) => E           n-ary uncurried function.
+  | Pexp_fun of {
+      newtypes: (string loc * attributes) list;
+      params: fun_param list;
+      body: expression;
+      async: bool;
+    }
+    (* (type t, P1, ~l:P2, ?l:P3=E0) => E   n-ary uncurried function.
        The function's arity is [List.length params]; a function returning
        another function is a nested [Pexp_fun] in [body].
+       [newtypes] are the function's locally abstract types, each with its
+       own attributes; the parser hoists them in front of the value
+       parameters.
 
        Notes:
        - A default expression is only allowed on Optional parameters.

compiler/ml/pprintast.ml

@@ -627,17 +627,23 @@ and expression ctxt f x =
     | (Pexp_let _ | Pexp_letmodule _ | Pexp_open _ | Pexp_letexception _)
       when ctxt.semi ->
       paren true (expression reset_ctxt) f x
-    | Pexp_fun {params; body; async} ->
+    | Pexp_fun {newtypes; params; body; async} ->
       let arity_str = "[arity:" ^ string_of_int (List.length params) ^ "]" in
       let async_str = if async then "async " else "" in
+      let rec pp_newtypes f = function
+        | [] -> ()
+        | ((name : string Location.loc), nt_attrs) :: rest ->
+          pp f "%a(type %s)@;" (attributes ctxt) nt_attrs name.txt;
+          pp_newtypes f rest
+      in
       let rec pp_params f = function
         | [] -> ()
         | {p_lbl; p_default; p_pat} :: rest ->
           pp f "%a" (label_exp ctxt) (p_lbl, p_default, p_pat);
           pp_params f rest
       in
-      pp f "@[<2>%sfun@;%s%a->@;%a@]" async_str arity_str pp_params params
-        (expression ctxt) body
+      pp f "@[<2>%sfun@;%a%s%a->@;%a@]" async_str pp_newtypes newtypes arity_str
+        pp_params params (expression ctxt) body
     | Pexp_match (e, l) ->
       pp f "@[<hv0>@[<hv0>@[<2>match %a@]@ with@]%a@]" (expression reset_ctxt) e
         (case_list ctxt) l
@@ -1062,9 +1068,15 @@ and binding ctxt f {pvb_pat = p; pvb_expr = x; _} =
     if x.pexp_attributes <> [] then pp f "=@;%a" (expression ctxt) x
     else
       match x.pexp_desc with
-      | Pexp_fun {params; body; async} ->
+      | Pexp_fun {newtypes; params; body; async} ->
         let arity_str = "[arity:" ^ string_of_int (List.length params) ^ "]" in
         let async_str = if async then "async " else "" in
+        let rec pp_newtypes f = function
+          | [] -> ()
+          | ((name : string Location.loc), nt_attrs) :: rest ->
+            pp f "%a(type@ %s)@ " (attributes ctxt) nt_attrs name.txt;
+            pp_newtypes f rest
+        in
         let pp_param f {p_lbl; p_default; p_pat} =
           if p_lbl = Nolabel then simple_pattern ctxt f p_pat
           else label_exp ctxt f (p_lbl, p_default, p_pat)
@@ -1075,8 +1087,8 @@ and binding ctxt f {pvb_pat = p; pvb_expr = x; _} =
             pp f "%a@ " pp_param param;
             pp_params f rest
         in
-        pp f "%s%s%a%a" async_str arity_str pp_params params
-          pp_print_pexp_function body
+        pp f "%s%a%s%a%a" async_str pp_newtypes newtypes arity_str pp_params
+          params pp_print_pexp_function body
       | Pexp_newtype (str, e) ->
         pp f "(type@ %s)@ %a" str.txt pp_print_pexp_function e
       | _ -> pp f "=@;%a" (expression ctxt) x

compiler/ml/printast.ml

@@ -249,10 +249,15 @@ and expression i ppf x =
     line i ppf "Pexp_let %a\n" fmt_rec_flag rf;
     list i value_binding ppf l;
     expression i ppf e
-  | Pexp_fun {params; body; async} ->
+  | Pexp_fun {newtypes; params; body; async} ->
     line i ppf "Pexp_fun\n";
     let () = if async then line i ppf "async\n" in
     line i ppf "arity:%d\n" (List.length params);
+    List.iter
+      (fun ((name : string loc), attrs) ->
+        attributes i ppf attrs;
+        line i ppf "newtype \"%s\"\n" name.txt)
+      newtypes;
     List.iter
       (fun {p_attrs; p_lbl; p_default; p_pat} ->
         attributes i ppf p_attrs;