Extend untyped and types ast with async attribute.

cristianoc · cristianoc · commit c90eb55177e3 · 2025-01-08T16:54:32.000+01:00
diff --git a/compiler/frontend/ast_compatible.ml b/compiler/frontend/ast_compatible.ml
@@ -64,13 +64,20 @@ 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 = []) ~arity pat exp =
+let fun_ ?(loc = default_loc) ?(attrs = []) ?(async = false) ~arity pat exp =
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
     pexp_desc =
       Pexp_fun
-        {arg_label = Nolabel; default = None; lhs = pat; rhs = exp; arity};
+        {
+          arg_label = Nolabel;
+          default = None;
+          lhs = pat;
+          rhs = exp;
+          arity;
+          async;
+        };
   }
 
 let const_exp_string ?(loc = default_loc) ?(attrs = []) ?delimiter (s : string)
diff --git a/compiler/frontend/ast_compatible.mli b/compiler/frontend/ast_compatible.mli
@@ -74,6 +74,7 @@ val apply_labels :
 val fun_ :
   ?loc:Location.t ->
   ?attrs:attrs ->
+  ?async:bool ->
   arity:int option ->
   pattern ->
   expression ->
diff --git a/compiler/frontend/ast_uncurry_gen.ml b/compiler/frontend/ast_uncurry_gen.ml
@@ -36,16 +36,16 @@ 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; lhs = arg; rhs = body} ->
+      | Pexp_fun {arg_label; lhs = arg; rhs = body; async} ->
         Bs_syntaxerr.optional_err loc arg_label;
-        aux ((arg_label, self.pat self arg, attrs) :: acc) body
+        aux ((arg_label, self.pat self arg, attrs, async) :: acc) body
       | _ -> (self.expr self body, acc))
     | _, _ -> (self.expr self body, acc)
   in
-  let result, rev_extra_args = aux [(label, self_pat, [])] body in
+  let result, rev_extra_args = aux [(label, self_pat, [], false)] body in
   let body =
-    Ext_list.fold_left rev_extra_args result (fun e (label, p, attrs) ->
-        Ast_helper.Exp.fun_ ~loc ~attrs ~arity:None label None p e)
+    Ext_list.fold_left rev_extra_args result (fun e (label, p, attrs, async) ->
+        Ast_helper.Exp.fun_ ~loc ~attrs ~arity:None ~async label None p e)
   in
   let arity = List.length rev_extra_args in
   let arity_s = string_of_int arity in
