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 ocaml-ppx#2507.

lib/Fmt_ast.ml

@@ -2354,6 +2354,13 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                         parenze_exp xbch && not symbol_parens
                       in
                       let parens_exp = false in
+                      let keyword_comments, has_keyword_comments =
+                        let exp_loc = xbch.ast.pexp_loc in
+                        let has = Cmts.has_before c.cmts exp_loc in
+                        let pro = break 1 0 in
+                        ( Cmts.fmt_before ~pro ~epi:noop ~eol:noop c exp_loc
+                        , has )
+                      in
                       let p =
                         Params.get_if_then_else c.conf ~first ~last
                           ~parens_bch ~parens_prev_bch:!parens_prev_bch
@@ -2364,6 +2371,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)
+                          ~keyword_comments ~has_keyword_comments
                       in
                       parens_prev_bch := parens_bch ;
                       p.box_branch
@@ -2924,8 +2932,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

@@ -698,7 +698,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
+    ~keyword_comments ~has_keyword_comments =
   let imd = c.fmt_opts.indicate_multiline_delimiters.v in
   let beginend, branch_expr =
     let ast = xbch.Ast.ast in
@@ -727,22 +728,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 $ str "then" )
+          $ keyword_comments )
+    | None -> hvbox 2 (str "else" $ keyword_comments)
+  in
+  let branch_pro ?(indent = 2) () =
+    if has_keyword_comments 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:
@@ -758,7 +765,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
@@ -775,7 +782,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:
@@ -798,7 +805,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:
@@ -814,6 +821,11 @@ 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 ->
+      let keyword =
+        hvbox 2
+          ( fmt_or (Option.is_some xcond) (str "then") (str "else")
+          $ keyword_comments )
+      in
       { box_branch= Fn.id
       ; cond=
           opt xcond (fun xcnd ->
@@ -823,11 +835,8 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
                     (str "else if")
                 $ fmt_attributes $ space_break $ fmt_cond xcnd )
               $ space_break )
-      ; box_keyword_and_expr=
-          (fun k ->
-            hvbox 2
-              (fmt_or (Option.is_some xcond) (str "then") (str "else") $ k) )
-      ; branch_pro= fmt_or (beginend || parens_bch) (str " ") space_break
+      ; box_keyword_and_expr= (fun k -> hovbox 2 (keyword $ k))
+      ; branch_pro= branch_pro ~indent:0 ()
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:
@@ -890,8 +899,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)
+  -> keyword_comments:Fmt.t
+  -> has_keyword_comments:bool
   -> 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/comments-no-wrap.ml.err

@@ -1,5 +1,5 @@
-Warning: tests/comments.ml:186 exceeds the margin
-Warning: tests/comments.ml:190 exceeds the margin
-Warning: tests/comments.ml:250 exceeds the margin
-Warning: tests/comments.ml:401 exceeds the margin
-Warning: tests/comments.ml:433 exceeds the margin
+Warning: tests/comments.ml:198 exceeds the margin
+Warning: tests/comments.ml:202 exceeds the margin
+Warning: tests/comments.ml:262 exceeds the margin
+Warning: tests/comments.ml:413 exceeds the margin
+Warning: tests/comments.ml:445 exceeds the margin

test/passing/tests/comments-no-wrap.ml.ref

@@ -37,16 +37,28 @@ let foo = function Blah, (x, (* old *) y) -> ()
 let foo = function (x, y) (* old *), z -> ()
 
 let _ =
-  if (* a0 *) b (* c0 *) then (* d0 *) e (* f0 *) else (* g0 *) h (* i0 *)
+  if (* a0 *) b (* c0 *) then (* d0 *)
+    e (* f0 *)
+  else (* g0 *)
+    h (* i0 *)
 
 let _ =
-  if (* a1 *) b (* c1 *) then (* d1 *) e (* f1 *) else (* g1 *) h (* i1 *)
+  if (* a1 *) b (* c1 *) then (* d1 *)
+    e (* f1 *)
+  else (* g1 *)
+    h (* i1 *)
 
 let _ =
-  if (* a2 *) B (* c2 *) then (* d2 *) E (* f2 *) else (* g2 *) H (* i2 *)
+  if (* a2 *) B (* c2 *) then (* d2 *)
+    E (* f2 *)
+  else (* g2 *)
+    H (* i2 *)
 
 let _ =
-  if (* a3 *) B (* c3 *) then (* d3 *) E (* f3 *) else (* g3 *) H (* i3 *)
+  if (* a3 *) B (* c3 *) then (* d3 *)
+    E (* f3 *)
+  else (* g3 *)
+    H (* i3 *)
 ;;
 
 match x with

test/passing/tests/comments.ml.err

@@ -1 +1 @@
-Warning: tests/comments.ml:252 exceeds the margin
+Warning: tests/comments.ml:264 exceeds the margin

test/passing/tests/comments.ml.ref

@@ -37,16 +37,28 @@ let foo = function Blah, (x, (* old *) y) -> ()
 let foo = function (x, y) (* old *), z -> ()
 
 let _ =
-  if (* a0 *) b (* c0 *) then (* d0 *) e (* f0 *) else (* g0 *) h (* i0 *)
+  if (* a0 *) b (* c0 *) then (* d0 *)
+    e (* f0 *)
+  else (* g0 *)
+    h (* i0 *)
 
 let _ =
-  if (* a1 *) b (* c1 *) then (* d1 *) e (* f1 *) else (* g1 *) h (* i1 *)
+  if (* a1 *) b (* c1 *) then (* d1 *)
+    e (* f1 *)
+  else (* g1 *)
+    h (* i1 *)
 
 let _ =
-  if (* a2 *) B (* c2 *) then (* d2 *) E (* f2 *) else (* g2 *) H (* i2 *)
+  if (* a2 *) B (* c2 *) then (* d2 *)
+    E (* f2 *)
+  else (* g2 *)
+    H (* i2 *)
 
 let _ =
-  if (* a3 *) B (* c3 *) then (* d3 *) E (* f3 *) else (* g3 *) H (* i3 *)
+  if (* a3 *) B (* c3 *) then (* d3 *)
+    E (* f3 *)
+  else (* g3 *)
+    H (* i3 *)
 ;;
 
 match x with

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

@@ -116,14 +116,16 @@ let foo =
     else some other action
   else some default action
 
-let foo = if cmp < 0 then (* foo *) a + b else (* foo *) a - b
+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 (* 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 +159,14 @@ 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

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

@@ -131,14 +131,16 @@ let foo =
     else some other action
   else some default action
 
-let foo = if cmp < 0 then (* foo *) a + b else (* foo *) a - b
+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 (* 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 +174,14 @@ 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

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

@@ -140,7 +140,11 @@ let foo =
   else
     some default action
 
-let foo = if cmp < 0 then (* foo *) a + b else (* foo *) a - b
+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 *)
@@ -191,3 +195,14 @@ let _ =
     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