begin fun x -> now can be printed on a single line (#2664)

    map x 
      begin 
        fun y ->
          print x;
          y
      end

is now

    map x 
      begin fun y ->
        print x;
        y
      end

Author: EmileTrotignon

CHANGES.md

@@ -33,8 +33,8 @@ profile. This started with version 0.26.0.
 
 ### Changed
 
-- `begin match` can now be printed on the same line, with one less indentation
-  level for the body of the match. (#2666, @EmileTrotignon)
+- `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. (#2666, @EmileTrotignon)
 
 ## 0.27.0
 

lib/Fmt_ast.ml

@@ -1494,7 +1494,7 @@ and fmt_indexop_access c ctx ~fmt_atrs ~has_attr ~parens x =
 
 (** Format a [Pexp_function]. [wrap_intro] wraps up to after the [->] and is
     responsible for breaking. *)
-and fmt_function ?(last_arg = false) ?force_closing_paren ~ctx ~ctx0
+and fmt_function ?(last_arg = false) ?force_closing_paren ~ctx ~ctx0 ?pro
     ~wrap_intro ?box:(should_box = true) ~label ?(parens = false) ?ext ~attrs
     ~loc c (args, typ, body) =
   let should_box =
@@ -1637,13 +1637,20 @@ and fmt_function ?(last_arg = false) ?force_closing_paren ~ctx ~ctx0
       wrap (fits_breaks "(" "") (fits_breaks ")" "")
     else Fn.id
   in
+  let pro_outer, pro_inner =
+    let pro = fmt_opt pro in
+    if Params.Exp.function_inner_pro ~has_cmts_outer ~ctx0 then (noop, pro)
+    else (pro, noop)
+  in
   let body =
     let pro =
       wrap_intro
-        (hvbox_if has_cmts_outer 0
-           ( cmts_outer
-           $ Params.Exp.box_fun_decl ~ctx0 c.conf
-               (fmt_label label label_sep $ cmts_inner $ opn_paren $ head) ) )
+        ( pro_outer
+        $ hvbox_if has_cmts_outer 0
+            ( cmts_outer
+            $ Params.Exp.box_fun_decl ~ctx0 c.conf
+                ( pro_inner $ fmt_label label label_sep $ cmts_inner
+                $ opn_paren $ head ) ) )
     in
     body ~pro $ cls_paren
   in
@@ -1911,17 +1918,17 @@ and fmt_match c ?pro ?eol ~loc ~parens ?ext ctx xexp cs e0 keyword =
   let cmts_before = Cmts.fmt_before c ?eol loc in
   let ctx0 = xexp.ctx in
   let indent = Params.match_indent c.conf ~parens ~ctx:ctx0 in
-  let pro_outside_parens, pro_inside_parens =
+  let pro_outer, pro_inner =
     let pro = fmt_opt pro in
-    if Params.Exp.box_pro_with_match ~ctx0 ~parens then (noop, pro)
+    if Params.Exp.match_inner_pro ~ctx0 ~parens then (noop, pro)
     else (pro, noop)
   in
   hvbox indent
-    ( cmts_before $ pro_outside_parens
+    ( cmts_before $ pro_outer
     $ Params.Exp.wrap c.conf ~parens ~disambiguate:true
       @@ Params.Align.match_ c.conf ~xexp
       @@ ( hvbox 0
-             ( hvbox 0 (pro_inside_parens $ keyword)
+             ( hvbox 0 (pro_inner $ keyword)
              $ fmt_extension_suffix c ext
              $ fmt_attributes c xexp.ast.pexp_attributes
              $ break 1 2
@@ -2358,11 +2365,10 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
              $ cut_break $ str "." $ fmt_longident_loc c lid $ fmt_atrs ) )
   | Pexp_function (args, typ, body) ->
       let wrap_intro intro =
-        hovbox ~name:"fmt_expression | Pexp_function" 2 (pro $ intro)
-        $ space_break
+        hovbox ~name:"fmt_expression | Pexp_function" 2 intro $ space_break
       in
-      fmt_function ~wrap_intro ~box ~ctx ~ctx0 ~label:Nolabel ~parens ?ext
-        ~attrs:pexp_attributes ~loc:pexp_loc c (args, typ, body)
+      fmt_function ~pro ~wrap_intro ~box ~ctx ~ctx0 ~label:Nolabel ~parens
+        ?ext ~attrs:pexp_attributes ~loc:pexp_loc c (args, typ, body)
   | Pexp_ident {txt; loc} ->
       let outer_parens = has_attr && parens in
       pro
@@ -2565,7 +2571,7 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
            && c.conf.fmt_opts.break_cases.v <> `Vertical ) ->
       let cmts_before = Cmts.fmt_before c ?eol pexp_loc in
       let pro_outer, pro_inner =
-        if Params.Exp.box_pro_with_match ~ctx0 ~parens then (noop, pro)
+        if Params.Exp.match_inner_pro ~ctx0 ~parens then (noop, pro)
         else (pro, noop)
       in
       (* side effects of Cmts.fmt_before before [fmt_pattern] is important *)
@@ -2921,11 +2927,11 @@ and fmt_beginend c ?(box = true) ?(pro = noop) ~ctx ~fmt_atrs ~ext
   let begin_ = str "begin" $ fmt_extension_suffix c ext $ fmt_atrs
   and end_ = str "end" in
   match e.pexp_desc with
-  | Pexp_match _ | Pexp_try _ ->
+  | Pexp_match _ | Pexp_try _ | Pexp_function _ ->
       pro
       $ hvbox 0
           ( fmt_expression c
-              ~pro:(begin_ $ break 1 0)
+              ~pro:(begin_ $ str " ")
               ~box ?eol ~parens:false ~indent_wrap (sub_exp ~ctx e)
           $ break 1 0 $ end_ )
   | Pexp_extension

lib/Params.ml

@@ -308,13 +308,26 @@ module Exp = struct
   let box_fun_decl ~ctx0 c k =
     match ctx0 with
     | _ when ocp c -> hvbox 2 k
+    (* Avoid large indentation for [let _ = function]. *)
+    | Lb
+        {pvb_body= Pfunction_body {pexp_desc= Pexp_function ([], _, _); _}; _}
+      ->
+        hovbox 2 k
     | Str _ | Lb _ | Clf _ | Exp {pexp_desc= Pexp_let _; _} -> hovbox 4 k
     | _ -> hvbox 2 k
 
-  let box_pro_with_match ~ctx0 ~parens =
+  let match_inner_pro ~ctx0 ~parens =
     if parens then false
     else
       match ctx0 with Exp {pexp_desc= Pexp_infix _; _} -> false | _ -> true
+
+  let function_inner_pro ~has_cmts_outer ~ctx0 =
+    if has_cmts_outer then false
+    else
+      match ctx0 with
+      | Str _ | Lb _ | Exp {pexp_desc= Pexp_ifthenelse _ | Pexp_let _; _} ->
+          false
+      | _ -> true
 end
 
 module Mod = struct

lib/Params.mli

@@ -81,8 +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_pro_with_match : ctx0:Ast.t -> parens:bool -> bool
-  (** whether the [~pro] argument should be in the same box as the [match] keyword. *)
+  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.*)
+
+  val function_inner_pro : has_cmts_outer:bool -> ctx0:Ast.t -> bool
+  (** whether the [pro] argument of [fmt_function] should be displayed as an
+    inner or outer prologue. *)
 end
 
 module Mod : sig

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

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

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

