Keep a line comment in a block construct's operand in the operand

changes/operand-line-comment.fix.md

@@ -0,0 +1 @@
+- Keep a line comment in the operand of `switch`, `if`, `try` or a loop in the operand instead of moving it into the block when formatting.

compiler/syntax/src/res_doc.ml

@@ -122,6 +122,25 @@ let rec will_break doc =
   | IfBreaks {yes; no} -> will_break yes || will_break no
   | _ -> false
 
+(* See documentation in interface file *)
+let has_pending_line_suffix doc =
+  (* [pending]: a line suffix is queued; [broken]: the innermost enclosing
+     group is printed in break mode, so its line breaks flush the queue. *)
+  let rec walk ~broken pending doc =
+    match doc with
+    | LineSuffix _ -> true
+    | LineBreak _ -> pending && not broken
+    | Nil | Text _ | BreakParent -> pending
+    | Indent doc -> walk ~broken pending doc
+    | Concat docs -> List.fold_left (walk ~broken) pending docs
+    | Group {should_break; doc} ->
+      walk ~broken:(should_break || will_break doc) pending doc
+    | IfBreaks {yes; no; broken = if_broken} ->
+      walk ~broken pending (if broken || if_broken then yes else no)
+    | CustomLayout docs -> List.exists (walk ~broken:false pending) docs
+  in
+  walk ~broken:false false doc
+
 let join ~sep docs =
   let rec loop acc sep docs =
     match docs with

compiler/syntax/src/res_doc.mli

@@ -50,6 +50,13 @@ val will_break : t -> bool
     [will_break] to detect a forced break in a custom layout and explicitly add
     [break_parent] to propagate it to the parent document. *)
 
+val has_pending_line_suffix : t -> bool
+(** [has_pending_line_suffix doc] checks whether a line suffix (a trailing line
+    comment) in [doc] is still queued at its end, because no line break of a
+    group that breaks follows it. Groups that only break for lack of width
+    count as flat. Text after [doc] on the same line would precede the
+    comment. *)
+
 val to_string : width:int -> t -> string
 val debug : t -> unit
 

compiler/syntax/src/res_printer.ml

@@ -3165,6 +3165,38 @@ and print_expression_args ~state (args : Parsetree.expression list) cmt_tbl =
         Doc.rparen;
       ]
 
+(* The operand of [switch], [try], [if], [while] or a loop, which is followed
+   by more code on the same line. A line comment still pending at its end
+   would move after that code, into the block, so the operand goes into
+   parentheses with the comment before the [)]. *)
+and print_operand_before_brace doc =
+  if Doc.has_pending_line_suffix doc then add_operand_parens doc else doc
+
+(* An operand followed by a block, see [print_operand_before_brace], with the
+   parentheses or braces it needs *)
+and print_block_operand ~state expr cmt_tbl =
+  let doc = print_expression_with_comments ~state expr cmt_tbl in
+  print_operand_before_brace
+    (match Parens.expr expr with
+    | Parens.Parenthesized -> add_operand_parens doc
+    | Braced braces -> print_braces doc expr braces
+    | Nothing -> doc)
+
+(* Parentheses around such an operand, or around one that needs them, such as
+   a coercion. They break when a line comment is still pending at the end of
+   the operand, so that the comment stays before the [)]. *)
+and add_operand_parens doc =
+  if Doc.has_pending_line_suffix doc then
+    Doc.breakable_group ~force_break:true
+      (Doc.concat
+         [
+           Doc.lparen;
+           Doc.indent (Doc.concat [Doc.soft_line; doc]);
+           Doc.soft_line;
+           Doc.rparen;
+         ])
+  else add_parens doc
+
 and print_if_chain ~state pexp_attributes ifs else_expr cmt_tbl =
   let if_docs =
     Doc.join ~sep:Doc.space
@@ -3180,14 +3212,14 @@ and print_if_chain ~state pexp_attributes ifs else_expr cmt_tbl =
                    print_expression_with_comments ~state if_expr cmt_tbl
                  in
                  match Parens.expr if_expr with
