wip: await AST

cristianoc · cristianoc · commit c90f3d9eb02a · 2025-04-03T10:45:07.000+02:00
diff --git a/analysis/src/CompletionFrontEnd.ml b/analysis/src/CompletionFrontEnd.ml
@@ -312,11 +312,12 @@ let rec exprToContextPathInner ~(inJsxContext : bool) (e : Parsetree.expression)
     if List.length exprs = List.length exprsAsContextPaths then
       Some (CTuple exprsAsContextPaths)
     else None
+  | Pexp_await e -> exprToContextPathInner ~inJsxContext e
   | _ -> None
 
 and exprToContextPath ~(inJsxContext : bool) (e : Parsetree.expression) =
   match
-    ( Res_parsetree_viewer.has_await_attribute e.pexp_attributes,
+    ( Res_parsetree_viewer.has_await_attribute e,
       exprToContextPathInner ~inJsxContext e )
   with
   | true, Some ctxPath -> Some (CPAwait ctxPath)
diff --git a/analysis/src/Utils.ml b/analysis/src/Utils.ml
@@ -111,6 +111,7 @@ let identifyPexp pexp =
   | Pexp_pack _ -> "Pexp_pack"
   | Pexp_extension _ -> "Pexp_extension"
   | Pexp_open _ -> "Pexp_open"
+  | Pexp_await _ -> "Pexp_await"
 
 let identifyPpat pat =
   match pat with
diff --git a/compiler/frontend/bs_ast_mapper.ml b/compiler/frontend/bs_ast_mapper.ml
@@ -366,6 +366,7 @@ module E = struct
     | Pexp_open (ovf, lid, e) ->
       open_ ~loc ~attrs ovf (map_loc sub lid) (sub.expr sub e)
     | Pexp_extension x -> extension ~loc ~attrs (sub.extension sub x)
+    | Pexp_await e -> await ~loc ~attrs (sub.expr sub e)
 end
 
 module P = struct
diff --git a/compiler/frontend/bs_builtin_ppx.ml b/compiler/frontend/bs_builtin_ppx.ml
@@ -177,7 +177,7 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
   (* module M = await Belt.List *)
   | Pexp_letmodule
       (lid, ({pmod_desc = Pmod_ident {txt}; pmod_attributes} as me), expr)
-    when Res_parsetree_viewer.has_await_attribute pmod_attributes ->
+    when Res_parsetree_viewer.has_await_attribute2 pmod_attributes ->
     let safe_module_type_lid : Ast_helper.lid =
       {txt = Lident (local_module_type_name txt); loc = me.pmod_loc}
     in
@@ -201,8 +201,8 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
            pmod_attributes = attrs2;
          } as me),
         expr )
