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 (#2156, @gpetiot)
 
 ### Changes
 

lib/Chunk.ml

@@ -0,0 +1,70 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                              OCamlFormat                               *)
+(*                                                                        *)
+(*            Copyright (c) Facebook, Inc. and its affiliates.            *)
+(*                                                                        *)
+(*      This source code is licensed under the MIT license found in       *)
+(*      the LICENSE file in the root directory of this source tree.       *)
+(*                                                                        *)
+(**************************************************************************)
+
+open Extended_ast
+
+type state = Enable | Disable of Location.t
+
+type 'a t =
+  | Structure : structure t
+  | Signature : signature t
+  | Use_file : use_file t
+
+let update_conf ?quiet c l = List.fold ~init:c l ~f:(Conf.update ?quiet)
+
+let disabling (c : Conf.t) attr =
+  (not c.opr_opts.disable)
+  && (update_conf ~quiet:true c [attr]).opr_opts.disable
+
+let enabling (c : Conf.t) attr =
+  c.opr_opts.disable
+  && not (update_conf ~quiet:true c [attr]).opr_opts.disable
+
+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) (fg : a list t) (x : a) =
+  match (fg, x) with
+  | Structure, {pstr_desc= Pstr_attribute x; pstr_loc} -> Some (x, pstr_loc)
+  | Signature, {psig_desc= Psig_attribute x; psig_loc} -> Some (x, psig_loc)
+  | Use_file, Ptop_def ({pstr_desc= Pstr_attribute x; pstr_loc} :: _) ->
+      Some (x, pstr_loc)
+  | _ -> None
+
+let is_state_attr fg ~f c x =
+  match is_attr fg x with
+  | Some (attr, loc) when f c attr -> Some loc
+  | _ -> None
+
+let split fg c l =
+  List.fold_left l ~init:([], c) ~f:(fun (acc, c) x ->
+      match is_state_attr fg ~f:disabling c x with
+      | Some loc -> ((Disable loc, [x]) :: acc, Conf.update_state c `Disable)
+      | None -> (
+        match is_state_attr fg ~f:enabling c x with
+        | Some _ -> ((Enable, [x]) :: acc, Conf.update_state c `Enable)
+        | None -> (
+          match acc with
+          | [] ->
+              let chunk =
+                if c.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.rev_map ~f:(function
+       | Enable, lx -> (Enable, List.rev lx)
+       | Disable loc, lx -> (Disable loc, List.rev lx) )

lib/Chunk.mli

@@ -0,0 +1,19 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                              OCamlFormat                               *)
+(*                                                                        *)
+(*            Copyright (c) Facebook, Inc. and its affiliates.            *)
+(*                                                                        *)
+(*      This source code is licensed under the MIT license found in       *)
+(*      the LICENSE file in the root directory of this source tree.       *)
+(*                                                                        *)
+(**************************************************************************)
+
+type state = Enable | Disable of Location.t
+
+type 'a t =
+  | Structure : Extended_ast.structure t
+  | Signature : Extended_ast.signature t
+  | Use_file : Extended_ast.use_file t
+
+val split : 'a list t -> Conf.t -> 'a list -> (state * 'a list) list

lib/Conf.ml

@@ -2093,6 +2093,11 @@ let update ?(quiet = false) c {attr_name= {txt; loc}; attr_payload; _} =
 let update_value config ~name ~value =
   C.update ~config ~from:`Commandline ~name ~value ~inline:false
 
+let update_state c state =
+  let disable = match state with `Enable -> false | `Disable -> true in
+  let opr_opts = {c.opr_opts with disable} in
+  {c with opr_opts}
+
 let print_config = C.print_config
 
 module UI = struct

lib/Conf.mli

@@ -129,6 +129,8 @@ val update : ?quiet:bool -> t -> Parsetree.attribute -> t
 val update_value :
   t -> name:string -> value:string -> (t, Config_option.Error.t) Result.t
 
+val update_state : t -> [`Enable | `Disable] -> t
+
 val print_config : t -> unit
 
 module UI : sig

lib/Fmt_ast.ml

@@ -4412,15 +4412,57 @@ let fmt_repl_file c _ itms =
 
 (** Entry points *)
 
+module Chunk = struct
+  open Chunk
+
+  let fmt_item (type a) (fg : a list t) : c -> Ast.t -> a list -> Fmt.t =
+    match fg with
+    | Structure -> fmt_structure
+    | Signature -> fmt_signature
+    | Use_file -> fmt_toplevel ?force_semisemi:None
+
+  let loc_end (type a) (fg : a list t) (l : a list) =
+    match fg 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 update_conf c state = {c with conf= Conf.update_state c.conf state}
+
+  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_conf c `Disable 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_conf c `Enable 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.conf 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 *)
 
 [@@@ocamlformat "disable"]
-let () = ()
+let () =
+  ()
 [@@@ocamlformat "enable"]
 
 (* cmts *)

test/passing/tests/comments_around_disabled.ml.ref

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

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