-                 | Parens.Parenthesized -> add_parens doc
+                 | Parens.Parenthesized -> add_operand_parens doc
                  | Braced braces -> print_braces doc if_expr braces
                  | Nothing -> Doc.if_breaks (add_parens doc) doc
              in
              Doc.concat
                [
                  if_txt;
-                 Doc.group condition;
+                 print_operand_before_brace (Doc.group condition);
                  Doc.space;
                  (let then_expr =
                     match Parsetree_viewer.process_braces then_expr with
@@ -3659,16 +3691,19 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
       let condition =
         let doc = print_expression_with_comments ~state expr1 cmt_tbl in
         match Parens.expr expr1 with
-        | Parens.Parenthesized -> add_parens doc
+        | Parens.Parenthesized -> add_operand_parens doc
         | Braced braces -> print_braces doc expr1 braces
         | Nothing -> doc
       in
       Doc.breakable_group ~force_break:true
         (Doc.concat
            [
              Doc.text "while ";
-             (if Parsetree_viewer.is_block_expr expr1 then condition
-              else Doc.group (Doc.if_breaks (add_parens condition) condition));
+             print_operand_before_brace
+               (match Parens.expr expr1 with
+               | Parens.Parenthesized -> condition
+               | _ when Parsetree_viewer.is_block_expr expr1 -> condition
+               | _ -> Doc.group (Doc.if_breaks (add_parens condition) condition));
              Doc.space;
              print_expression_block ~state ~braces:true expr2 cmt_tbl;
            ])
@@ -3679,19 +3714,9 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
              Doc.text "for ";
              print_pattern ~state pattern cmt_tbl;
              Doc.text " in ";
-             (let doc =
-                print_expression_with_comments ~state from_expr cmt_tbl
-              in
-              match Parens.expr from_expr with
-              | Parens.Parenthesized -> add_parens doc
-              | Braced braces -> print_braces doc from_expr braces
-              | Nothing -> doc);
+             print_block_operand ~state from_expr cmt_tbl;
              print_direction_flag direction_flag;
-             (let doc = print_expression_with_comments ~state to_expr cmt_tbl in
-              match Parens.expr to_expr with
-              | Parens.Parenthesized -> add_parens doc
-              | Braced braces -> print_braces doc to_expr braces
-              | Nothing -> doc);
+             print_block_operand ~state to_expr cmt_tbl;
              Doc.space;
              print_expression_block ~state ~braces:true body cmt_tbl;
            ])
@@ -3702,13 +3727,7 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
              Doc.text "for ";
              print_pattern ~state pattern cmt_tbl;
              Doc.text " of ";
-             (let doc =
-                print_expression_with_comments ~state array_expr cmt_tbl
-              in
-              match Parens.expr array_expr with
-              | Parens.Parenthesized -> add_parens doc
-              | Braced braces -> print_braces doc array_expr braces
-              | Nothing -> doc);
+             print_block_operand ~state array_expr cmt_tbl;
              Doc.space;
              print_expression_block ~state ~braces:true body cmt_tbl;
            ])
@@ -3719,13 +3738,7 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
              Doc.text "for await ";
              print_pattern ~state pattern cmt_tbl;
              Doc.text " of ";
-             (let doc =
-                print_expression_with_comments ~state iterable_expr cmt_tbl
-              in
-              match Parens.expr iterable_expr with
-              | Parens.Parenthesized -> add_parens doc
-              | Braced braces -> print_braces doc iterable_expr braces
-              | Nothing -> doc);
+             print_block_operand ~state iterable_expr cmt_tbl;
              Doc.space;
              print_expression_block ~state ~braces:true body cmt_tbl;
            ])
@@ -3805,13 +3818,7 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
     | Pexp_sequence _ -> print_expression_block ~state ~braces:true e cmt_tbl
     | Pexp_let _ -> print_expression_block ~state ~braces:true e cmt_tbl
     | Pexp_try (expr, cases, finally) ->
-      let expr_doc =
-        let doc = print_expression_with_comments ~state expr cmt_tbl in
-        match Parens.expr expr with
-        | Parens.Parenthesized -> add_parens doc
-        | Braced braces -> print_braces doc expr braces
-        | Nothing -> doc
-      in
+      let expr_doc = print_block_operand ~state expr cmt_tbl in
       let catch_doc =
         match cases with
         | [] -> Doc.nil
@@ -3830,13 +3837,7 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
       in
       Doc.concat [Doc.text "try "; expr_doc; catch_doc; finally_doc]
     | Pexp_match (expr, cases) ->
-      let expr_doc =
-        let doc = print_expression_with_comments ~state expr cmt_tbl in
-        match Parens.expr expr with
-        | Parens.Parenthesized -> add_parens doc
-        | Braced braces -> print_braces doc expr braces
-        | Nothing -> doc
-      in
+      let expr_doc = print_block_operand ~state expr cmt_tbl in
       Doc.concat
         [
           Doc.text "switch ";