Preserve comment placement after a then or else (#2589)

* Improve comment placement after a `then` or `else`

Allow comments on the same line as `then` and `else`, as it was the case
in 0.26.2. Also, make sure to avoid formatting any code after a comment
in that position.

This was broken since #2507.

* Preserve placement of comments after a `then` or `else`

To avoid generating a large number of diffs, the locations of then and else
keywords are added to the AST. This new information is used to preserve
different cases:

    if cond then (* attached to then *)
      expr;

    if cond then
      (* attached to expr *)
      expr;

    if cond then (* attached to expr *) expr;

CHANGES.md

@@ -43,7 +43,7 @@ profile. This started with version 0.26.0.
 - \* Force a break around comments following an infix operator (fix non-stabilizing comments) (#2478, @gpetiot)
 - \* Fix the indentation of tuples in attributes and extensions (#2488, @Julow)
 - Fix unstable comment around docked functor argument (#2506, @Julow)
-- \* Fix unwanted alignment after comment (#2507, @Julow)
+- \* Fix unwanted alignment after comment (#2507, #2589, @Julow)
 - \* Fix unwanted alignment in if-then-else (#2511, @Julow)
 - Fix position of comments around and within `(type ...)` function arguments (#2503, @gpetiot)
 - Fix missing parentheses around constraint expressions with attributes (#2513, @alanechang)

lib/Ast.ml

@@ -1424,7 +1424,7 @@ end = struct
         | Pexp_ifthenelse (eN, e) ->
             assert (
               List.exists eN ~f:(fun x -> f x.if_cond || f x.if_body)
-              || Option.exists e ~f )
+              || Option.exists e ~f:(fun (x, _) -> f x) )
         | Pexp_for (_, e1, e2, _, e3) ->
             assert (e1 == exp || e2 == exp || e3 == exp)
         | Pexp_override e1N -> assert (List.exists e1N ~f:snd_f) )
@@ -1991,7 +1991,7 @@ end = struct
         | Pexp_assert e
          |Pexp_construct (_, Some e)
          |Pexp_function (_, _, Pfunction_body e)
-         |Pexp_ifthenelse (_, Some e)
+         |Pexp_ifthenelse (_, Some (e, _))
          |Pexp_prefix (_, e)
          |Pexp_infix (_, _, e)
          |Pexp_lazy e
@@ -2066,7 +2066,7 @@ end = struct
       match exp.pexp_desc with
       | Pexp_assert e
        |Pexp_construct (_, Some e)
-       |Pexp_ifthenelse (_, Some e)
+       |Pexp_ifthenelse (_, Some (e, _))
        |Pexp_prefix (_, e)
        |Pexp_infix (_, _, e)
        |Pexp_lazy e
@@ -2204,7 +2204,7 @@ end = struct
            && List.exists eN ~f:(fun x -> x.if_body == exp)
            && ifthenelse pexp_desc ->
         true
-    | Exp {pexp_desc= Pexp_ifthenelse (_, Some e); _}, {pexp_desc; _}
+    | Exp {pexp_desc= Pexp_ifthenelse (_, Some (e, _)); _}, {pexp_desc; _}
       when !parens_ite && e == exp && ifthenelse pexp_desc ->
         true
     | ( Exp {pexp_desc= Pexp_infix (_, _, e1); _}
@@ -2294,7 +2294,8 @@ end = struct
       | Pexp_ifthenelse (eN, _)
         when List.exists eN ~f:(fun x -> x.if_body == exp) ->
           exposed_right_exp ThenElse exp
-      | Pexp_ifthenelse (_, Some els) when els == exp -> Exp.is_sequence exp
+      | Pexp_ifthenelse (_, Some (els, _)) when els == exp ->
+          Exp.is_sequence exp
       | Pexp_apply (({pexp_desc= Pexp_new _; _} as exp2), _) when exp2 == exp
         ->
           false

lib/Fmt_ast.ml

@@ -2326,7 +2326,7 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
   | Pexp_ifthenelse (if_branches, else_) ->
       let last_loc =
         match else_ with
-        | Some e -> e.pexp_loc
+        | Some (e, _) -> e.pexp_loc
         | None -> (List.last_exn if_branches).if_body.pexp_loc
       in
       Cmts.relocate c.cmts ~src:pexp_loc ~before:pexp_loc ~after:last_loc ;
@@ -2335,25 +2335,36 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
         let with_conds =
           List.map if_branches ~f:(fun x ->
               ( Some (sub_exp ~ctx x.if_cond)
+              , x.if_loc_then
               , sub_exp ~ctx x.if_body
               , x.if_attrs ) )
         in
         match else_ with
-        | Some x ->
-            List.rev ((None, sub_exp ~ctx x, []) :: List.rev with_conds)
+        | Some (x, loc_else) ->
+            List.rev
+              ((None, loc_else, sub_exp ~ctx x, []) :: List.rev with_conds)
         | None -> with_conds
       in
       pro
       $ hvbox 0
           ( Params.Exp.wrap c.conf ~parens:(parens || has_attr)
               (hvbox 0
                  (list_fl cnd_exps
-                    (fun ~first ~last (xcond, xbch, pexp_attributes) ->
+                    (fun
+                      ~first
+                      ~last
+                      (xcond, keyword_loc, xbch, pexp_attributes)
+                    ->
                       let symbol_parens = Exp.is_symbol xbch.ast in
                       let parens_bch =
                         parenze_exp xbch && not symbol_parens
                       in
-                      let parens_exp = false in
+                      let cmts_before_kw = Cmts.fmt_before c keyword_loc in
+                      let cmts_after_kw =
+                        if Cmts.has_after c.cmts keyword_loc then
+                          Some (Cmts.fmt_after c keyword_loc)
+                        else None
+                      in
                       let p =
                         Params.get_if_then_else c.conf ~first ~last
                           ~parens_bch ~parens_prev_bch:!parens_prev_bch
@@ -2364,6 +2375,7 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                           ~fmt_attributes:
                             (fmt_attributes c ~pre:Blank pexp_attributes)
                           ~fmt_cond:(fmt_expression ~box:false c)
+                          ~cmts_before_kw ~cmts_after_kw
                       in
                       parens_prev_bch := parens_bch ;
                       p.box_branch
@@ -2372,7 +2384,7 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                             ( p.branch_pro
                             $ p.wrap_parens
                                 ( fmt_expression c ?box:p.box_expr
-                                    ~parens:parens_exp ?pro:p.expr_pro
+                                    ~parens:false ?pro:p.expr_pro
                                     ?eol:p.expr_eol p.branch_expr
                                 $ p.break_end_branch ) ) )
                       $ fmt_if (not last) p.space_between_branches ) ) )
@@ -2924,8 +2936,7 @@ and fmt_class_signature c ~ctx ~pro ~epi ?ext self_ fields =
   in
   let ast x = Ctf x in
   let cmts_within =
-    if List.is_empty fields then
-      (* Side effect order is important. *)
+    if List.is_empty fields then (* Side effect order is important. *)
       Cmts.fmt_within ~pro:noop c (Ast.location ctx)
     else noop
   in

lib/Params.ml

@@ -692,7 +692,8 @@ type if_then_else =
   ; space_between_branches: Fmt.t }
 
 let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
-    ~xcond ~xbch ~expr_loc ~fmt_extension_suffix ~fmt_attributes ~fmt_cond =
+    ~xcond ~xbch ~expr_loc ~fmt_extension_suffix ~fmt_attributes ~fmt_cond
+    ~cmts_before_kw ~cmts_after_kw =
   let imd = c.fmt_opts.indicate_multiline_delimiters.v in
   let beginend, branch_expr =
     let ast = xbch.Ast.ast in
@@ -721,22 +722,28 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
   let cond () =
     match xcond with
     | Some xcnd ->
-        hvbox 0
-          ( hvbox 2
-              ( fmt_if (not first) (str "else ")
-              $ str "if"
-              $ fmt_if first (fmt_opt fmt_extension_suffix)
-              $ fmt_attributes $ space_break $ fmt_cond xcnd )
-          $ space_break $ str "then" )
-    | None -> str "else"
+        hvbox 2
+          ( hvbox 0
+              ( hvbox 2
+                  ( fmt_if (not first) (str "else ")
+                  $ str "if"
+                  $ fmt_if first (fmt_opt fmt_extension_suffix)
+                  $ fmt_attributes $ space_break $ fmt_cond xcnd )
+              $ space_break $ cmts_before_kw $ str "then" )
+          $ opt cmts_after_kw Fn.id )
+    | None -> cmts_before_kw $ hvbox 2 (str "else" $ opt cmts_after_kw Fn.id)
+  in
+  let branch_pro ?(indent = 2) () =
+    if Option.is_some cmts_after_kw then break 1000 indent
+    else if beginend || parens_bch then str " "
+    else break 1 indent
   in
-  let branch_pro = fmt_or (beginend || parens_bch) (str " ") (break 1 2) in
   match c.fmt_opts.if_then_else.v with
   | `Compact ->
       { box_branch= hovbox ~name:"Params.get_if_then_else `Compact" 2
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro= fmt_or (beginend || parens_bch) (str " ") space_break
+      ; branch_pro= branch_pro ~indent:0 ()
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:
@@ -752,7 +759,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
       { box_branch= Fn.id
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro
+      ; branch_pro= branch_pro ()
       ; wrap_parens= wrap_parens ~wrap_breaks:(wrap (break 1000 2) noop)
       ; box_expr= Some false
       ; expr_pro= None
@@ -769,7 +776,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
             | _ -> 0 )
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro
+      ; branch_pro= branch_pro ()
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:
@@ -792,7 +799,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
       { box_branch= Fn.id
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro
+      ; branch_pro= branch_pro ()
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:
@@ -808,20 +815,26 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
           | `Closing_on_separate_line when parens_bch -> str " "
           | _ -> space_break ) }
   | `Keyword_first ->
-      { box_branch= Fn.id
-      ; cond=
-          opt xcond (fun xcnd ->
-              hvbox 2
-                ( fmt_or first
-                    (str "if" $ fmt_opt fmt_extension_suffix)
-                    (str "else if")
-                $ fmt_attributes $ space_break $ fmt_cond xcnd )
-              $ space_break )
-      ; box_keyword_and_expr=
-          (fun k ->
+      let keyword =
+        hvbox 2
+          ( fmt_or (Option.is_some xcond) (str "then") (str "else")
+          $ opt cmts_after_kw Fn.id )
+      and cond =
+        match xcond with
+        | Some xcond ->
             hvbox 2
-              (fmt_or (Option.is_some xcond) (str "then") (str "else") $ k) )
-      ; branch_pro= fmt_or (beginend || parens_bch) (str " ") space_break
+              ( fmt_or first
+                  (str "if" $ fmt_opt fmt_extension_suffix)
+                  (str "else if")
+              $ fmt_attributes $ space_break $ fmt_cond xcond
+              $ cmts_before_kw )
+            $ space_break
+        | None -> cmts_before_kw
+      in
+      { box_branch= Fn.id
+      ; cond
+      ; box_keyword_and_expr= (fun k -> hovbox 2 (keyword $ k))
+      ; branch_pro= branch_pro ~indent:0 ()
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:
@@ -884,8 +897,7 @@ module Align = struct
 
   let module_pack (c : Conf.t) ~me =
     if not c.fmt_opts.ocp_indent_compat.v then false
-    else
-      (* Align when the constraint is not desugared. *)
+    else (* Align when the constraint is not desugared. *)
       match me.pmod_desc with
       | Pmod_structure _ | Pmod_ident _ -> false
       | _ -> true

lib/Params.mli

@@ -191,6 +191,8 @@ val get_if_then_else :
   -> fmt_extension_suffix:Fmt.t option
   -> fmt_attributes:Fmt.t
   -> fmt_cond:(expression Ast.xt -> Fmt.t)
+  -> cmts_before_kw:Fmt.t
+  -> cmts_after_kw:Fmt.t option
   -> if_then_else
 
 val match_indent : ?default:int -> Conf.t -> parens:bool -> ctx:Ast.t -> int

lib/Translation_unit.ml

@@ -370,8 +370,7 @@ let format (type ext std) (ext_fg : ext Extended_ast.t)
         Error
           (Unstable {iteration= i; prev= prev_source; next= fmted; input_name}
           ) )
-      else
-        (* All good, continue *)
+      else (* All good, continue *)
         print_check ~i:(i + 1) ~conf ~prev_source:fmted ext_t_new std_t_new
   in
   try print_check ~i:1 ~conf ~prev_source ext_parsed std_parsed with

test/passing/tests/break_string_literals-never.ml.err

@@ -1,7 +1,7 @@
-Warning: tests/break_string_literals.ml:4 exceeds the margin
-Warning: tests/break_string_literals.ml:7 exceeds the margin
-Warning: tests/break_string_literals.ml:11 exceeds the margin
-Warning: tests/break_string_literals.ml:48 exceeds the margin
-Warning: tests/break_string_literals.ml:51 exceeds the margin
-Warning: tests/break_string_literals.ml:63 exceeds the margin
-Warning: tests/break_string_literals.ml:68 exceeds the margin
+Warning: tests/break_string_literals.ml:3 exceeds the margin
+Warning: tests/break_string_literals.ml:6 exceeds the margin
+Warning: tests/break_string_literals.ml:10 exceeds the margin
+Warning: tests/break_string_literals.ml:47 exceeds the margin
+Warning: tests/break_string_literals.ml:50 exceeds the margin
+Warning: tests/break_string_literals.ml:62 exceeds the margin
+Warning: tests/break_string_literals.ml:67 exceeds the margin

test/passing/tests/break_string_literals-never.ml.ref

@@ -1,6 +1,5 @@
 let () =
-  if true then
-    (* Shrinking the margin a bit *)
+  if true then (* Shrinking the margin a bit *)
     Format.printf
       "@[<v 2>@{<warning>@{<title>Warning@}@}@,@,\        These are @{<warning>NOT@} the Droids you are looking for!@,@,\        Some more text. Just more letters and words.@,\        All this text is left-aligned because it's part of the UI.@,\        It'll be easier for the user to read this message.@]@\n@."
 

test/passing/tests/break_string_literals.ml.ref

@@ -1,6 +1,5 @@
 let () =
-  if true then
-    (* Shrinking the margin a bit *)
+  if true then (* Shrinking the margin a bit *)
     Format.printf
       "@[<v 2>@{<warning>@{<title>Warning@}@}@,\
        @,\

test/passing/tests/ite-compact.ml.ref

@@ -119,11 +119,15 @@ let foo =
 let foo = if cmp < 0 then (* foo *) a + b else (* foo *) a - b
 
 let foo =
-  if cmp < 0 then
-    (* ast higher precedence than context: no parens *)
+  if cmp < 0 then (* foo *)
+    a + b
+  else (* foo *)
+    a - b
+
+let foo =
+  if cmp < 0 then (* ast higher precedence than context: no parens *)
     false
-  else if cmp > 0 then
-    (* context higher prec than ast: add parens *)
+  else if cmp > 0 then (* context higher prec than ast: add parens *)
     true
   else if Poly.(assoc_of_prec prec_ast = which_child && which_child <> Non)
   then foo
@@ -157,3 +161,18 @@ let _ =
     bar
   then 1
   else 2
+
+let compare s1 s2 =
+  if String.equal s1 s2 then (* this simplifies the next two cases *)
+    0
+  else if String.equal s1 Cmdliner.Manpage.s_options then
+    (* ensure OPTIONS section is last (hence first in the manual) *)
+    1
+  else if String.equal s2 Cmdliner.Manpage.s_options then (* same as above *)
+    -1
+  else (* reverse order *)
+    String.compare s2 s1
+
+let _ = if x then 42 (* dummy *) else y
+
+let _ = if x then 42 (* dummy *) else if y then z else w

test/passing/tests/ite-compact_closing.ml.ref

@@ -134,11 +134,15 @@ let foo =
 let foo = if cmp < 0 then (* foo *) a + b else (* foo *) a - b
 
 let foo =
-  if cmp < 0 then
-    (* ast higher precedence than context: no parens *)
+  if cmp < 0 then (* foo *)
+    a + b
+  else (* foo *)
+    a - b
+
+let foo =
+  if cmp < 0 then (* ast higher precedence than context: no parens *)
     false
-  else if cmp > 0 then
-    (* context higher prec than ast: add parens *)
+  else if cmp > 0 then (* context higher prec than ast: add parens *)
     true
   else if Poly.(assoc_of_prec prec_ast = which_child && which_child <> Non)
   then foo
@@ -172,3 +176,18 @@ let _ =
     bar
   then 1
   else 2
+
+let compare s1 s2 =
+  if String.equal s1 s2 then (* this simplifies the next two cases *)
+    0
+  else if String.equal s1 Cmdliner.Manpage.s_options then
+    (* ensure OPTIONS section is last (hence first in the manual) *)
+    1
+  else if String.equal s2 Cmdliner.Manpage.s_options then (* same as above *)
+    -1
+  else (* reverse order *)
+    String.compare s2 s1
+
+let _ = if x then 42 (* dummy *) else y
+
+let _ = if x then 42 (* dummy *) else if y then z else w