Fix module expressions losing attributes, await or parens when formatting

CHANGELOG.md

@@ -23,6 +23,7 @@
 
 #### :bug: Bug fix
 
+- Fix the formatter dropping or moving attributes and `await` on module expressions, such as `module M = @attr F(X)`, and printing invalid code for module constraints and functors in some positions, such as `module M: T = (X: S)`. https://github.com/rescript-lang/rescript/pull/8735
 - Make rewatch compile independent modules after an unrelated failure and recompile blocked dependents when a changed interface survives a failed implementation, including across full watcher rebuilds. https://github.com/rescript-lang/rescript/pull/8667
 - Fix `RescriptTools.binaryPath` failing with `ERR_PACKAGE_IMPORT_NOT_DEFINED` when used from a user project. https://github.com/rescript-lang/rescript/pull/8694
 - GenType: type inline-record cases of untagged variants as objects instead of their single field's type. https://github.com/rescript-lang/rescript/pull/8693

compiler/syntax/src/res_parens.ml

@@ -428,11 +428,23 @@ let braced_expr expr =
   | Pexp_constraint _ -> true
   | _ -> false
 
-let include_mod_expr mod_expr =
+(* A constraint that needs parens where [M: S] would not parse or would mean
+   something else, e.g. after [include] or [module type of]. With attributes
+   or [await] it prints its own. *)
+let mod_constraint mod_expr =
   match mod_expr.Parsetree.pmod_desc with
-  | Parsetree.Pmod_constraint _ -> true
+  | Parsetree.Pmod_constraint _ ->
+    not (Parsetree_viewer.mod_expr_has_attributes mod_expr)
   | _ -> false
 
+(* An applied module expression that needs parens: [(M: S)(X)],
+   [((Y) => M)(X)], [(%ext)(X)], and with attributes or [await], which
+   would otherwise apply to the whole application *)
+let mod_apply_callee callee =
+  match callee.Parsetree.pmod_desc with
+  | Pmod_constraint _ | Pmod_functor _ | Pmod_extension (_, PStr []) -> true
+  | _ -> Parsetree_viewer.mod_expr_has_attributes callee
+
 let mod_expr_parens mod_expr =
   match mod_expr with
   | {

compiler/syntax/src/res_parens.mli

@@ -33,7 +33,8 @@ val mod_expr_functor_constraint : Parsetree.module_type -> bool
 val braced_expr : Parsetree.expression -> bool
 val call_expr : Parsetree.expression -> kind
 
-val include_mod_expr : Parsetree.module_expr -> bool
+val mod_constraint : Parsetree.module_expr -> bool
+val mod_apply_callee : Parsetree.module_expr -> bool
 
 val mod_expr_parens : Parsetree.module_expr -> bool
 

compiler/syntax/src/res_parsetree_viewer.ml

@@ -225,19 +225,20 @@ let rec unwrap_braces expr =
   | Pexp_braces {expr = inner} -> unwrap_braces inner
   | _ -> expr
 
+(* Attributes the parser adds to encode syntax; they are never printed *)
+let is_parsing_attr (attr : Parsetree.attribute) =
+  match attr with
+  | ( {
+        Location.txt =
+          ( "res.iflet" | "res.ternary" | "res.await" | "res.patVariantSpread"
+          | "res.dictPattern" | "res.dictSpread" );
+      },
+      _ ) ->
+    true
+  | _ -> false
+
 let filter_parsing_attrs attrs =
-  List.filter
-    (fun attr ->
-      match attr with
-      | ( {
-            Location.txt =
-              ( "res.iflet" | "res.ternary" | "res.await"
-              | "res.patVariantSpread" | "res.dictPattern" | "res.dictSpread" );
-          },
-          _ ) ->
-        false
-      | _ -> true)
-    attrs
+  List.filter (fun attr -> not (is_parsing_attr attr)) attrs
 
 let is_block_expr expr =
   match (unwrap_braces expr).pexp_desc with
@@ -548,26 +549,37 @@ let is_printable_attribute attr =
 
 let has_printable_attributes attrs = List.exists is_printable_attribute attrs
 
+(* Attributes or [await] on a module expression: either one prints before it
+   and binds less tightly than an application or a constraint. The attributes
+   are the ones [Res_printer.print_attributes] prints. *)
+let mod_expr_has_attributes (mod_expr : Parsetree.module_expr) =
+  List.exists
+    (fun attr ->
+      match attr with
+      | {Location.txt = "res.await"}, _ -> true
+      | _ -> not (is_parsing_attr attr))
+    mod_expr.pmod_attributes
+
 let filter_printable_attributes attrs = List.filter is_printable_attribute attrs
 
 let partition_printable_attributes attrs =
   List.partition is_printable_attribute attrs
 
+let is_doc_comment_attribute ((id, payload) : Parsetree.attribute) =
+  match (id, payload) with
+  | ( {txt = "res.doc"},
+      PStr
+        [
+          {
+            pstr_desc =
+              Pstr_eval ({pexp_desc = Pexp_constant (Pconst_string _)}, _);
+          };
+        ] ) ->
+    true
+  | _ -> false
+
 let partition_doc_comment_attributes attrs =
-  List.partition
-    (fun ((id, payload) : Parsetree.attribute) ->
-      match (id, payload) with
-      | ( {txt = "res.doc"},
-          PStr
-            [
-              {
-                pstr_desc =
-                  Pstr_eval ({pexp_desc = Pexp_constant (Pconst_string _)}, _);
-              };
-            ] ) ->
-        true
-      | _ -> false)
-    attrs
+  List.partition is_doc_comment_attribute attrs
 
 let rec is_fun_expr expr =
   match expr.pexp_desc with
@@ -600,18 +612,24 @@ let requires_special_callback_printing_first_arg args =
 let mod_expr_apply mod_expr =
   let rec loop acc mod_expr =
     match mod_expr with
-    | {pmod_desc = Pmod_apply (next, arg)} -> loop (arg :: acc) next
+    (* An inner application with attributes is kept as the callee, so the
+       printer can parenthesize it with its attributes *)
+    | {pmod_desc = Pmod_apply (next, arg)} as apply
+      when acc = [] || not (mod_expr_has_attributes apply) ->
+      loop (arg :: acc) next
     | _ -> (acc, mod_expr)
   in
   loop [] mod_expr
 
 let mod_expr_functor mod_expr =
   let rec loop acc mod_expr =
     match mod_expr with
+    (* An awaited inner functor is kept as the result, so [await] is printed *)
     | {
      pmod_desc = Pmod_functor (lbl, mod_type, return_mod_expr);
      pmod_attributes = attrs;
-    } ->
+    }
+      when acc = [] || not (has_await_attribute attrs) ->
       let param = (attrs, lbl, mod_type) in
       loop (param :: acc) return_mod_expr
     | return_mod_expr -> (List.rev acc, return_mod_expr)

compiler/syntax/src/res_parsetree_viewer.mli

@@ -96,9 +96,11 @@ val filter_fragile_match_attributes :
 val should_indent_binary_expr : Parsetree.expression -> bool
 val should_inline_rhs_binary_expr : Parsetree.expression -> bool
 val has_printable_attributes : Parsetree.attributes -> bool
+val mod_expr_has_attributes : Parsetree.module_expr -> bool
 val filter_printable_attributes : Parsetree.attributes -> Parsetree.attributes
 val partition_printable_attributes :
   Parsetree.attributes -> Parsetree.attributes * Parsetree.attributes
+val is_doc_comment_attribute : Parsetree.attribute -> bool
 val partition_doc_comment_attributes :
   Parsetree.attributes -> Parsetree.attributes * Parsetree.attributes