Fix formatting oscillation in if-then-else branches

CHANGES.md

@@ -8,6 +8,10 @@ profile. This started with version 0.26.0.
 
 ### Fixed
 
+- Fix formatting oscillation with `if-then-else=fit-or-vertical` and
+  `begin...end`.
+  (#2800, @MisterDA)
+
 - Fix instability on long `if-then-else` with `if-then-else=fit-or-vertical`
   (#2797, @MisterDA)
 

lib/Fmt_ast.ml

@@ -2559,10 +2559,18 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                          let cmts_before_kw =
                            Cmts.fmt_before c keyword_loc
                          in
-                         let cmts_after_kw =
+                         let cmts_after_kw, raw_cmts_after_kw =
                            if Cmts.has_after c.cmts keyword_loc then
-                             Some (Cmts.fmt_after c keyword_loc)
-                           else None
+                             if
+                               Params.needs_raw_cmts_after_kw
+                                 xbch.ast.pexp_desc
+                             then
+                               ( None
+                               , Some
+                                   (Cmts.fmt_after ~pro:noop ~epi:noop c
+                                      keyword_loc ) )
+                             else (Some (Cmts.fmt_after c keyword_loc), None)
+                           else (None, None)
                          in
                          let cmts_before_opt loc =
                            if Cmts.has_before c.cmts loc then
@@ -2582,6 +2590,12 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                              ~fmt_cond:(fmt_expression ~box:false c)
                              ~cmts_before_kw ~cmts_after_kw
                          in
+                         let branch_pro =
+                           match raw_cmts_after_kw with
+                           | Some cmts ->
+                               Params.raw_cmts_branch_pro c.conf cmts
+                           | None -> p.branch_pro
+                         in
                          let wrap_beginend =
                            match p.beginend_loc with
                            | Some loc -> Cmts.fmt c loc
@@ -2591,7 +2605,7 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                          p.box_branch
                            ( p.cond
                            $ p.box_keyword_and_expr
-                               ( p.branch_pro
+                               ( branch_pro
                                $ wrap_beginend
                                    (p.wrap_parens
                                       ( fmt_expression c ?box:p.box_expr

lib/Params.ml

@@ -512,6 +512,21 @@ let is_special_beginend exp =
   | Pexp_match _ | Pexp_try _ | Pexp_function _ | Pexp_ifthenelse _ -> true
   | _ -> false
 
+let needs_raw_cmts_after_kw = function
+  | Pexp_beginend
+      ( { pexp_desc=
+            Pexp_match _ | Pexp_try _ | Pexp_function _ | Pexp_ifthenelse _
+        ; _ }
+      , _ )
+   |Pexp_match _ | Pexp_try _ | Pexp_function _ | Pexp_ifthenelse _ ->
+      true
+  | _ -> false
+
+let raw_cmts_branch_pro (c : Conf.t) cmts =
+  match c.fmt_opts.if_then_else.v with
+  | `Compact -> break 1000 0 $ cmts $ break 1000 0
+  | _ -> break 1000 2 $ cmts
+
 type cases =
   { leading_space: Fmt.t
   ; bar: Fmt.t
@@ -922,10 +937,29 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
   in
   match c.fmt_opts.if_then_else.v with
   | `Compact ->
+      let branch_pro_with_cmts =
+        if
+          (not has_beginend) && (not parens_bch)
+          && (not (Location.is_single_line expr_loc c.fmt_opts.margin.v))
+          && (not has_cmts_after_kw)
+          && needs_raw_cmts_after_kw xbch.ast.pexp_desc
+        then
+          match cmts_before_opt xbch.ast.pexp_loc with
+          | Some cmts -> break 1000 0 $ cmts $ break 1000 0
+          | None -> (
+            match xbch.ast.pexp_desc with
+            | Pexp_beginend ({pexp_loc; pexp_desc; _}, _)
+              when is_special_beginend pexp_desc -> (
+              match cmts_before_opt pexp_loc with
+              | Some cmts -> break 1000 0 $ cmts $ break 1000 0
+              | None -> branch_pro ~indent:0 () )
+            | _ -> branch_pro ~indent:0 () )
+        else branch_pro ~indent:0 ()
+      in
       { box_branch= hovbox ~name:"Params.get_if_then_else `Compact" 2
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro= branch_pro ~indent:0 ()
+      ; branch_pro= branch_pro_with_cmts
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:
@@ -954,22 +988,32 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
       ; space_between_branches= fmt_if (has_beginend || parens_bch) (str " ")
       }
   | `Fit_or_vertical ->
+      let branch_pro_with_cmts =
+        if
+          (not has_beginend)
+          && (not (Location.is_single_line expr_loc c.fmt_opts.margin.v))
+          && not has_cmts_after_kw
+        then
+          match cmts_before_opt xbch.ast.pexp_loc with
+          | Some cmts -> break 1000 2 $ cmts
+          | None -> (
+            match xbch.ast.pexp_desc with
+            | Pexp_beginend ({pexp_loc; pexp_desc; _}, _)
+              when is_special_beginend pexp_desc -> (
+              match cmts_before_opt pexp_loc with
+              | Some cmts -> break 1000 2 $ cmts
+              | None -> branch_pro ~begin_end_offset:0 () )
+            | _ -> branch_pro ~begin_end_offset:0 () )
+        else branch_pro ~begin_end_offset:0 ()
+      in
       { box_branch=
           hovbox
             ( match imd with
             | `Closing_on_separate_line when parens_prev_bch -> -2
             | _ -> 0 )
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro=
-          ( if
-              (not has_beginend) && (not has_cmts_after_kw)
-              && not (Location.is_single_line expr_loc c.fmt_opts.margin.v)
-            then
-              match cmts_before_opt xbch.ast.pexp_loc with
-              | Some cmts -> break 1000 2 $ cmts
-              | None -> branch_pro ~begin_end_offset:0 ()
-            else branch_pro ~begin_end_offset:0 () )
+      ; branch_pro= branch_pro_with_cmts
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:

lib/Params.mli

@@ -237,6 +237,16 @@ val get_if_then_else :
   -> if_then_else
 (** [cmts_before_opt] return the comment before the given location with no breaks around it. *)
 
+val needs_raw_cmts_after_kw : expression_desc -> bool
+(** [needs_raw_cmts_after_kw desc] returns [true] when comments after the
+    keyword should be extracted without breaks (raw) to prevent oscillation
+    between "after keyword" and "before expression" placements. *)
+
+val raw_cmts_branch_pro : Conf.t -> Fmt.t -> Fmt.t
+(** [raw_cmts_branch_pro c cmts] returns the branch_pro for raw comments
+    extracted after a keyword, using the correct indentation for the current
+    if-then-else mode. *)
+
 val match_indent : ?default:int -> Conf.t -> parens:bool -> ctx:Ast.t -> int
 (** [match_indent c ~ctx ~default] returns the indentation used for the
     pattern-matching in context [ctx], depending on the `match-indent-nested`

test/passing/refs.ahrefs/ite-compact.ml.err

@@ -1,3 +1,5 @@
 Warning: ite-compact.ml:95 exceeds the margin
 Warning: ite-compact.ml:100 exceeds the margin
 Warning: ite-compact.ml:105 exceeds the margin
+Warning: ite-compact.ml:210 exceeds the margin
+Warning: ite-compact.ml:219 exceeds the margin

test/passing/refs.ahrefs/ite-compact.ml.ref

@@ -201,3 +201,32 @@ let test =
     (* comment *)
     1
   else 2
+
+(* Begin-end with long comment before nested if (regression test for
+   formatting oscillation) *)
+let f =
+  match x with
+  | A ->
+    if gf then
+      (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
+      begin if a then b
+      end
+
+(* Long comment before match in else branch (regression test for
+   formatting oscillation) *)
+let _ =
+  if a then b
+  else (
+    (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
+    match x with
+    | A -> f)
+
+(* Comment before nested if in then branch - Compact mode oscillation *)
+let _ =
+  if aaaa then
+    if bbbbb then
+      (* in this case, b <- s *)
+      if abs_float fa < abs_float fs then
+        iter (i + 1) ~a:s ~b:a ~c:b ~d:c ~fa:fs ~fb:fa ~fc:fb mflag
+      else iter (i + 1) ~a ~b:s ~c:b ~d:c ~fa ~fb:fs ~fc:fb mflag
+    else g

test/passing/refs.ahrefs/ite-compact_closing.ml.err

@@ -0,0 +1,2 @@
+Warning: ite-compact_closing.ml:229 exceeds the margin
+Warning: ite-compact_closing.ml:238 exceeds the margin

test/passing/refs.ahrefs/ite-compact_closing.ml.ref

@@ -220,3 +220,33 @@ let test =
     (* comment *)
     1
   else 2
+
+(* Begin-end with long comment before nested if (regression test for
+   formatting oscillation) *)
+let f =
+  match x with
+  | A ->
+    if gf then
+      (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
+      begin if a then b
+      end
+
+(* Long comment before match in else branch (regression test for
+   formatting oscillation) *)
+let _ =
+  if a then b
+  else (
+    (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
+    match x with
+    | A -> f
+  )
+
+(* Comment before nested if in then branch - Compact mode oscillation *)
+let _ =
+  if aaaa then
+    if bbbbb then
+      (* in this case, b <- s *)
+      if abs_float fa < abs_float fs then
+        iter (i + 1) ~a:s ~b:a ~c:b ~d:c ~fa:fs ~fb:fa ~fc:fb mflag
+      else iter (i + 1) ~a ~b:s ~c:b ~d:c ~fa ~fb:fs ~fc:fb mflag
+    else g

test/passing/refs.ahrefs/ite-fit_or_vertical.ml.err

@@ -1,3 +1,5 @@
 Warning: ite-fit_or_vertical.ml:116 exceeds the margin
 Warning: ite-fit_or_vertical.ml:121 exceeds the margin
 Warning: ite-fit_or_vertical.ml:126 exceeds the margin
+Warning: ite-fit_or_vertical.ml:247 exceeds the margin
+Warning: ite-fit_or_vertical.ml:260 exceeds the margin

test/passing/refs.ahrefs/ite-fit_or_vertical.ml.ref

@@ -238,3 +238,38 @@ let test =
     1
   else
     2
+
+(* Begin-end with long comment before nested if (regression test for
+   formatting oscillation) *)
+let f =
+  match x with
+  | A ->
+    if gf then
+      (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
+    begin if
+      a
+    then
+      b
+    end
+
+(* Long comment before match in else branch (regression test for
+   formatting oscillation) *)
+let _ =
+  if a then
+    b
+  else
+    (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)(
+    match x with
+    | A -> f)
+
+(* Comment before nested if in then branch - Compact mode oscillation *)
+let _ =
+  if aaaa then
+    if bbbbb then
+      (* in this case, b <- s *)
+      if abs_float fa < abs_float fs then
+        iter (i + 1) ~a:s ~b:a ~c:b ~d:c ~fa:fs ~fb:fa ~fc:fb mflag
+      else
+        iter (i + 1) ~a ~b:s ~c:b ~d:c ~fa ~fb:fs ~fc:fb mflag
+    else
+      g

test/passing/refs.ahrefs/ite-fit_or_vertical_closing.ml.err

@@ -0,0 +1,2 @@
+Warning: ite-fit_or_vertical_closing.ml:257 exceeds the margin
+Warning: ite-fit_or_vertical_closing.ml:270 exceeds the margin

test/passing/refs.ahrefs/ite-fit_or_vertical_closing.ml.ref

@@ -248,3 +248,39 @@ let test =
     1
   else
     2
+
+(* Begin-end with long comment before nested if (regression test for
+   formatting oscillation) *)
+let f =
+  match x with
+  | A ->
+    if gf then
+      (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
+    begin if
+      a
+    then
+      b
+    end
+
+(* Long comment before match in else branch (regression test for
+   formatting oscillation) *)
+let _ =
+  if a then
+    b
+  else
+    (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)(
+    match x with
+    | A -> f
+  )
+
+(* Comment before nested if in then branch - Compact mode oscillation *)
+let _ =
+  if aaaa then
+    if bbbbb then
+      (* in this case, b <- s *)
+      if abs_float fa < abs_float fs then
+        iter (i + 1) ~a:s ~b:a ~c:b ~d:c ~fa:fs ~fb:fa ~fc:fb mflag
+      else
+        iter (i + 1) ~a ~b:s ~c:b ~d:c ~fa ~fb:fs ~fc:fb mflag
+    else
+      g

test/passing/refs.ahrefs/ite-fit_or_vertical_no_indicate.ml.err

@@ -1,3 +1,5 @@
 Warning: ite-fit_or_vertical_no_indicate.ml:116 exceeds the margin
 Warning: ite-fit_or_vertical_no_indicate.ml:121 exceeds the margin
 Warning: ite-fit_or_vertical_no_indicate.ml:126 exceeds the margin
+Warning: ite-fit_or_vertical_no_indicate.ml:247 exceeds the margin
+Warning: ite-fit_or_vertical_no_indicate.ml:260 exceeds the margin

test/passing/refs.ahrefs/ite-fit_or_vertical_no_indicate.ml.ref

@@ -238,3 +238,38 @@ let test =
     1
   else
     2
+
+(* Begin-end with long comment before nested if (regression test for
+   formatting oscillation) *)
+let f =
+  match x with
+  | A ->
+    if gf then
+      (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
+    begin if
+      a
+    then
+      b
+    end
+
+(* Long comment before match in else branch (regression test for
+   formatting oscillation) *)
+let _ =
+  if a then
+    b
+  else
+    (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)(
+    match x with
+    | A -> f)
+
+(* Comment before nested if in then branch - Compact mode oscillation *)
+let _ =
+  if aaaa then
+    if bbbbb then
+      (* in this case, b <- s *)
+      if abs_float fa < abs_float fs then
+        iter (i + 1) ~a:s ~b:a ~c:b ~d:c ~fa:fs ~fb:fa ~fc:fb mflag
+      else
+        iter (i + 1) ~a ~b:s ~c:b ~d:c ~fa ~fb:fs ~fc:fb mflag
+    else
+      g

test/passing/refs.ahrefs/ite-kr.ml.err

@@ -0,0 +1,2 @@
+Warning: ite-kr.ml:286 exceeds the margin
+Warning: ite-kr.ml:296 exceeds the margin