Additional fixes on top of #2810

CHANGES.md

@@ -8,6 +8,18 @@ profile. This started with version 0.26.0.
 
 ### Fixed
 
+- Fix `begin match … end` (and `begin if … end`) branches: with
+  `if-then-else=fit-or-vertical` the `match … with` header no longer splits
+  over several lines and `end` is aligned with `begin`; and a leading comment
+  on the body no longer reindents it — `begin` keeps the body one indent in
+  instead of gluing the keyword to the comment.
+  (#2810, @MisterDA)
+
+- Fix indentation of a bare `match`/`function`/`try` branch preceded by a
+  comment with `if-then-else=fit-or-vertical`: the branch is no longer
+  indented relative to the comment's end column.
+  (#2810, @MisterDA)
+
 - Fix formatting oscillation with `if-then-else=fit-or-vertical` and
   `begin...end`.
   (#2800, @MisterDA)

lib/Fmt_ast.ml

@@ -2581,6 +2581,7 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                          in
                          let p =
                            Params.get_if_then_else c.conf ~cmts_before_opt
+                             ~has_cmts_before:(Cmts.has_before c.cmts)
                              ~pro:(fmt_if first pro_inner) ~first ~last
                              ~parens_bch ~parens_prev_bch:!parens_prev_bch
                              ~xcond ~xbch ~expr_loc:pexp_loc
@@ -2590,11 +2591,20 @@ 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 =
+                         let branch_pro, wrap_parens =
                            match raw_cmts_after_kw with
                            | Some cmts ->
-                               Params.raw_cmts_branch_pro c.conf cmts
-                           | None -> p.branch_pro
+                               let bp =
+                                 Params.raw_cmts_branch_pro c.conf cmts
+                               in
+                               if
+                                 parens_bch
+                                 && Params
+                                    .is_special_or_nested_special_beginend
+                                      xbch.ast.pexp_desc
+                               then (bp $ str "(", fun k -> k $ str ")")
+                               else (bp, p.wrap_parens)
+                           | None -> (p.branch_pro, p.wrap_parens)
                          in
                          let wrap_beginend =
                            match p.beginend_loc with
@@ -2607,7 +2617,7 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
                            $ p.box_keyword_and_expr
                                ( branch_pro
                                $ wrap_beginend
-                                   (p.wrap_parens
+                                   (wrap_parens
                                       ( fmt_expression c ?box:p.box_expr
                                           ~parens:false ?pro:p.expr_pro
                                           ?eol:p.expr_eol p.branch_expr
@@ -3034,11 +3044,23 @@ and fmt_beginend c ~loc ?(box = true) ?(pro = noop) ~ctx ~ctx0 ~fmt_atrs
   $
   match e.pexp_desc with
   | Pexp_match _ | Pexp_try _ | Pexp_function _ | Pexp_ifthenelse _ ->
-      beginend_box
-        (fmt_expression c
-           ~pro:(pro $ begin_ $ str " ")
-           ~box:false ?eol ~parens:false ~indent_wrap (sub_exp ~ctx e) )
-      $ end_
+      (* In an [if-then-else] branch the branch break indents [begin] one
+         level in (e.g. [fit-or-vertical], whose branch box is [hovbox 0]),
+         while [end] follows the branch box; wrap the body and [end] together
+         so [end] lines up with [begin]. Other contexts (e.g. [map x begin
+         fun … end] as an application argument) must keep their own
+         indentation. *)
+      let box =
+        match ctx0 with
+        | Exp {pexp_desc= Pexp_ifthenelse _; _} -> hvbox 0
+        | _ -> Fn.id
+      in
+      box
+        ( beginend_box
+            (fmt_expression c
+               ~pro:(pro $ begin_ $ str " ")
+               ~box:false ?eol ~parens:false ~indent_wrap (sub_exp ~ctx e) )
+        $ end_ )
   | _ ->
       beginend_box
         ( hvbox 0 (pro $ begin_)

lib/Params.ml

@@ -519,7 +519,7 @@ let is_special_or_nested_special_beginend = function
 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
+  | _ -> break 1000 2 $ cmts $ break 1000 2
 
 type cases =
   { leading_space: Fmt.t
@@ -858,15 +858,22 @@ type if_then_else =
   ; break_end_branch: Fmt.t
   ; space_between_branches: Fmt.t }
 
-let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
-    ~parens_bch ~parens_prev_bch ~xcond ~xbch ~expr_loc ~fmt_infix_ext_attrs
-    ~infix_ext_attrs ~fmt_cond ~cmts_before_kw ~cmts_after_kw =
+let get_if_then_else (c : Conf.t) ~cmts_before_opt ~has_cmts_before ~pro
+    ~first ~last ~parens_bch ~parens_prev_bch ~xcond ~xbch ~expr_loc
+    ~fmt_infix_ext_attrs ~infix_ext_attrs ~fmt_cond ~cmts_before_kw
+    ~cmts_after_kw =
   let imd = c.fmt_opts.indicate_multiline_delimiters.v in
   let beginend_loc, infix_ext_attrs_beginend, branch_expr =
     let ast = xbch.Ast.ast in
     match ast with
-    | {pexp_desc= Pexp_beginend ({pexp_desc; _}, _); _}
-      when is_special_beginend pexp_desc ->
+    | {pexp_desc= Pexp_beginend (nested_exp, _); _}
+      when is_special_beginend nested_exp.pexp_desc
+           && not (has_cmts_before nested_exp.pexp_loc) ->
+        (* [begin match/try/function/if … end] shortcut: keep [begin] glued
+           to the body keyword. A leading comment on the body breaks the
+           glue, so fall through to the plain [begin]/[end] machinery below,
+           which puts [begin] on its own line and the body (comment included)
+           one indent in. *)
         (None, None, xbch)
     | { pexp_desc= Pexp_beginend (nested_exp, infix_ext_attrs)
       ; pexp_attributes= []
@@ -878,6 +885,18 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
     | _ -> (None, None, xbch)
   in
   let has_beginend = Option.is_some beginend_loc in
+  (* A [begin]/[end] branch whose body is a [match]/[try]/[function]/[if]
+     (the [begin match … end] shortcut, or its plain form when a leading
+     comment de-glues [begin] from the body). Such a body has a header
+     ([match … with], [if … then]) that the branch [break_unless_newline
+     1000] would wrongly split, so it provides its own break after [begin]
+     instead. Simple-bodied [begin … end] keeps the regular branch break. *)
+  let is_special_beginend_branch =
+    match xbch.ast.pexp_desc with
+    | Pexp_beginend (nested_exp, _) ->
+        is_special_beginend nested_exp.pexp_desc
+    | _ -> false
+  in
   let wrap_parens ~wrap_breaks k =
     if has_beginend then
       let infix_ext_attrs_beginend =
@@ -929,31 +948,44 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
     | None when parens_bch -> str " "
     | None -> break 1 indent
   in
+  (* When a comment precedes a multi-line
+     [match]/[function]/[try]/[if-then-else] branch body (or a [begin]/[end]
+     wrapping one), emit it from [branch_pro] with forced breaks, so the
+     branch indentation stays stable regardless of where the comment was
+     initially attached. [default] is the comment-less [branch_pro] for the
+     current mode, [guard] an extra mode-specific applicability condition.
+     Returns [(branch_pro, has_cmts_before)] where [has_cmts_before] is true
+     when a comment was consumed into [branch_pro]. *)
+  let branch_pro_with_cmts ~default ~guard =
+    if
+      (not has_beginend)
+      && (not (Location.is_single_line expr_loc c.fmt_opts.margin.v))
+      && (not has_cmts_after_kw) && guard
+    then
+      match cmts_before_opt xbch.ast.pexp_loc with
+      | Some cmts -> (raw_cmts_branch_pro c cmts, true)
+      | 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 -> (raw_cmts_branch_pro c cmts, true)
+          | None -> (default, false) )
+        | _ -> (default, false) )
+    else (default, false)
+  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)
-          && is_special_or_nested_special_beginend 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 ()
+      let branch_pro, _has_cmts =
+        branch_pro_with_cmts ~default:(branch_pro ~indent:0 ())
+          ~guard:
+            ( (not parens_bch)
+            && is_special_or_nested_special_beginend xbch.ast.pexp_desc )
       in
       { box_branch= hovbox ~name:"Params.get_if_then_else `Compact" 2
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro= branch_pro_with_cmts
+      ; branch_pro
       ; wrap_parens=
           wrap_parens
             ~wrap_breaks:
@@ -973,7 +1005,7 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
       ; branch_pro= branch_pro ~begin_end_offset:0 ()
       ; wrap_parens= wrap_parens ~wrap_breaks:(wrap (break 1000 2) noop)
       ; beginend_loc
-      ; box_expr= Some has_beginend
+      ; box_expr= Some (has_beginend && not is_special_beginend_branch)
       ; expr_pro= None
       ; expr_eol= Some (break 1 2)
       ; branch_expr
@@ -982,23 +1014,29 @@ 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 ()
+      let branch_pro_default, has_cmts =
+        branch_pro_with_cmts
+          ~default:(branch_pro ~begin_end_offset:0 ())
+          ~guard:true
+      in
+      let paren_glue =
+        has_cmts && parens_bch
+        && is_special_or_nested_special_beginend xbch.ast.pexp_desc
+      in
+      let branch_pro, wrap_parens =
+        if paren_glue then
+          let cls =
+            match imd with
+            | `Closing_on_separate_line -> break 1000 0 $ str ")"
+            | _ -> str ")"
+          in
+          (branch_pro_default $ str "(", fun k -> k $ cls)
+        else
+          ( branch_pro_default
+          , wrap_parens
+              ~wrap_breaks:
+                (get_parens_breaks ~opn_hint_indent:2
+                   ~cls_hint:((1, 0), (1000, 0)) ) )
       in
       { box_branch=
           hovbox
@@ -1007,19 +1045,18 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
             | _ -> 0 )
       ; cond= cond ()
       ; box_keyword_and_expr= Fn.id
-      ; branch_pro= branch_pro_with_cmts
-      ; wrap_parens=
-          wrap_parens
-            ~wrap_breaks:
-              (get_parens_breaks ~opn_hint_indent:2
-                 ~cls_hint:((1, 0), (1000, 0)) )
+      ; branch_pro
+      ; wrap_parens
       ; beginend_loc
       ; box_expr= Some false
       ; expr_pro=
-          Some
-            (fmt_if
-               (not (Location.is_single_line expr_loc c.fmt_opts.margin.v))
-               (break_unless_newline 1000 2) )
+          ( if is_special_beginend_branch || paren_glue then None
+            else
+              Some
+                (fmt_if
+                   (not
+                      (Location.is_single_line expr_loc c.fmt_opts.margin.v) )
+                   (break_unless_newline 1000 2) ) )
       ; expr_eol= Some (break 1 2)
       ; branch_expr
       ; break_end_branch= noop
@@ -1040,7 +1077,9 @@ let get_if_then_else (c : Conf.t) ~cmts_before_opt ~pro ~first ~last
                  ~cls_hint:((1, 0), (1000, 0)) )
       ; beginend_loc
       ; box_expr= None
-      ; expr_pro= Some (break_unless_newline 1000 2)
+      ; expr_pro=
+          ( if is_special_beginend_branch then None
+            else Some (break_unless_newline 1000 2) )
       ; expr_eol= None
       ; branch_expr
       ; break_end_branch= noop

lib/Params.mli

@@ -221,6 +221,7 @@ type if_then_else =
 val get_if_then_else :
      Conf.t
   -> cmts_before_opt:(Location.t -> Fmt.t option)
+  -> has_cmts_before:(Location.t -> bool)
   -> pro:Fmt.t
   -> first:bool
   -> last:bool

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

@@ -207,10 +207,10 @@ let test =
 let f =
   match x with
   | A ->
-    if gf then
+    if gf then begin
       (* aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa *)
-      begin if a then b
-      end
+      if a then b
+    end
 
 (* Long comment before match in else branch (regression test for
    formatting oscillation) *)
@@ -230,3 +230,35 @@ let _ =
         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
+
+(* Short comment before a bare match in the else branch: the match must stay
+   indented under the comment, not pushed to its end column (regression test) *)
+let _ =
+  if cond then x
+  else (
+    (* a comment *)
+    match e with
+    | A -> a
+    | B -> b)
+
+(* Comment before [begin match ... end]: comment must not glue to [begin]
+   (regression test for paren-glue in profiles with exp_grouping=parens) *)
+let _ =
+  if cond then x
+  else
+    (* a comment *)
+    begin match some_long_scrutinee_expression with
+    | A -> a
+    | B -> b
+    end
+
+(* A bare [begin match ... end] branch must keep [match ... with] on one line,
+   not split it as [begin match] / scrutinee / [with], and must align [end]
+   with [begin] (regression test) *)
+let _ =
+  if cond then x
+  else
+    begin match e with
+    | A -> function_with_a_long_enough_name_to_force_a_vertical_break a
+    | B -> b
+    end