-    when Res_parsetree_viewer.has_await_attribute attrs1
-         || Res_parsetree_viewer.has_await_attribute attrs2 ->
+    when Res_parsetree_viewer.has_await_attribute2 attrs1
+         || Res_parsetree_viewer.has_await_attribute2 attrs2 ->
     {
       e with
       pexp_desc =
@@ -242,10 +242,12 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
          || Ast_attributes.has_await_payload attrs2 ->
     check_await ();
     result
-  | _ when Ast_attributes.has_await_payload e.pexp_attributes ->
-    check_await ();
-    Ast_await.create_await_expression result
-  | _ -> result
+  | _ -> (
+    match result.pexp_desc with
+    | Pexp_await e ->
+      check_await ();
+      Ast_await.create_await_expression e
+    | _ -> result)
 
 let typ_mapper (self : mapper) (typ : Parsetree.core_type) =
   Ast_core_type_class_type.typ_mapper self typ
@@ -493,7 +495,7 @@ let rec structure_mapper ~await_context (self : mapper) (stru : Ast_structure.t)
     | Pstr_module
         ({pmb_expr = {pmod_desc = Pmod_ident {txt; loc}; pmod_attributes} as me}
          as mb)
-      when Res_parsetree_viewer.has_await_attribute pmod_attributes ->
+      when Res_parsetree_viewer.has_await_attribute2 pmod_attributes ->
       let item = self.structure_item self item in
       let safe_module_type_name = local_module_type_name txt in
       let has_local_module_name =
@@ -544,7 +546,8 @@ let rec structure_mapper ~await_context (self : mapper) (stru : Ast_structure.t)
                 ( _,
                   ({pmod_desc = Pmod_ident {txt; loc}; pmod_attributes} as me),
                   expr )
-              when Res_parsetree_viewer.has_await_attribute pmod_attributes -> (
+              when Res_parsetree_viewer.has_await_attribute2 pmod_attributes
+              -> (
               let safe_module_type_name = local_module_type_name txt in
               let has_local_module_name =
                 Hashtbl.find_opt !await_context safe_module_type_name
diff --git a/compiler/ml/ast_helper.ml b/compiler/ml/ast_helper.ml
@@ -180,7 +180,7 @@ module Exp = struct
   let pack ?loc ?attrs a = mk ?loc ?attrs (Pexp_pack a)
   let open_ ?loc ?attrs a b c = mk ?loc ?attrs (Pexp_open (a, b, c))
   let extension ?loc ?attrs a = mk ?loc ?attrs (Pexp_extension a)
-
+  let await ?loc ?attrs a = mk ?loc ?attrs (Pexp_await a)
   let case lhs ?guard rhs = {pc_lhs = lhs; pc_guard = guard; pc_rhs = rhs}
 end
 
diff --git a/compiler/ml/ast_helper.mli b/compiler/ml/ast_helper.mli
@@ -210,6 +210,7 @@ module Exp : sig
   val extension : ?loc:loc -> ?attrs:attrs -> extension -> expression
 
   val case : pattern -> ?guard:expression -> expression -> case
+  val await : ?loc:loc -> ?attrs:attrs -> expression -> expression
 end
 
 (** Value declarations *)
diff --git a/compiler/ml/ast_iterator.ml b/compiler/ml/ast_iterator.ml
@@ -344,6 +344,7 @@ module E = struct
       iter_loc sub lid;
       sub.expr sub e
     | Pexp_extension x -> sub.extension sub x
+    | Pexp_await e -> sub.expr sub e
 end
 
 module P = struct
diff --git a/compiler/ml/ast_mapper.ml b/compiler/ml/ast_mapper.ml
@@ -329,6 +329,7 @@ module E = struct
     | Pexp_open (ovf, lid, e) ->
       open_ ~loc ~attrs ovf (map_loc sub lid) (sub.expr sub e)
     | Pexp_extension x -> extension ~loc ~attrs (sub.extension sub x)
+    | Pexp_await e -> await ~loc ~attrs (sub.expr sub e)
 end
 
 module P = struct
diff --git a/compiler/ml/ast_mapper_to0.ml b/compiler/ml/ast_mapper_to0.ml
@@ -407,6 +407,7 @@ module E = struct
     | Pexp_open (ovf, lid, e) ->
       open_ ~loc ~attrs ovf (map_loc sub lid) (sub.expr sub e)
     | Pexp_extension x -> extension ~loc ~attrs (sub.extension sub x)
+    | Pexp_await _ -> (* TODO *) assert false
 end
 
 module P = struct
diff --git a/compiler/ml/depend.ml b/compiler/ml/depend.ml
@@ -289,6 +289,7 @@ let rec add_expr bv exp =
     | Pstr_eval ({pexp_desc = Pexp_construct (c, None)}, _) -> add bv c
     | _ -> handle_extension e)
   | Pexp_extension e -> handle_extension e
+  | Pexp_await e -> add_expr bv e
 
 and add_cases bv cases = List.iter (add_case bv) cases
 
diff --git a/compiler/ml/parsetree.ml b/compiler/ml/parsetree.ml
@@ -313,8 +313,9 @@ and expression_desc =
        let open M in E
        let! open M in E *)
   | Pexp_extension of extension
-(* [%id] *)
-(* . *)
+  (* [%id] *)
+  (* . *)
+  | Pexp_await of expression
 
 and case = {
   (* (P -> E) or (P when E0 -> E) *)
diff --git a/compiler/ml/pprintast.ml b/compiler/ml/pprintast.ml
@@ -722,6 +722,7 @@ and expression ctxt f x =
         (expression ctxt) e
     | Pexp_variant (l, Some eo) -> pp f "@[<2>`%s@;%a@]" l (simple_expr ctxt) eo
     | Pexp_extension e -> extension ctxt f e
+    | Pexp_await e -> pp f "@[<hov2>await@ %a@]" (simple_expr ctxt) e
     | _ -> expression1 ctxt f x
 
 and expression1 ctxt f x =
diff --git a/compiler/ml/printast.ml b/compiler/ml/printast.ml
@@ -345,6 +345,9 @@ and expression i ppf x =
   | Pexp_extension (s, arg) ->
     line i ppf "Pexp_extension \"%s\"\n" s.txt;
     payload i ppf arg
+  | Pexp_await e ->
+    line i ppf "Pexp_await\n";
+    expression i ppf e
 
 and value_description i ppf x =
   line i ppf "value_description %a %a\n" fmt_string_loc x.pval_name fmt_location
diff --git a/compiler/ml/typecore.ml b/compiler/ml/typecore.ml
@@ -178,6 +178,7 @@ let iter_expression f e =
       expr e;
       module_expr me
     | Pexp_pack me -> module_expr me
+    | Pexp_await _ -> assert false (* should be handled earlier *)
   and case {pc_lhs = _; pc_guard; pc_rhs} =
     may expr pc_guard;
     expr pc_rhs
@@ -3195,6 +3196,7 @@ and type_expect_ ?type_clash_context ?in_function ?(recarg = Rejected) env sexp
     | _ -> raise (Error (loc, env, Invalid_extension_constructor_payload)))
   | Pexp_extension ext ->
     raise (Error_forward (Builtin_attributes.error_of_extension ext))
+  | Pexp_await _ -> (* should be handled earlier *) assert false
 
 and type_function ?in_function ~arity ~async loc attrs env ty_expected_ l
     caselist =
diff --git a/compiler/syntax/src/res_ast_debugger.ml b/compiler/syntax/src/res_ast_debugger.ml
@@ -707,6 +707,7 @@ module SexpAst = struct
           ]
       | Pexp_extension ext ->
         Sexp.list [Sexp.atom "Pexp_extension"; extension ext]