diff --git a/compiler/frontend/bs_ast_mapper.ml b/compiler/frontend/bs_ast_mapper.ml
@@ -315,8 +315,9 @@ module E = struct
            sub vbs)
         (sub.expr sub e)
     (* #end *)
-    | Pexp_fun {arg_label = lab; default = def; lhs = p; rhs = e; arity} ->
-      fun_ ~loc ~attrs ~arity lab
+    | Pexp_fun {arg_label = lab; default = def; lhs = p; rhs = e; arity; async}
+      ->
+      fun_ ~loc ~attrs ~arity ~async lab
         (map_opt (sub.expr sub) def)
         (sub.pat sub p) (sub.expr sub e)
     | Pexp_apply (e, l) ->
diff --git a/compiler/frontend/bs_builtin_ppx.ml b/compiler/frontend/bs_builtin_ppx.ml
@@ -113,8 +113,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 {arg_label = label; lhs = pat; rhs = body} -> (
-    let async = Ast_async.has_async_payload e.pexp_attributes in
+  | Pexp_fun {arg_label = label; lhs = pat; rhs = body; async} -> (
     match Ast_attributes.process_attributes_rev e.pexp_attributes with
     | Nothing, _ ->
       (* Handle @async x => y => ... is in async context *)
diff --git a/compiler/ml/ast_async.ml b/compiler/ml/ast_async.ml
@@ -1,29 +1,16 @@
-let has_async_payload attrs =
-  Ext_list.exists attrs (fun ({Location.txt}, _) -> txt = "res.async")
-
 let rec dig_async_payload_from_function (expr : Parsetree.expression) =
   match expr.pexp_desc with
-  | Pexp_fun _ -> has_async_payload expr.pexp_attributes
+  | Pexp_fun {async} -> async
   | Pexp_newtype (_, body) -> dig_async_payload_from_function body
   | _ -> false
 
 let add_async_attribute ~async (body : Parsetree.expression) =
-  let add (exp : Parsetree.expression) =
-    if has_async_payload exp.pexp_attributes then exp
-    else
-      {
-        exp with
-        pexp_attributes =
-          ({txt = "res.async"; loc = Location.none}, PStr [])
-          :: exp.pexp_attributes;
-      }
-  in
   if async then
     let rec add_to_fun (exp : Parsetree.expression) =
       match exp.pexp_desc with
       | Pexp_newtype (txt, e) ->
         {exp with pexp_desc = Pexp_newtype (txt, add_to_fun e)}
-      | Pexp_fun _ -> add exp
+      | Pexp_fun f -> {exp with pexp_desc = Pexp_fun {f with async}}
       | _ -> exp
     in
     add_to_fun body
diff --git a/compiler/ml/ast_helper.ml b/compiler/ml/ast_helper.ml
@@ -151,9 +151,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 ~arity a b c d =
+  let fun_ ?loc ?attrs ?(async = false) ~arity a b c d =
     mk ?loc ?attrs
-      (Pexp_fun {arg_label = a; default = b; lhs = c; rhs = d; arity})
+      (Pexp_fun {arg_label = a; default = b; lhs = c; rhs = d; arity; async})
   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))
   let try_ ?loc ?attrs a b = mk ?loc ?attrs (Pexp_try (a, b))
diff --git a/compiler/ml/ast_helper.mli b/compiler/ml/ast_helper.mli
@@ -138,6 +138,7 @@ module Exp : sig
   val fun_ :
     ?loc:loc ->
     ?attrs:attrs ->
+    ?async:bool ->
     arity:int option ->
     arg_label ->
     expression option ->
diff --git a/compiler/ml/ast_mapper.ml b/compiler/ml/ast_mapper.ml
@@ -278,8 +278,9 @@ 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 {arg_label = lab; default = def; lhs = p; rhs = e; arity} ->
-      fun_ ~loc ~attrs ~arity lab
+    | Pexp_fun {arg_label = lab; default = def; lhs = p; rhs = e; arity; async}
+      ->
+      fun_ ~loc ~attrs ~arity ~async lab
         (map_opt (sub.expr sub) def)
         (sub.pat sub p) (sub.expr sub e)
     | Pexp_apply (e, l) ->
diff --git a/compiler/ml/ast_mapper_from0.ml b/compiler/ml/ast_mapper_from0.ml
@@ -304,7 +304,8 @@ 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 ~arity:None lab
+      let async = Ext_list.exists attrs (fun ({txt}, _) -> txt = "res.async") in
+      fun_ ~loc ~attrs ~async ~arity:None lab
         (map_opt (sub.expr sub) def)
         (sub.pat sub p) (sub.expr sub e)
     | Pexp_function _ -> assert false
diff --git a/compiler/ml/ast_mapper_to0.ml b/compiler/ml/ast_mapper_to0.ml
@@ -295,7 +295,13 @@ 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 {arg_label = lab; default = def; lhs = p; rhs = e; arity} -> (
+    | Pexp_fun {arg_label = lab; default = def; lhs = p; rhs = e; arity; async}
+      -> (
+      let attrs =
+        if async then
+          ({txt = "res.async"; loc = Location.none}, Pt.PStr []) :: attrs
+        else attrs
+      in
       let e =
         fun_ ~loc ~attrs lab
           (map_opt (sub.expr sub) def)
diff --git a/compiler/ml/parsetree.ml b/compiler/ml/parsetree.ml
@@ -230,6 +230,7 @@ and expression_desc =
       lhs: pattern;
       rhs: expression;
       arity: arity;
+      async: bool;
     }
     (* fun P -> E1                          (Simple, None)
        fun ~l:P -> E1                       (Labelled l, None)
diff --git a/compiler/ml/pprintast.ml b/compiler/ml/pprintast.ml
@@ -605,14 +605,15 @@ 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 {arg_label = l; default = e0; lhs = p; rhs = e; arity} ->
+    | Pexp_fun {arg_label = l; default = e0; lhs = p; rhs = e; arity; async} ->
       let arity_str =
         match arity with
         | None -> ""
         | Some arity -> "[arity:" ^ string_of_int arity ^ "]"
       in
-      pp f "@[<2>fun@;%s%a->@;%a@]" arity_str (label_exp ctxt) (l, e0, p)
-        (expression ctxt) e
+      let async_str = if async then "async " else "" in
+      pp f "@[<2>%sfun@;%s%a->@;%a@]" async_str arity_str (label_exp ctxt)
+        (l, e0, p) (expression ctxt) e
     | Pexp_match (e, l) ->
       pp f "@[<hv0>@[<hv0>@[<2>match %a@]@ with@]%a@]" (expression reset_ctxt) e
         (case_list ctxt) l
diff --git a/compiler/ml/printast.ml b/compiler/ml/printast.ml
@@ -238,8 +238,9 @@ 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 {arg_label = l; default = eo; lhs = p; rhs = e; arity} ->
+  | Pexp_fun {arg_label = l; default = eo; lhs = p; rhs = e; arity; async} ->
     line i ppf "Pexp_fun\n";
+    let () = if async then line i ppf "async\n" in
     let () =
       match arity with
       | None -> ()
diff --git a/compiler/ml/printtyped.ml b/compiler/ml/printtyped.ml
@@ -285,8 +285,10 @@ and expression i ppf x =
     line i ppf "Texp_let %a\n" fmt_rec_flag rf;
     list i value_binding ppf l;
     expression i ppf e
-  | Texp_function {arg_label = p; arity; param; case = case_; partial = _} ->
+  | Texp_function
+      {arg_label = p; arity; async; param; case = case_; partial = _} ->
     line i ppf "Texp_function\n";
+    if async then line i ppf "async\n";
     (match arity with
     | Some arity -> line i ppf "arity: %d\n" arity
     | None -> ());
diff --git a/compiler/ml/tast_mapper.ml b/compiler/ml/tast_mapper.ml
@@ -199,8 +199,9 @@ let expr sub x =
     | Texp_let (rec_flag, list, exp) ->
       let rec_flag, list = sub.value_bindings sub (rec_flag, list) in
       Texp_let (rec_flag, list, sub.expr sub exp)
-    | Texp_function {arg_label; arity; param; case; partial} ->
-      Texp_function {arg_label; arity; param; case = sub.case sub case; partial}
+    | Texp_function {arg_label; arity; param; case; partial; async} ->
+      Texp_function
+        {arg_label; arity; param; case = sub.case sub case; partial; async}
     | Texp_apply (exp, list) ->
       Texp_apply
         (sub.expr sub exp, List.map (tuple2 id (opt (sub.expr sub))) list)
diff --git a/compiler/ml/translcore.ml b/compiler/ml/translcore.ml
@@ -555,7 +555,8 @@ let rec push_defaults loc bindings case partial =
    c_lhs = pat;
    c_guard = None;
    c_rhs =
-     {exp_desc = Texp_function {arg_label; arity; param; case; partial}} as exp;
+     {exp_desc = Texp_function {arg_label; arity; param; case; partial; async}}
+     as exp;
   } ->
     let case = push_defaults exp.exp_loc bindings case partial in
 
@@ -565,7 +566,8 @@ let rec push_defaults loc bindings case partial =
       c_rhs =
         {
           exp with
-          exp_desc = Texp_function {arg_label; arity; param; case; partial};
+          exp_desc =
+            Texp_function {arg_label; arity; param; case; partial; async};
         };
     }
   | {
@@ -671,8 +673,7 @@ and transl_exp0 (e : Typedtree.expression) : Lambda.lambda =
   | Texp_constant cst -> Lconst (Const_base cst)
   | Texp_let (rec_flag, pat_expr_list, body) ->
     transl_let rec_flag pat_expr_list (transl_exp body)
-  | Texp_function {arg_label = _; arity; param; case; partial} -> (
-    let async = Ast_async.has_async_payload e.exp_attributes in
+  | Texp_function {arg_label = _; arity; param; case; partial; async} -> (
     let directive =
       match extract_directive_for_fn e with
       | None -> None
@@ -1050,11 +1051,11 @@ and transl_function loc partial param case =
              param = param';
              case;
              partial = partial';
+             async;
            };
      } as exp;
   }
-    when Parmatch.inactive ~partial pat
-         && not (Ast_async.has_async_payload exp.exp_attributes) ->
+    when Parmatch.inactive ~partial pat && not async ->
     let params, body, return_unit =
       transl_function exp.exp_loc partial' param' case
     in
diff --git a/compiler/ml/typecore.ml b/compiler/ml/typecore.ml
@@ -2364,7 +2364,14 @@ and type_expect_ ?type_clash_context ?in_function ?(recarg = Rejected) env sexp
         exp_env = env;
       }
   | Pexp_fun
