Fix dropped attributes on begin-end in a match case (ocaml-ppx#2421)

The special formatting of a begin-end in a match case was dropping
attributes.

Author: Julow

CHANGES.md

@@ -8,6 +8,7 @@ profile. This started with version 0.26.0.
 
 ### Bug fixes
 
+- Fix dropped attributes on a begin-end in a match case (#2421, @Julow)
 - Fix non-stabilizing comments before a functor type argument (#2420, @Julow)
 - Fix crash caused by module types with nested `with module` (#2419, @Julow)
 

lib/Fmt_ast.ml

@@ -2641,14 +2641,6 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
           when List.exists eN ~f:(fun x ->
                    Base.phys_equal xexp.ast x.if_body ) ->
             Fn.id
-        (* begin-end keywords are handled when printing pattern-matching
-           cases *)
-        | Exp
-            { pexp_desc=
-                Pexp_function xs | Pexp_match (_, xs) | Pexp_try (_, xs)
-            ; _ }
-          when List.exists xs ~f:(fun x -> Poly.(x.pc_rhs = exp)) ->
-            Fn.id
         | _ ->
             fun k ->
               let opn = str "begin" $ fmt_extension_suffix c ext
@@ -3058,7 +3050,7 @@ and fmt_case c ctx ~first ~last case =
           $ p.open_paren_branch )
       $ p.break_after_opening_paren
       $ hovbox 0
-          ( fmt_expression ?eol c ?parens:p.expr_parens xrhs
+          ( fmt_expression ?eol c ?parens:p.expr_parens p.branch_expr
           $ p.close_paren_branch ) )
 
 and fmt_value_description ?ext c ctx vd =

lib/Params.ml

@@ -135,6 +135,7 @@ type cases =
   ; open_paren_branch: Fmt.t
   ; break_after_opening_paren: Fmt.t
   ; expr_parens: bool option
+  ; branch_expr: expression Ast.xt
   ; close_paren_branch: Fmt.t }
 
 let get_cases (c : Conf.t) ~ctx ~first ~last ~xbch:({ast; _} as xast) =
@@ -165,24 +166,25 @@ let get_cases (c : Conf.t) ~ctx ~first ~last ~xbch:({ast; _} as xast) =
     else (parenze_exp xast && not body_has_parens, Some false)
   in
   let indent = if align_nested_match then 0 else indent in
-  let beginend =
-    match ast with {pexp_desc= Pexp_beginend _; _} -> true | _ -> false
-  in
-  let open_paren_branch =
-    if beginend then fmt "@;<1 0>begin" else fmt_if parens_branch " ("
-  in
-  let close_paren_branch =
-    if beginend then
-      let offset =
-        match c.fmt_opts.break_cases.v with `Nested -> 0 | _ -> -2
-      in
-      fits_breaks " end" ~level:1 ~hint:(1000, offset) "end"
-    else
-      fmt_if_k parens_branch
-        ( match c.fmt_opts.indicate_multiline_delimiters.v with
-        | `Space -> fmt "@ )"
-        | `No -> fmt "@,)"
-        | `Closing_on_separate_line -> fmt "@;<1000 -2>)" )
+  let open_paren_branch, close_paren_branch, branch_expr =
+    match ast with
+    | {pexp_desc= Pexp_beginend nested_exp; pexp_attributes= []; _} ->
+        let close_paren =
+          let offset =
+            match c.fmt_opts.break_cases.v with `Nested -> 0 | _ -> -2
+          in
+          fits_breaks " end" ~level:1 ~hint:(1000, offset) "end"
+        in
+        (fmt "@;<1 0>begin", close_paren, sub_exp ~ctx:(Exp ast) nested_exp)
+    | _ ->
+        let close_paren =
+          fmt_if_k parens_branch
+            ( match c.fmt_opts.indicate_multiline_delimiters.v with
+            | `Space -> fmt "@ )"
+            | `No -> fmt "@,)"
+            | `Closing_on_separate_line -> fmt "@;<1000 -2>)" )
+        in
+        (fmt_if parens_branch " (", close_paren, xast)
   in
   match c.fmt_opts.break_cases.v with
   | `Fit ->
@@ -195,6 +197,7 @@ let get_cases (c : Conf.t) ~ctx ~first ~last ~xbch:({ast; _} as xast) =
       ; open_paren_branch
       ; break_after_opening_paren= fmt "@ "
       ; expr_parens
+      ; branch_expr
       ; close_paren_branch }
   | `Nested ->
       { leading_space= fmt_if (not first) "@ "
@@ -206,6 +209,7 @@ let get_cases (c : Conf.t) ~ctx ~first ~last ~xbch:({ast; _} as xast) =
       ; open_paren_branch
       ; break_after_opening_paren= fmt_or (indent > 2) "@;<1 4>" "@;<1 2>"
       ; expr_parens
+      ; branch_expr
       ; close_paren_branch }
   | `Fit_or_vertical ->
       { leading_space= break_unless_newline 1000 0
@@ -217,6 +221,7 @@ let get_cases (c : Conf.t) ~ctx ~first ~last ~xbch:({ast; _} as xast) =
       ; open_paren_branch
       ; break_after_opening_paren= fmt "@ "
       ; expr_parens
+      ; branch_expr
       ; close_paren_branch }
   | `Toplevel | `All ->
       { leading_space= break_unless_newline 1000 0
@@ -228,6 +233,7 @@ let get_cases (c : Conf.t) ~ctx ~first ~last ~xbch:({ast; _} as xast) =
       ; open_paren_branch
       ; break_after_opening_paren= fmt "@ "
       ; expr_parens
+      ; branch_expr
       ; close_paren_branch }
   | `Vertical ->
       { leading_space= break_unless_newline 1000 0
@@ -239,6 +245,7 @@ let get_cases (c : Conf.t) ~ctx ~first ~last ~xbch:({ast; _} as xast) =
       ; open_paren_branch
       ; break_after_opening_paren= break 1000 0
       ; expr_parens
+      ; branch_expr
       ; close_paren_branch }
 
 let wrap_collec c ~space_around opn cls =