+      | Pexp_await e -> Sexp.list [Sexp.atom "Pexp_await"; expression e]
     in
     Sexp.list [Sexp.atom "expression"; desc]
 
diff --git a/compiler/syntax/src/res_core.ml b/compiler/syntax/src/res_core.ml
@@ -3307,15 +3307,12 @@ and parse_async_arrow_expression ?(arrow_attrs = []) p =
 
 and parse_await_expression p =
   let await_loc = mk_loc p.Parser.start_pos p.end_pos in
-  let await_attr = make_await_attr await_loc in
   Parser.expect Await p;
   let token_prec = Token.precedence MinusGreater in
   let expr = parse_binary_expr ~context:OrdinaryExpr p token_prec in
-  {
-    expr with
-    pexp_attributes = await_attr :: expr.pexp_attributes;
-    pexp_loc = {expr.pexp_loc with loc_start = await_loc.loc_start};
-  }
+  Ast_helper.Exp.await
+    ~loc:{expr.pexp_loc with loc_start = await_loc.loc_start}
+    ~attrs:[] expr
 
 and parse_try_expression p =
   let start_pos = p.Parser.start_pos in
diff --git a/compiler/syntax/src/res_parens.ml b/compiler/syntax/src/res_parens.ml
@@ -56,8 +56,7 @@ let call_expr expr =
     } ->
       Parenthesized
     | _ when Ast_uncurried.expr_is_uncurried_fun expr -> Parenthesized
-    | _ when ParsetreeViewer.has_await_attribute expr.pexp_attributes ->
-      Parenthesized
+    | _ when ParsetreeViewer.has_await_attribute expr -> Parenthesized
     | _ -> Nothing)
 
 let structure_expr expr =
@@ -109,8 +108,7 @@ let unary_expr_operand expr =
        | Pexp_try _ | Pexp_while _ | Pexp_for _ | Pexp_ifthenelse _ );
     } ->
       Parenthesized
-    | _ when ParsetreeViewer.has_await_attribute expr.pexp_attributes ->
-      Parenthesized
+    | _ when ParsetreeViewer.has_await_attribute expr -> Parenthesized
     | _ -> Nothing)
 
 let binary_expr_operand ~is_lhs expr =
