better implementation

lib/Fmt_ast.ml

@@ -1483,7 +1483,8 @@ and fmt_function ?force_closing_paren ~ctx ~ctx0 ?(wrap_intro = fun x -> hovbox
     match args, typ, body with
     | (_ :: _), _, Pfunction_body body ->
         (* Only [fun]. *)
-        fmt_fun_args_typ args typ, fmt_expression c (sub_exp ~ctx body), hovbox (Params.Indent.fun_ c.conf ~ctx0)
+        fmt_fun_args_typ args typ, fmt_expression c (sub_exp ~ctx body),
+     (Params.Exp.box_fun_expr c.conf ~ctx0)
     | [], _, Pfunction_body _ -> assert false
     | args, typ, Pfunction_cases (cs, _loc, cs_attrs) ->
         (* [fun _ -> function] or [function]. [spilled_attrs] are extra attrs
@@ -1507,7 +1508,6 @@ and fmt_function ?force_closing_paren ~ctx ~ctx0 ?(wrap_intro = fun x -> hovbox
     else noop, noop
   in
   let box k = if should_box then box k else k in
-  let box = match ctx0   with Str _ -> hvbox_if should_box (Params.Indent.fun_ ~ctx0 c.conf) | _ ->  box in
   box
     ( wrap_intro
         (hvbox_if has_cmts_outer 0

lib/Params.ml

@@ -26,6 +26,15 @@ let is_labelled_arg args exp =
       | Labelled _, x | Optional _, x -> phys_equal x exp )
     args
 
+let ctx_is_infix = function
+| Exp { pexp_desc= Pexp_infix ({txt= ":="; _}, _, _); _ } -> false
+    | Exp { pexp_desc= Pexp_infix _; _ } -> true
+    | _ -> false
+
+  let ctx_is_let = function
+    | Lb _ | Str _ | Exp { pexp_desc= Pexp_let _; _ } | Bo _ -> true
+    | _ -> false
+
 let parens_if parens (c : Conf.t) ?(disambiguate = false) k =
   if disambiguate && c.fmt_opts.disambiguate_non_breaking_match.v then
     wrap_if_fits_or parens "(" ")" k
@@ -99,6 +108,23 @@ module Exp = struct
     in
     box_decl (kw $ hvbox_if should_box_args 0 args $ fmt_opt annot)
 
+  let box_fun_expr (c: Conf.t) ~ctx0 =
+    let indent =
+    if ctx_is_infix ctx0 then
+      0
+    else
+        match c.fmt_opts.function_indent_nested.v with
+        | `Always -> c.fmt_opts.function_indent.v
+        | _ ->
+            if ctx_is_let ctx0 then
+              if c.fmt_opts.let_binding_deindent_fun.v then 1
+              else 0
+            else
+              2 in
+    ( match ctx0 with
+      | Str _ -> hvbox indent
+      | _ -> hovbox indent)
+
   (* if the function is the last argument of an apply and no other arguments
      are "complex" (approximation). *)
   let function_attrs_sp c ~ctx0 ~ctx =
@@ -707,14 +733,7 @@ module Align = struct
 end
 
 module Indent = struct
-  let ctx_is_infix = function
-    | Exp { pexp_desc= Pexp_infix ({txt= ":="; _}, _, _); _ } -> false
-    | Exp { pexp_desc= Pexp_infix _; _ } -> true
-    | _ -> false
 
-  let ctx_is_let = function
-    | Lb _ | Str _ | Exp { pexp_desc= Pexp_let _; _ } | Bo _ -> true
-    | _ -> false
 
   let function_ ?(default = 0) (c : Conf.t) ~ctx0 ~parens ~has_label =
     if ctx_is_infix ctx0 then
@@ -725,18 +744,7 @@ module Indent = struct
       | _ when ocp c && parens && not has_label -> default + 1
       | _ -> default
 
-  let fun_ (c : Conf.t) ~ctx0 =
-    if ctx_is_infix ctx0 then
-      0
-    else
-        match c.fmt_opts.function_indent_nested.v with
-        | `Always -> c.fmt_opts.function_indent.v
-        | _ ->
-            if ctx_is_let ctx0 then
-              if c.fmt_opts.let_binding_deindent_fun.v then 1
-              else 0
-            else
-              2
+
 
   let fun_type_annot c = if ocp c then 2 else 4
 

lib/Params.mli

@@ -47,6 +47,8 @@ module Exp : sig
   (** Box and assemble the parts [kw] (up to the arguments), [args] and
       [annot]. *)
 
+  val box_fun_expr : Conf.t -> ctx0:Ast.t -> Fmt.t -> Fmt.t
+
   val function_attrs_sp : Conf.t -> ctx0:Ast.t -> ctx:Ast.t -> bool
   (** Whether a space should be added between the [function] keyword and the
       attributes. *)
@@ -201,9 +203,6 @@ module Indent : sig
   (** Check the [function-indent-nested] option, or return [default] (0 if
       not provided) if the option does not apply. *)
 
-  val fun_ : Conf.t -> ctx0:Ast.t -> int
-  (** Handle [function-indent-nested]. *)
-
   val fun_args : Conf.t -> int
 
   val fun_type_annot : Conf.t -> int