Extend untyped and types ast with async attribute.

Author: cristianoc

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)

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 ->

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

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) ->

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 *)

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

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))

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 ->

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) ->

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