@@ -133,8 +131,7 @@ let binary_expr_operand ~is_lhs expr =
     | expr when ParsetreeViewer.is_binary_expression expr -> Parenthesized
     | expr when ParsetreeViewer.is_ternary_expr expr -> Parenthesized
     | {pexp_desc = Pexp_lazy _ | Pexp_assert _} when is_lhs -> Parenthesized
-    | _ when ParsetreeViewer.has_await_attribute expr.pexp_attributes ->
-      Parenthesized
+    | _ when ParsetreeViewer.has_await_attribute expr -> Parenthesized
     | {Parsetree.pexp_attributes = attrs} ->
       if ParsetreeViewer.has_printable_attributes attrs then Parenthesized
       else Nothing)
@@ -229,9 +226,7 @@ let lazy_or_assert_or_await_expr_rhs ?(in_await = false) expr =
        | Pexp_while _ | Pexp_for _ | Pexp_ifthenelse _ );
     } ->
       Parenthesized
-    | _
-      when (not in_await)
-           && ParsetreeViewer.has_await_attribute expr.pexp_attributes ->
+    | _ when (not in_await) && ParsetreeViewer.has_await_attribute expr ->
       Parenthesized
     | _ -> Nothing)
 
@@ -277,8 +272,7 @@ let field_expr expr =
        | Pexp_try _ | Pexp_while _ | Pexp_for _ | Pexp_ifthenelse _ );
     } ->
       Parenthesized
-    | _ when ParsetreeViewer.has_await_attribute expr.pexp_attributes ->
-      Parenthesized
+    | _ when ParsetreeViewer.has_await_attribute expr -> Parenthesized
     | _ -> Nothing)
 
 let set_field_expr_rhs expr =
@@ -339,8 +333,7 @@ let jsx_prop_expr expr =
       }
         when starts_with_minus x ->
         Parenthesized
