Fix formatting oscillation with if-then-else=fit-or-vertical and begin...end

When `if-then-else=fit-or-vertical` formats a branch containing
`begin (* long comment *) if ... then ... end`, comments between the
`then` keyword and `begin` would get attached to different AST
locations on each formatting pass ("after" the keyword vs "before"
the begin...end expression). This caused `branch_pro` to compute
different indentation each time, producing an infinite oscillation
that triggered "BUG: formatting did not stabilize after 10 iterations".

The fix has two parts:

1. In Fmt_ast, when the if-then-else branch is a special begin...end
   (wrapping match/try/function/ifthenelse), relocate "after keyword"
   comments to "before branch" before formatting. This ensures
   consistent comment placement regardless of source layout.

2. In Params, when computing branch_pro for Fit_or_vertical, also
   check for comments before the inner expression of a special
   begin...end, not just before the begin...end node itself.

A new helper `Cmts.relocate_after_to_before` moves comments from
the "after" position of one location to the "before" position of
another.

Author: MisterDA

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`.
+  (#2780, @MisterDA)
+
 - Fix instability on long `if-then-else` with `if-then-else=fit-or-vertical`
   (#2797, @MisterDA)
 

lib/Cmts.ml

@@ -328,6 +328,20 @@ let relocate_cmts_before (t : t) ~src ~sep ~dst =
   in
   update_cmts t `Before ~f ; update_cmts t `Within ~f
 
+let relocate_after_to_before (t : t) ~src ~dst =
+  let merge_and_sort x y =
+    List.rev_append x y
+    |> List.sort ~compare:(Comparable.lift Location.compare_start ~f:Cmt.loc)
+  in
+  match Map.find t.cmts_after src with
+  | None -> ()
+  | Some cmts ->
+      update_cmts t `After ~f:(fun map -> Map.remove map src) ;
+      update_cmts t `Before ~f:(fun map ->
+          Map.update map dst ~f:(function
+            | None -> cmts
+            | Some existing -> merge_and_sort cmts existing ) )
+
 let relocate_pattern_matching_cmts (t : t) src tok ~whole_loc ~matched_loc =
   if not (whole_loc.Location.loc_ghost || matched_loc.Location.loc_ghost)
   then

lib/Cmts.mli

@@ -44,6 +44,10 @@ val relocate_wrongfully_attached_cmts :
     pattern-matching expressions ([match-with] or [try-with] expressions) but
     are wrongfully attached to the matched expression. *)
 
+val relocate_after_to_before : t -> src:Location.t -> dst:Location.t -> unit
+(** [relocate_after_to_before t ~src ~dst] moves comments that are after
+    [src] to be before [dst]. *)
+
 val fmt_before :
      t
   -> Conf.t

lib/Fmt_ast.ml

@@ -2556,6 +2556,22 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                          let parens_bch =
                            parenze_exp xbch && not symbol_parens
                          in
+                         (* When the branch is begin...end wrapping a special
+                            expression (match/try/function/if), comments
+                            between the keyword and begin can end up attached
+                            to different locations depending on source
+                            layout. Relocate them to a consistent position to
+                            avoid oscillation. *)
+                         ( match xbch.ast.pexp_desc with
+                         | Pexp_beginend ({pexp_desc; _}, _)
+                           when match pexp_desc with
+                                | Pexp_match _ | Pexp_try _
+                                 |Pexp_function _ | Pexp_ifthenelse _ ->
+                                    true
+                                | _ -> false ->
+                             Cmts.relocate_after_to_before c.cmts
+                               ~src:keyword_loc ~dst:xbch.ast.pexp_loc
+                         | _ -> () ) ;
                          let cmts_before_kw =
                            Cmts.fmt_before c keyword_loc
                          in

lib/Params.ml

@@ -968,7 +968,14 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
             then
               match cmts_before_opt xbch.ast.pexp_loc with
               | Some cmts -> break 1000 2 $ cmts
-              | None -> branch_pro ~begin_end_offset:0 ()
+              | None -> (
+                match xbch.ast.pexp_desc with
+                | Pexp_beginend (nested_exp, _)
+                  when is_special_beginend nested_exp.pexp_desc -> (
+                  match cmts_before_opt nested_exp.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 () )
       ; wrap_parens=
           wrap_parens

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

@@ -1,3 +1,4 @@
 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:199 exceeds the margin

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

@@ -201,3 +201,13 @@ 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

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

@@ -0,0 +1 @@
+Warning: ite-compact_closing.ml:216 exceeds the margin

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

@@ -220,3 +220,13 @@ 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

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

@@ -1,3 +1,4 @@
 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:236 exceeds the margin