lib/Params.mli

@@ -61,6 +61,7 @@ type cases =
   ; open_paren_branch: Fmt.t
   ; break_after_opening_paren: Fmt.t
   ; expr_parens: bool option
+  ; branch_expr: expression Ast.xt  (** Expression on the RHS of the [->]. *)
   ; close_paren_branch: Fmt.t }
 
 val get_cases :

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

@@ -307,3 +307,24 @@ let _ =
   [@landmark "parse_constant_dividends"]
 
 let () = if a then b (* asd *)
+
+let x =
+  let get_path_and_distance pv1 pv2 =
+    if is_loop pv1 pv2 then
+      Some ([], 0)
+    else
+      match Tbl.find dist_tbl (pv1, pv2) with
+      | None ->
+        (* FIXME: temporary hack to avoid Jane Street's annoying warnings. *)
+        begin
+          try
+            let path', dist = Dijkstra.shortest_path pgraph pv1 pv2 in
+            let path = unwrap_path path' in
+            Tbl.set dist_tbl ~key:(pv1, pv2) ~data:(path, dist) ;
+            Some (path, dist)
+          with Not_found | Not_found_s _ -> None
+        end
+        [@warning "-3"]
+      | pd -> pd
+  in
+  ()

test/passing/tests/exp_grouping.ml

@@ -249,3 +249,21 @@ let () =
   if a then begin b
     (* asd *)
   end
+
+let x =
+  let get_path_and_distance pv1 pv2 =
+    if is_loop pv1 pv2 then Some ([],0) else
+      match Tbl.find dist_tbl (pv1, pv2) with
+      | None ->
+        (* FIXME: temporary hack to avoid Jane Street's annoying warnings. *)
+        begin[@warning "-3"] try
+          let path', dist = Dijkstra.shortest_path pgraph pv1 pv2 in
+          let path = unwrap_path path' in
+          Tbl.set dist_tbl ~key:(pv1, pv2) ~data:(path, dist);
+          Some (path, dist)
+        with Not_found | Not_found_s _ ->
+          None
+        end
+      | pd -> pd
+  in
+  ()

test/passing/tests/exp_grouping.ml.ref

@@ -361,3 +361,24 @@ let () =
     b
   (* asd *)
   end
+
+let x =
+  let get_path_and_distance pv1 pv2 =
+    if is_loop pv1 pv2 then
+      Some ([], 0)
+    else
+      match Tbl.find dist_tbl (pv1, pv2) with
+      | None ->
+        (* FIXME: temporary hack to avoid Jane Street's annoying warnings. *)
+        begin
+          try
+            let path', dist = Dijkstra.shortest_path pgraph pv1 pv2 in
+            let path = unwrap_path path' in
+            Tbl.set dist_tbl ~key:(pv1, pv2) ~data:(path, dist) ;
+            Some (path, dist)
+          with Not_found | Not_found_s _ -> None
+        end
+        [@warning "-3"]
+      | pd -> pd
+  in
+  ()