-      {arg_label = l; default = Some default; lhs = spat; rhs = sbody; arity} ->
+      {
+        arg_label = l;
+        default = Some default;
+        lhs = spat;
+        rhs = sbody;
+        arity;
+        async;
+      } ->
     assert (is_optional l);
     (* default allowed only with optional argument *)
     let open Ast_helper in
@@ -2402,10 +2409,13 @@ and type_expect_ ?type_clash_context ?in_function ?(recarg = Rejected) env sexp
         [Vb.mk spat smatch]
         sbody
     in
-    type_function ?in_function ~arity loc sexp.pexp_attributes env ty_expected l
+    type_function ?in_function ~arity ~async loc sexp.pexp_attributes env
+      ty_expected l
       [Exp.case pat body]
-  | Pexp_fun {arg_label = l; default = None; lhs = spat; rhs = sbody; arity} ->
-    type_function ?in_function ~arity loc sexp.pexp_attributes env ty_expected l
+  | Pexp_fun
+      {arg_label = l; default = None; lhs = spat; rhs = sbody; arity; async} ->
+    type_function ?in_function ~arity ~async loc sexp.pexp_attributes env
+      ty_expected l
       [Ast_helper.Exp.case spat sbody]
   | Pexp_apply (sfunct, sargs) ->
     assert (sargs <> []);
