map begin fun on one line, second attempt (#2693)

* apply begin fun on one line

* fmt

* changelog

* showcase infix apply begin function issue

* fix infix apply begin function issue

Author: EmileTrotignon

CHANGES.md

@@ -69,10 +69,11 @@ profile. This started with version 0.26.0.
 
 ### Changed
 
-- \* `|> begin`, `~arg:begin`, `begin if`, `lazy begin`, `begin match` and
-  `begin fun` can now be printed on the same line, with one less indentation
-  level for the body of the inner expression. (#2664, #2666, #2671, #2672,
-  #2681, #2685, @EmileTrotignon) For example :
+- `|> begin`, `~arg:begin`, `begin if`, `lazy begin`, `begin match`,
+  `begin fun` and `map li begin fun`  can now be printed on the same line, with
+  one less indentation level for the body of the inner expression.
+  (#2664, #2666, #2671, #2672, #2681, #2685, #2693, @EmileTrotignon)
+  For example :
   ```ocaml
   (* before *)
   begin

bin/ocamlformat/dune

@@ -18,7 +18,8 @@
   (:standard -open Ocamlformat_stdlib))
  (instrumentation
   (backend bisect_ppx))
- (libraries ocamlformat-lib bin_conf))
+ (libraries ocamlformat-lib bin_conf)
+ (modes byte native))
 
 (rule
  (with-stdout-to

lib/Fmt_ast.ml

@@ -1665,8 +1665,9 @@ and fmt_function ?force_closing_paren ~ctx ~ctx0 ?pro ~wrap_intro
         $ hvbox_if has_cmts_outer 0
             ( cmts_outer
             $ Params.Exp.box_fun_decl ~ctx0 c.conf
-                (pro_inner $ cmts_inner $ opn_attr_paren $ opn_paren $ head)
-            ) )
+                ( pro_inner $ cmts_inner
+                $ Params.Exp.box_fun_decl_after_pro ~ctx0
+                    (opn_attr_paren $ opn_paren $ head) ) ) )
     in
     body ~pro $ cls_paren
     $ fmt_attributes c ~pre:Space attrs
@@ -1874,7 +1875,8 @@ and fmt_infix_op_args c ~parens xexp op_args =
     in
     if Params.Exp.Infix_op_arg.dock c.conf xarg then
       (* Indentation of docked fun or function start before the operator. *)
-      hovbox 2 (fmt_expression c ~parens ~box:false ~pro xarg)
+      hovbox ~name:"Infix_op_arg docked" 2
+        (fmt_expression c ~parens ~box:false ~pro xarg)
     else
       match xarg.ast.pexp_desc with
       | Pexp_function _ | Pexp_beginend _ ->
@@ -2289,6 +2291,26 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
             ( expr_epi
             $ Params.parens_if outer_parens c.conf
                 (args $ fmt_atrs $ fmt_if inner_parens (str ")")) )
+      | Pexp_beginend ({pexp_desc= Pexp_function _; _}, _) ->
+          let fmt_atrs =
+            fmt_attributes c ~pre:(Break (1, -2)) pexp_attributes
+          in
+          let force =
+            if Location.is_single_line pexp_loc c.conf.fmt_opts.margin.v then
+              Fit
+            else Break
+          in
+          let pro =
+            intro_epi
+            $ fmt_if parens (str "(")
+            $ ( fmt_args_grouped ~epi:fmt_atrs e0 args_before
+              $ fmt_if parens (closing_paren c ~force ~offset:(-3)) )
+          in
+          let label_sep = Params.Exp.fun_label_sep c.conf in
+          let pro = pro $ break 1 0 $ fmt_label lbl label_sep in
+          expr_epi
+          $ hovbox 4
+              (fmt_expression c ~pro ~box:false (sub_exp ~ctx last_arg))
       | _ ->
           let fmt_atrs =
             fmt_attributes c ~pre:(Break (1, -2)) pexp_attributes
@@ -2927,8 +2949,8 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
       pro $ fmt_indexop_access c ctx ~fmt_atrs ~has_attr ~parens x
   | Pexp_hole -> pro $ hvbox 0 (fmt_hole () $ fmt_atrs)
   | Pexp_beginend (e, infix_ext_attrs) ->
-      fmt_beginend c ~loc:pexp_loc ~box ~pro ~ctx ~fmt_atrs ~infix_ext_attrs
-        ~indent_wrap ?eol e
+      fmt_beginend c ~loc:pexp_loc ~box ~pro ~ctx ~ctx0 ~fmt_atrs
+        ~infix_ext_attrs ~indent_wrap ?eol e
   | Pexp_parens e ->
       pro
       $ hvbox 0
@@ -2959,27 +2981,36 @@ and fmt_lazy c ~ctx ?(pro = noop) ~fmt_atrs ~infix_ext_attrs ~parens e =
          $ fmt_expression c ~pro:kw_inner (sub_exp ~ctx e)
          $ fmt_atrs ) )
 
-and fmt_beginend c ~loc ?(box = true) ?(pro = noop) ~ctx ~fmt_atrs
+and fmt_beginend c ~loc ?(box = true) ?(pro = noop) ~ctx ~ctx0 ~fmt_atrs
     ~infix_ext_attrs ~indent_wrap ?eol e =
   let cmts_before = Cmts.fmt_before c ?eol loc in
   let begin_ = fmt_infix_ext_attrs c ~pro:(str "begin") infix_ext_attrs
-  and end_ = str "end" $ fmt_atrs in
+  and end_ =
+    (if not box then break 1000 (-2) else break 1000 0)
+    $ str "end" $ fmt_atrs
+  in
+  let box_beginend_sb = Params.Exp.box_beginend_subexpr c.conf ~ctx ~ctx0 in
+  let beginend_box =
+    if Params.Exp.box_beginend c.conf ~ctx ~ctx0 then
+      hvbox ~name:"beginend" 2
+    else Fn.id
+  in
   cmts_before
   $
   match e.pexp_desc with
   | Pexp_match _ | Pexp_try _ | Pexp_function _ | Pexp_ifthenelse _ ->
-      hvbox 0
-        ( fmt_expression c
-            ~pro:(pro $ begin_ $ str " ")
-            ~box ?eol ~parens:false ~indent_wrap (sub_exp ~ctx e)
-        $ break 1 0 $ end_ )
+      beginend_box
+        (fmt_expression c
+           ~pro:(pro $ begin_ $ str " ")
+           ~box:false ?eol ~parens:false ~indent_wrap (sub_exp ~ctx e) )
+      $ end_
   | _ ->
-      hvbox 0
+      beginend_box
         ( hvbox 0 (pro $ begin_)
         $ break 1 2
-        $ fmt_expression c ~box ?eol ~parens:false ~indent_wrap
-            (sub_exp ~ctx e)
-        $ force_break $ end_ )
+        $ fmt_expression c ~box:box_beginend_sb ?eol ~parens:false
+            ~indent_wrap (sub_exp ~ctx e) )
+      $ end_
 
 and fmt_let_bindings c ~ctx0 ~parens ~has_attr ~fmt_atrs ~fmt_expr ~loc_in
     rec_flag bindings body =

lib/Params.ml

@@ -145,7 +145,11 @@ module Exp = struct
         | Pexp_apply (_, args) -> (
           (* Rhs is an apply and it ends with a [fun]. *)
           match List.last_exn args with
-          | _, {pexp_desc= Pexp_function _; _} -> true
+          | _, {pexp_desc= Pexp_function _; _}
+           |( _
+            , { pexp_desc= Pexp_beginend ({pexp_desc= Pexp_function _; _}, _)
+              ; _ } ) ->
+              true
           | _ -> false )
         | Pexp_match _ | Pexp_try _ -> true
         | _ -> false
@@ -172,10 +176,16 @@ module Exp = struct
       | Some ((Labelled _ | Optional _), _, _) -> true
       | _ -> false
     in
+    let is_ctx_beginend =
+      match ctx0 with
+      | Exp {pexp_desc= Pexp_beginend _; _} -> true
+      | _ -> false
+    in
     if Conf.(c.fmt_opts.ocp_indent_compat.v) then
       if last_arg || is_labelled_arg then break 1 2 else str " "
     else if is_labelled_arg then break 1 2
     else if last_arg then break 1 0
+    else if is_ctx_beginend then break 1 0
     else str " "
 
   let box_fun_decl_args ~ctx ~ctx0 ?(last_arg = false) ?epi c ~parens ~kw
@@ -224,38 +234,42 @@ module Exp = struct
         $ hvbox_if should_box_args 0 (args $ fmt_opt annot $ fmt_opt epi) )
 
   let box_fun_expr (c : Conf.t) ~source ~ctx0 ~ctx =
-    let indent =
-      if ctx_is_rhs_of_infix ~ctx0 ~ctx then 0
-      else if Poly.equal c.fmt_opts.function_indent_nested.v `Always then
-        c.fmt_opts.function_indent.v
-      else if ctx_is_let_or_fun ~ctx ctx0 then
-        if c.fmt_opts.let_binding_deindent_fun.v then 1 else 0
-      else if ocp c then
-        let begins_line loc =
-          Source.begins_line ~ignore_spaces:true source loc
+    match ctx0 with
+    | Exp {pexp_desc= Pexp_beginend _; _} -> (Fn.id, 0)
+    | _ ->
+        let indent =
+          if ctx_is_rhs_of_infix ~ctx0 ~ctx then 0
+          else if Poly.equal c.fmt_opts.function_indent_nested.v `Always then
+            c.fmt_opts.function_indent.v
+          else if ctx_is_let_or_fun ~ctx ctx0 then
+            if c.fmt_opts.let_binding_deindent_fun.v then 1 else 0
+          else if ocp c then
+            let begins_line loc =
+              Source.begins_line ~ignore_spaces:true source loc
+            in
+            match ctx_is_apply_and_exp_is_arg ~ctx ~ctx0 with
+            | Some (Nolabel, fun_exp, is_last_arg) ->
+                if begins_line fun_exp.pexp_loc then
+                  if is_last_arg then 5 else 3
+                else 2
+            | Some ((Labelled x | Optional x), fun_exp, is_last_arg) ->
+                if begins_line fun_exp.pexp_loc then
+                  (* The [fun] had to break after the label, nested boxes
+                     must be indented less. The last argument is special as
+                     the box structure is different. *)
+                  if is_last_arg then 4 else 2
+                else if begins_line x.loc then 4
+                else 2
+            | None -> if ctx_is_apply_and_exp_is_func ~ctx ctx0 then 3 else 2
+          else if
+            ctx_is_apply_and_exp_is_last_arg_and_other_args_are_simple c ~ctx
+              ~ctx0
+          then 4
+          else 2
         in
-        match ctx_is_apply_and_exp_is_arg ~ctx ~ctx0 with
-        | Some (Nolabel, fun_exp, is_last_arg) ->
-            if begins_line fun_exp.pexp_loc then if is_last_arg then 5 else 3
-            else 2
-        | Some ((Labelled x | Optional x), fun_exp, is_last_arg) ->
-            if begins_line fun_exp.pexp_loc then
-              (* The [fun] had to break after the label, nested boxes must be
-                 indented less. The last argument is special as the box
-                 structure is different. *)
-              if is_last_arg then 4 else 2
-            else if begins_line x.loc then 4
-            else 2
-        | None -> if ctx_is_apply_and_exp_is_func ~ctx ctx0 then 3 else 2
-      else if
-        ctx_is_apply_and_exp_is_last_arg_and_other_args_are_simple c ~ctx
-          ~ctx0
-      then 4
-      else 2
-    in
-    let name = "Params.box_fun_expr" in
-    let mkbox = if ctx_is_let_or_fun ~ctx ctx0 then hvbox else hovbox in
-    (mkbox ~name indent, ~-indent)
+        let name = "Params.box_fun_expr" in
+        let mkbox = if ctx_is_let_or_fun ~ctx ctx0 then hvbox else hovbox in
+        (mkbox ~name indent, ~-indent)
 
   (* if the function is the last argument of an apply and no other arguments
      are "complex" (approximation). *)
@@ -317,6 +331,7 @@ module Exp = struct
 
   let box_fun_decl ~ctx0 c k =
     match ctx0 with
+    | Exp {pexp_desc= Pexp_beginend _; _} -> hovbox 2 k
     | _ when ocp c -> hvbox 2 k
     (* Avoid large indentation for [let _ = function]. *)
     | Lb
@@ -326,6 +341,30 @@ module Exp = struct
     | Str _ | Lb _ | Clf _ | Exp {pexp_desc= Pexp_let _; _} -> hovbox 4 k
     | _ -> hvbox 2 k
 
+  let box_fun_decl_after_pro ~ctx0 =
+    match ctx0 with
+    | Exp {pexp_desc= Pexp_beginend _; _} ->
+        hvbox (2 - String.length "begin ")
+    | _ -> Fn.id
+
+  let box_beginend c ~ctx0 ~ctx =
+    let contains_fun =
+      match ctx with
+      | Exp {pexp_desc= Pexp_beginend ({pexp_desc= Pexp_function _; _}, _); _}
+        ->
+          true
+      | _ -> false
+    in
+    contains_fun
+    && not
+         (ctx_is_apply_and_exp_is_last_arg_and_other_args_are_simple c ~ctx
+            ~ctx0 )
+
+  let box_beginend_subexpr c ~ctx0 ~ctx =
+    not
+      (ctx_is_apply_and_exp_is_last_arg_and_other_args_are_simple c ~ctx
+         ~ctx0 )
+
   let match_inner_pro ~ctx0 ~parens =
     if parens then false
     else
@@ -854,7 +893,7 @@ let get_if_then_else (c : Conf.t) ~pro ~first ~last ~parens_bch
       ; box_keyword_and_expr= Fn.id
       ; branch_pro= branch_pro ()
       ; wrap_parens= wrap_parens ~wrap_breaks:(wrap (break 1000 2) noop)
-      ; box_expr= Some false
+      ; box_expr= Some beginend
       ; expr_pro= None
       ; expr_eol= Some (break 1 2)
       ; branch_expr

lib/Params.mli

@@ -81,6 +81,13 @@ module Exp : sig
   val box_fun_decl : ctx0:Ast.t -> Conf.t -> Fmt.t -> Fmt.t
   (** Box a function decl from the label to the arrow. *)
 
+  val box_fun_decl_after_pro : ctx0:Ast.t -> Fmt.t -> Fmt.t
+  (** Box a function decl from after the [pro] to the arrow. *)
+
+  val box_beginend : Conf.t -> ctx0:Ast.t -> ctx:Ast.t -> bool
+
+  val box_beginend_subexpr : Conf.t -> ctx0:Ast.t -> ctx:Ast.t -> bool
+
   val match_inner_pro : ctx0:Ast.t -> parens:bool -> bool
   (**  whether the [pro] argument of [fmt_match] should be displayed as an inner
      or outer prologue.*)

test/passing/refs.default/attributes.ml.ref

@@ -395,8 +395,7 @@ include [@foo] M [@boo]
 let () =
   let () =
     S.ntyp Cbor_type.Reserved
-    @@ S.tok
-         begin[@warning "-4"] fun ev ->
+    @@ S.tok begin[@warning "-4"] fun ev ->
            match ev with Cbor_event.Reserved int -> Some int | _ -> None
          end
   in

test/passing/refs.default/cases_exp_grouping.ml.ref

@@ -116,7 +116,9 @@ let a =
   | B -> bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
   end
 
-let a = begin match f x i with A -> a | B -> b end
+let a =
+  begin match f x i with A -> a | B -> b
+  end
 
 let a =
   begin[@a] match[@b]
@@ -175,7 +177,9 @@ let a =
   | B -> bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
   end
 
-let a = begin try f x i with A -> a | B -> b end
+let a =
+  begin try f x i with A -> a | B -> b
+  end
 
 let a =
   begin[@a] try[@b]

test/passing/refs.default/exp_grouping-parens.ml.ref

@@ -388,8 +388,7 @@ let v =
       a f b)
 
 let v =
-  map x
-    begin%ext1 fun%ext2 x y z ->
+  map x begin%ext1 fun%ext2 x y z ->
       ya f;
       a f b
     end
@@ -554,3 +553,47 @@ let _ =
     a
   end
   [@a]
+
+let () =
+  fooooo
+  |>>>>> List.iter (fun a ->
+             let x =
+               some_really_really_really_long_name_that_doesn't_fit_on_the_line
+               $ y
+             in
+             fooooooooooo x)
+
+let () =
+  fooooo
+  |>>>>> List.iter
+           (fun aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa ->
+             let x =
+               some_really_really_really_long_name_that_doesn't_fit_on_the_line
+               $ y
+             in
+             fooooooooooo x)
+
+let () =
+  fooooo
+  |>>>>> List.iter (fun a ->
+             let x =
+               some_really_really_really_long_name_that_doesn't_fit_on_the_line
+               $ y
+             in
+             fooooooooooo x)
+
+let () =
+  fooooo
+  |>>>>> List.iter aaaaaaaaaaaa aaaaaaaaaaaaaaaaaaaaa aaaaaaaaaaaaaaa
+           aaaaaaaaaaaaaa
+           (fun
+             aaaaaaaaaaaaaaaaaaaaaaaaaaa
+             aa
+             aaaaaaaaaaa
+             aaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+           ->
+             let x =
+               some_really_really_really_long_name_that_doesn't_fit_on_the_line
+               $ y
+             in
+             fooooooooooo x)