-      | _ when ParsetreeViewer.has_await_attribute expr.pexp_attributes ->
-        Parenthesized
+      | _ when ParsetreeViewer.has_await_attribute expr -> Parenthesized
       | {
        Parsetree.pexp_desc =
          ( Pexp_ident _ | Pexp_constant _ | Pexp_field _ | Pexp_construct _
@@ -377,8 +370,7 @@ let jsx_child_expr expr =
       }
         when starts_with_minus x ->
         Parenthesized
-      | _ when ParsetreeViewer.has_await_attribute expr.pexp_attributes ->
-        Parenthesized
+      | _ when ParsetreeViewer.has_await_attribute expr -> Parenthesized
       | {
        Parsetree.pexp_desc =
          ( Pexp_ident _ | Pexp_constant _ | Pexp_field _ | Pexp_construct _
diff --git a/compiler/syntax/src/res_parsetree_viewer.ml b/compiler/syntax/src/res_parsetree_viewer.ml
@@ -65,13 +65,23 @@ let functor_type modtype =
   in
   process [] modtype
 
-let has_await_attribute attrs =
+let has_await_attribute2 attrs =
   List.exists
     (function
       | {Location.txt = "res.await"}, _ -> true
       | _ -> false)
     attrs
 
+let has_await_attribute e =
+  (match e.pexp_desc with
+  | Pexp_await _ -> true
+  | _ -> false)
+  || List.exists
+       (function
+         | {Location.txt = "res.await"}, _ -> true
+         | _ -> false)
+       e.pexp_attributes
+
 let has_inline_record_definition_attribute attrs =
   List.exists
     (function
diff --git a/compiler/syntax/src/res_parsetree_viewer.mli b/compiler/syntax/src/res_parsetree_viewer.mli
@@ -14,7 +14,8 @@ val functor_type :
   list
   * Parsetree.module_type
 
-val has_await_attribute : Parsetree.attributes -> bool
+val has_await_attribute : Parsetree.expression -> bool
+val has_await_attribute2 : Parsetree.attributes -> bool
 val has_inline_record_definition_attribute : Parsetree.attributes -> bool
 val has_res_pat_variant_spread_attribute : Parsetree.attributes -> bool
 val has_dict_pattern_attribute : Parsetree.attributes -> bool
diff --git a/compiler/syntax/src/res_printer.ml b/compiler/syntax/src/res_printer.ml
@@ -700,7 +700,7 @@ and print_module_binding ~state ~is_rec module_binding cmt_tbl i =
     match module_binding.pmb_expr with
     | {pmod_desc = Pmod_constraint (mod_expr, mod_type)}
       when not
-             (ParsetreeViewer.has_await_attribute
+             (ParsetreeViewer.has_await_attribute2
                 module_binding.pmb_expr.pmod_attributes) ->
       ( print_mod_expr ~state mod_expr cmt_tbl,
         Doc.concat [Doc.text ": "; print_mod_type ~state mod_type cmt_tbl] )
@@ -3432,9 +3432,10 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
         Doc.concat [Doc.text "\""; member_doc; Doc.text "\""]
       in
       Doc.group (Doc.concat [parent_doc; Doc.lbracket; member; Doc.rbracket])
-  in
-  let expr_with_await =
-    if ParsetreeViewer.has_await_attribute e.pexp_attributes then
+    | Pexp_await e ->
+      let printed_expression =
+        print_expression_with_comments ~state e cmt_tbl
+      in
       let rhs =
         match
           Parens.lazy_or_assert_or_await_expr_rhs ~in_await:true
@@ -3453,7 +3454,6 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
         | Nothing -> printed_expression
       in
       Doc.concat [Doc.text "await "; rhs]
-    else printed_expression
   in
   let should_print_its_own_attributes =
     match e.pexp_desc with
@@ -3467,11 +3467,11 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
     | _ -> false
   in
   match e.pexp_attributes with
-  | [] -> expr_with_await
+  | [] -> printed_expression
   | attrs when not should_print_its_own_attributes ->
     Doc.group
-      (Doc.concat [print_attributes ~state attrs cmt_tbl; expr_with_await])
-  | _ -> expr_with_await
+      (Doc.concat [print_attributes ~state attrs cmt_tbl; printed_expression])
+  | _ -> printed_expression
 
 and print_pexp_fun ~state ~in_callback e cmt_tbl =
   let async, parameters, return_expr = ParsetreeViewer.fun_expr e in
@@ -3761,7 +3761,10 @@ and print_binary_expression ~state (expr : Parsetree.expression) cmt_tbl =
               | _ -> add_parens doc
             in
             let is_await =
-              ParsetreeViewer.has_await_attribute expr.pexp_attributes
+              (match expr.pexp_desc with
+              | Pexp_await _ -> true
+              | _ -> false)
+              || ParsetreeViewer.has_await_attribute expr
             in
             let doc =
               if is_await then
@@ -5298,7 +5301,7 @@ and print_expression_block ~state ~braces expr cmt_tbl =
         match mod_expr.pmod_desc with
         | Pmod_constraint (mod_expr2, mod_type)
           when not
-                 (ParsetreeViewer.has_await_attribute mod_expr.pmod_attributes)
+                 (ParsetreeViewer.has_await_attribute2 mod_expr.pmod_attributes)
           ->
           let name =
             Doc.concat
@@ -5792,7 +5795,7 @@ and print_mod_expr ~state mod_expr cmt_tbl =
     | Pmod_functor _ -> print_mod_functor ~state mod_expr cmt_tbl
   in
   let doc =
-    if ParsetreeViewer.has_await_attribute mod_expr.pmod_attributes then
+    if ParsetreeViewer.has_await_attribute2 mod_expr.pmod_attributes then
       match mod_expr.pmod_desc with
       | Pmod_constraint _ ->
         Doc.concat [Doc.text "await "; Doc.lparen; doc; Doc.rparen]
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/syntax_tests/data/parsing/grammar/expressions/expected/await.res.txt b/tests/syntax_tests/data/parsing/grammar/expressions/expected/await.res.txt
diff --git a/tests/syntax_tests/data/printer/expr/asyncAwait.res b/tests/syntax_tests/data/printer/expr/asyncAwait.res
diff --git a/tests/syntax_tests/data/printer/expr/expected/asyncAwait.res.txt b/tests/syntax_tests/data/printer/expr/expected/asyncAwait.res.txt