@@ -3246,7 +3256,8 @@ and type_expect_ ?type_clash_context ?in_function ?(recarg = Rejected) env sexp
   | Pexp_extension ext ->
     raise (Error_forward (Builtin_attributes.error_of_extension ext))
 
-and type_function ?in_function ~arity loc attrs env ty_expected_ l caselist =
+and type_function ?in_function ~arity ~async loc attrs env ty_expected_ l
+    caselist =
   let state = Warnings.backup () in
   (* Disable Unerasable_optional_argument for uncurried functions *)
   let unerasable_optional_argument =
@@ -3304,7 +3315,8 @@ and type_function ?in_function ~arity loc attrs env ty_expected_ l caselist =
   Warnings.restore state;
   re
     {
-      exp_desc = Texp_function {arg_label = l; arity; param; case; partial};
+      exp_desc =
+        Texp_function {arg_label = l; arity; param; case; partial; async};
       exp_loc = loc;
       exp_extra = [];
       exp_type;
diff --git a/compiler/ml/typedtree.ml b/compiler/ml/typedtree.ml
@@ -82,6 +82,7 @@ and expression_desc =
       param: Ident.t;
       case: case;
       partial: partial;
+      async: bool;
     }
   | Texp_apply of expression * (arg_label * expression option) list
   | Texp_match of expression * case list * case list * partial
diff --git a/compiler/ml/typedtree.mli b/compiler/ml/typedtree.mli
@@ -136,6 +136,7 @@ and expression_desc =
       param: Ident.t;
       case: case;
       partial: partial;
+      async: bool;
     }
       (** [Pexp_fun] and [Pexp_function] both translate to [Texp_function].
             See {!Parsetree} for more details.
diff --git a/tests/syntax_tests/data/parsing/grammar/expressions/expected/UncurriedByDefault.res.txt b/tests/syntax_tests/data/parsing/grammar/expressions/expected/UncurriedByDefault.res.txt
@@ -37,9 +37,9 @@ type nonrec unested = (string -> unit (a:1)) -> unit (a:1)
 let (uannpoly : 'a -> string (a:1)) = xx
 let (uannint : int -> string (a:1)) = xx
 let _ = ((fun [arity:1]x -> 34)[@att ])
-let _ = ((fun [arity:1]x -> 34)[@res.async ][@att ])
+let _ = ((async fun [arity:1]x -> 34)[@att ])
 let _ = preserveAttr ((fun [arity:1]x -> 34)[@att ])
-let _ = preserveAttr ((fun [arity:1]x -> 34)[@res.async ][@att ])
+let _ = preserveAttr ((async fun [arity:1]x -> 34)[@att ])
 let t0 (type a) (type b) [arity:2](l : a list) (x : a) = x :: l
 let t1 (type a) (type b) [arity:2](l : a list) (x : a) = x :: l
 let t2 (type a) (type b) [arity:2](l : a list) (x : a) = x :: l
@@ -95,9 +95,9 @@ let pipe1 = 3 |.u f
 let (uannpoly : 'a -> string (a:1)) = xx
 let (uannint : int -> string (a:1)) = xx
 let _ = ((fun [arity:1]x -> 34)[@att ])
-let _ = ((fun [arity:1]x -> 34)[@res.async ][@att ])
+let _ = ((async fun [arity:1]x -> 34)[@att ])
 let _ = preserveAttr ((fun [arity:1]x -> 34)[@att ])
-let _ = preserveAttr ((fun [arity:1]x -> 34)[@res.async ][@att ])
+let _ = preserveAttr ((async fun [arity:1]x -> 34)[@att ])
 let t0 (type a) (type b) [arity:2](l : a list) (x : a) = x :: l
 let t1 (type a) (type b) [arity:2](l : a list) (x : a) = x :: l
 let t2 (type a) (type b) [arity:2](l : a list) (x : a) = x :: l
diff --git a/tests/syntax_tests/data/parsing/grammar/expressions/expected/async.res.txt b/tests/syntax_tests/data/parsing/grammar/expressions/expected/async.res.txt
diff --git a/tests/tools_tests/ppx/TestPpx.res b/tests/tools_tests/ppx/TestPpx.res
diff --git a/tests/tools_tests/src/expected/TestPpx.res.jsout b/tests/tools_tests/src/expected/TestPpx.res.jsout
