Fix dropped attributes on begin-end in an if-then-else branch (#2436)

Author: gpetiot

CHANGES.md

@@ -29,6 +29,7 @@ Tags:
 ### Fixed
 
 - Fix dropped attributes on a begin-end in a match case (#2421, @Julow)
+- Fix dropped attributes on begin-end in an if-then-else branch (#2436, @gpetiot)
 - Fix non-stabilizing comments before a functor type argument (#2420, @Julow)
 - Fix crash caused by module types with nested `with module` (#2419, @Julow)
 - Fix ';;' formatting between doc-comments and toplevel directives (#2432, @gpetiot)

lib/Fmt_ast.ml

@@ -2232,7 +2232,7 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
                           $ p.wrap_parens
                               ( fmt_expression c ?box:p.box_expr
                                   ~parens:parens_exp ?pro:p.expr_pro
-                                  ?eol:p.expr_eol xbch
+                                  ?eol:p.expr_eol p.branch_expr
                               $ p.break_end_branch ) ) )
                     $ fmt_if_k (not last) p.space_between_branches ) ) )
         $ fmt_atrs )
@@ -2655,24 +2655,11 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
       impossible "only used for methods, handled during method formatting"
   | Pexp_hole -> hvbox 0 (fmt_hole () $ fmt_atrs)
   | Pexp_beginend e ->
-      let wrap_beginend =
-        match ctx0 with
-        (* begin-end keywords are handled when printing if-then-else
-           branch *)
-        | Exp {pexp_desc= Pexp_ifthenelse (_, Some z); _}
-          when Base.phys_equal xexp.ast z ->
-            Fn.id
-        | Exp {pexp_desc= Pexp_ifthenelse (eN, _); _}
-          when List.exists eN ~f:(fun x ->
-                   Base.phys_equal xexp.ast x.if_body ) ->
-            Fn.id
-        | _ ->
-            fun k ->
-              let opn = str "begin" $ fmt_extension_suffix c ext
-              and cls = str "end" in
-              hvbox 0
-                ( wrap_k opn cls (wrap_k (break 1 2) (break 1000 0) k)
-                $ fmt_atrs )
+      let wrap_beginend k =
+        let opn = str "begin" $ fmt_extension_suffix c ext
+        and cls = str "end" in
+        hvbox 0
+          (wrap_k opn cls (wrap_k (break 1 2) (break 1000 0) k) $ fmt_atrs)
       in
       wrap_beginend
       @@ fmt_expression c ~box ?pro ?epi ?eol ~parens:false ~indent_wrap ?ext

lib/Params.ml

@@ -462,16 +462,19 @@ type if_then_else =
   ; box_expr: bool option
   ; expr_pro: Fmt.t option
   ; expr_eol: Fmt.t option
+  ; branch_expr: expression Ast.xt
   ; break_end_branch: Fmt.t
   ; 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 =
   let imd = c.fmt_opts.indicate_multiline_delimiters.v in
-  let beginend =
-    match xbch.Ast.ast with
-    | {pexp_desc= Pexp_beginend _; _} -> true
-    | _ -> false
+  let beginend, branch_expr =
+    let ast = xbch.Ast.ast in
+    match ast with
+    | {pexp_desc= Pexp_beginend nested_exp; pexp_attributes= []; _} ->
+        (true, sub_exp ~ctx:(Exp ast) nested_exp)
+    | _ -> (false, xbch)
   in
   let wrap_parens ~wrap_breaks k =
     if beginend then wrap "begin" "end" (wrap_breaks k)
@@ -517,6 +520,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
       ; box_expr= Some false
       ; expr_pro= None
       ; expr_eol= None
+      ; branch_expr
       ; break_end_branch= noop
       ; space_between_branches= fmt "@ " }
   | `K_R ->
@@ -528,6 +532,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
       ; box_expr= Some false
       ; expr_pro= None
       ; expr_eol= Some (fmt "@;<1 2>")
+      ; branch_expr
       ; break_end_branch=
           fmt_if_k (parens_bch || beginend || not last) (break 1000 0)
       ; space_between_branches= fmt_if (beginend || parens_bch) " " }
@@ -552,6 +557,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
                (not (Location.is_single_line expr_loc c.fmt_opts.margin.v))
                (break_unless_newline 1000 2) )
       ; expr_eol= Some (fmt "@;<1 2>")
+      ; branch_expr
       ; break_end_branch= noop
       ; space_between_branches=
           fmt
@@ -571,6 +577,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
       ; box_expr= None
       ; expr_pro= Some (break_unless_newline 1000 2)
       ; expr_eol= None
+      ; branch_expr
       ; break_end_branch= noop
       ; space_between_branches=
           fmt
@@ -600,6 +607,7 @@ let get_if_then_else (c : Conf.t) ~first ~last ~parens_bch ~parens_prev_bch
       ; box_expr= Some false
       ; expr_pro= None
       ; expr_eol= None
+      ; branch_expr
       ; break_end_branch= noop
       ; space_between_branches= fmt "@ " }
 

lib/Params.mli

@@ -139,6 +139,7 @@ type if_then_else =
   ; box_expr: bool option
   ; expr_pro: Fmt.t option
   ; expr_eol: Fmt.t option
+  ; branch_expr: expression Ast.xt
   ; break_end_branch: Fmt.t
   ; space_between_branches: Fmt.t }
 

test/passing/tests/exp_grouping-parens.ml.ref

@@ -328,3 +328,10 @@ let x =
       | pd -> pd
   in
   ()
+
+let _ =
+  if something_changed then
+    begin
+      loop
+    end
+    [@attr]

test/passing/tests/exp_grouping.ml

@@ -267,3 +267,8 @@ let x =
       | pd -> pd
   in
   ()
+
+let _ =
+  if something_changed then begin[@attr]
+    loop
+  end

test/passing/tests/exp_grouping.ml.ref

@@ -382,3 +382,10 @@ let x =
       | pd -> pd
   in
   ()
+
+let _ =
+  if something_changed then
+    begin
+      loop
+    end
+    [@attr]