Correctly format chunks of file in presence of enable/disable floating attributes

CHANGES.md

@@ -6,6 +6,7 @@
 - Reset max-indent when the `max-indent` option is not set (#2131, @hhugo, @gpetiot)
 - Add missing parentheses around immediate objects having attributes attached in 4.14 (#2144, @gpetiot)
 - Fix dropped comment attached to the identifier of an open-expression (#2155, @gpetiot)
+- Correctly format chunks of file in presence of `enable`/`disable` floating attributes (#<PR_NUMBER>, @gpetiot)
 
 ### Changes
 

lib/Fmt_ast.ml

@@ -4412,15 +4412,125 @@ let fmt_repl_file c _ itms =
 
 (** Entry points *)
 
+module Chunk = struct
+  type state = Enable | Disable of Location.t
+
+  type 'a t =
+    | Structure : structure t
+    | Signature : signature t
+    | Use_file : use_file t
+
+  let disabling (c : c) attr =
+    (not c.conf.opr_opts.disable)
+    && (update_config ~quiet:true c [attr]).conf.opr_opts.disable
+
+  let enabling (c : c) attr =
+    c.conf.opr_opts.disable
+    && not (update_config ~quiet:true c [attr]).conf.opr_opts.disable
+
+  let update c disable =
+    let opr_opts = {c.conf.opr_opts with disable} in
+    {c with conf= {c.conf with opr_opts}}
+
+  let init_loc =
+    let pos =
+      Lexing.
+        { pos_cnum= 0
+        ; pos_bol= 0
+        ; pos_lnum= 0
+        ; pos_fname= !Location.input_name }
+    in
+    Location.{loc_ghost= false; loc_start= pos; loc_end= pos}
+
+  let is_attr (type a) (x : a list t) : a -> _ =
+    match x with
+    | Structure -> (
+        function
+        | {pstr_desc= Pstr_attribute attr; pstr_loc} -> Some (attr, pstr_loc)
+        | _ -> None )
+    | Signature -> (
+        function
+        | {psig_desc= Psig_attribute attr; psig_loc} -> Some (attr, psig_loc)
+        | _ -> None )
+    | Use_file -> (
+        function
+        | Ptop_def ({pstr_desc= Pstr_attribute attr; pstr_loc} :: _) ->
+            Some (attr, pstr_loc)
+        | _ -> None )
+
+  let fmt_item (type a) (x : a list t) : c -> Ast.t -> a list -> Fmt.t =
+    match x with
+    | Structure -> fmt_structure
+    | Signature -> fmt_signature
+    | Use_file -> fmt_toplevel ?force_semisemi:None
+
+  let loc_end (type a) (x : a list t) (l : a list) =
+    match x with
+    | Structure -> (List.last_exn l).pstr_loc.loc_end
+    | Signature -> (List.last_exn l).psig_loc.loc_end
+    | Use_file ->
+        let item =
+          match List.last_exn l with
+          | Ptop_def x -> `Item (List.last_exn x)
+          | Ptop_dir x -> `Directive x
+        in
+        (Ast.location (Tli item)).loc_end
+
+  let split fg c l =
+    let is_state_attr ~f c x =
+      match is_attr fg x with
+      | Some (attr, loc) when f c attr -> Some loc
+      | _ -> None
+    in
+    List.fold_left l ~init:([], c) ~f:(fun (acc, c) x ->
+        match is_state_attr ~f:disabling c x with
+        | Some loc -> ((Disable loc, [x]) :: acc, update c true)
+        | None -> (
+          match is_state_attr ~f:enabling c x with
+          | Some _ -> ((Enable, [x]) :: acc, update c false)
+          | None -> (
+            match acc with
+            | [] ->
+                let chunk =
+                  if c.conf.opr_opts.disable then (Disable init_loc, [x])
+                  else (Enable, [x])
+                in
+                (chunk :: acc, c)
+            | (st, h) :: t -> ((st, x :: h) :: t, c) ) ) )
+    |> fst
+    |> List.map ~f:(function
+         | Enable, lx -> (Enable, List.rev lx)
+         | Disable loc, lx -> (Disable loc, List.rev lx) )
+    |> List.rev
+
+  let fmt fg c ctx chunks =
+    List.foldi chunks ~init:(c, noop) ~f:(fun i (c, output) -> function
+      | Disable item_loc, lx ->
+          let c = update c true in
+          let loc_end = loc_end fg lx in
+          let loc = Location.{item_loc with loc_end} in
+          ( c
+          , output
+            $ Cmts.fmt_before c item_loc ~eol:(fmt "\n@;<1000 0>")
+            $ fmt_if (i > 0) "\n@;<1000 0>"
+            $ str (String.strip (Source.string_at c.source loc)) )
+      | Enable, lx ->
+          let c = update c false in
+          (c, output $ fmt_if (i > 0) "@;<1000 0>" $ fmt_item fg c ctx lx) )
+    |> snd
+
+  let split_and_fmt fg c ctx l = fmt fg c ctx @@ split fg c l
+end
+
 let fmt_file (type a) ~ctx ~fmt_code ~debug (fragment : a Extended_ast.t)
     source cmts conf (itms : a) =
   let c = {source; cmts; conf; debug; fmt_code} in
   match (fragment, itms) with
   | Structure, [] | Signature, [] | Use_file, [] ->
       Cmts.fmt_after ~pro:noop c Location.none
-  | Structure, l -> fmt_structure c ctx l
-  | Signature, l -> fmt_signature c ctx l
-  | Use_file, l -> fmt_toplevel c ctx l
+  | Structure, l -> Chunk.split_and_fmt Structure c ctx l
+  | Signature, l -> Chunk.split_and_fmt Signature c ctx l
+  | Use_file, l -> Chunk.split_and_fmt Use_file c ctx l
   | Core_type, ty -> fmt_core_type c (sub_typ ~ctx:(Pld (PTyp ty)) ty)
   | Module_type, mty ->
       compose_module ~f:Fn.id

test/passing/dune.inc

@@ -1289,6 +1289,24 @@
  (package ocamlformat)
  (action (diff tests/directives.mlt.err directives.mlt.stderr)))
 
+(rule
+ (deps tests/.ocamlformat )
+ (package ocamlformat)
+ (action
+  (with-stdout-to disable_attr.ml.stdout
+   (with-stderr-to disable_attr.ml.stderr
+     (run %{bin:ocamlformat} --margin-check %{dep:tests/disable_attr.ml})))))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/disable_attr.ml disable_attr.ml.stdout)))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/disable_attr.ml.err disable_attr.ml.stderr)))
+
 (rule
  (deps tests/.ocamlformat )
  (package ocamlformat)

test/passing/tests/comments_around_disabled.ml

@@ -1,7 +1,8 @@
-(* cmts *)
+(* cmts before *)
 
 [@@@ocamlformat "disable"]
-let () = ()
+let () =
+  ()
 [@@@ocamlformat "enable"]
 
-(* cmts *)
+(* cmts after *)

test/passing/tests/comments_around_disabled.ml.ref

@@ -1,8 +1,8 @@
-(* cmts *)
-[@@@ocamlformat "disable"]
-
-let () = ()
+(* cmts before *)
 
+[@@@ocamlformat "disable"]
+let () =
+  ()
 [@@@ocamlformat "enable"]
 
-(* cmts *)
+(* cmts after *)

test/passing/tests/disable_attr.ml

@@ -0,0 +1,4 @@
+[@@@ocamlformat "disable"]
+
+(** hello *)
+let foo = 42

test/passing/tests/skip.ml

@@ -2,7 +2,6 @@
 
 let this_won't_be_formatted =
                1
-
 [@@@ocamlformat "enable"]
 
 let x = function