@@ -348,3 +348,47 @@ let _ =
       begin [@foo]
         y
       end
+
+let v = map x (fun x y z -> y)
+
+let v =
+  map x (fun x arggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg ->
+      y)
+
+let v =
+  map x
+    (fun x arggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg ->
+      y)
+
+let v =
+  map x
+    (fun x yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy z ->
+      print y;
+      z)
+
+let v =
+  map x
+    (fun
+      x
+      argggggggggggggggggggggggggggggggggg
+      gggggggggggggggggggg
+      ggggggggggggggg
+    -> y)
+
+let v =
+  map x
+    (fun x yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy z ->
+      print y;
+      z)
+
+let v =
+  map x (fun x y z ->
+      ya f;
+      a f b)
+
+let v =
+  map x
+    [%ext1
+      fun%ext2 x y z ->
+        ya f;
+        a f b]

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

@@ -16,13 +16,12 @@ let () =
 
 let () =
   List.iter
-    begin
-      fun v ->
-        (* do a lot of things *)
-        let a = "a" in
-        let b = "b" in
-        let c = "c" in
-        ()
+    begin fun v ->
+      (* do a lot of things *)
+      let a = "a" in
+      let b = "b" in
+      let c = "c" in
+      ()
     end
     values
 
@@ -406,3 +405,54 @@ let _ =
       begin [@foo]
         y
       end
+
+let v = map x begin fun x y z -> y end
+
+let v =
+  map x
+    begin fun x arggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg ->
+      y
+    end
+
+let v =
+  map x
+    begin fun x
+      arggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
+    -> y
+    end
+
+let v =
+  map x
+    (fun x yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy z ->
+      print y;
+      z)
+
+let v =
+  map x
+    begin fun x
+      argggggggggggggggggggggggggggggggggg
+      gggggggggggggggggggg
+      ggggggggggggggg
+    -> y
+    end
+
+let v =
+  map x
+    (fun x yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy z ->
+      print y;
+      z)
+
+let v =
+  map x
+    begin fun x y z ->
+      ya f;
+      a f b
+    end
+
+let v =
+  map x
+    begin%ext1
+      fun%ext2 x y z ->
+        ya f;
+        a f b
+    end

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

@@ -444,11 +444,10 @@ let () =
   let () =
     S.ntyp Cbor_type.Reserved
     @@ S.tok
-         begin [@warning "-4"]
-           fun ev ->
-             match ev with
-             | Cbor_event.Reserved int -> Some int
-             | _ -> None
+         begin [@warning "-4"] fun ev ->
+           match ev with
+           | Cbor_event.Reserved int -> Some int
+           | _ -> None
          end
   in
   ()

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

@@ -408,3 +408,43 @@ let _ =
       y
     end
 ;;
+
+let v = map x (fun x y z -> y)
+let v = map x (fun x arggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg -> y)
+
+let v =
+  map x (fun x arggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg -> y)
+;;
+
+let v =
+  map x (fun x yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy z ->
+    print y;
+    z)
+;;
+
+let v =
+  map
+    x
+    (fun x argggggggggggggggggggggggggggggggggg gggggggggggggggggggg ggggggggggggggg -> y)
+;;
+
+let v =
+  map x (fun x yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy z ->
+    print y;
+    z)
+;;
+
+let v =
+  map x (fun x y z ->
+    ya f;
+    a f b)
+;;
+
+let v =
+  map
+    x
+    [%ext1
+      fun%ext2 x y z ->
+        ya f;
+        a f b]
+;;