Fix formatting of comments in "disable" chunks

CHANGES.md

@@ -8,6 +8,7 @@
 
 - Avoid adding breaks inside `~label:(fun` and base the indentation on the label. (#2271, @Julow)
 - Fix non-stabilizing comments attached to private/virtual/mutable keywords (#2272, @gpetiot)
+- Fix formatting of comments in "disable" chunks (#2279, @gpetiot)
 
 ### Changes
 

lib/Chunk.ml

@@ -11,22 +11,16 @@
 
 open Extended_ast
 
-type state = Enable | Disable of Location.t
+type 'a item =
+  | Structure : Extended_ast.structure item
+  | Signature : Extended_ast.signature item
+  | Use_file : Extended_ast.use_file item
 
 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.v)
-  && (update_conf ~quiet:true c [attr]).opr_opts.disable.v
-
-let enabling (c : Conf.t) attr =
-  c.opr_opts.disable.v
-  && not (update_conf ~quiet:true c [attr]).opr_opts.disable.v
+  { attr_loc: Location.t
+  ; chunk_loc: Location.t
+  ; state: [`Enable | `Disable]
+  ; items: 'a list }
 
 let init_loc =
   let pos =
@@ -35,34 +29,84 @@ let init_loc =
   in
   Location.{loc_ghost= false; loc_start= pos; loc_end= pos}
 
-let is_attr (type a) (fg : a list t) (x : a) =
+let is_attr (type a) (fg : a list item) (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
+let is_state_attr fg ~state x =
+  let open Option.Monad_infix in
+  is_attr fg x
+  >>= fun (attr, loc) ->
+  Conf.parse_state_attr attr
+  >>= fun new_state ->
+  match (state, new_state) with
+  | `Enable, `Disable -> Some (`Disable, loc)
+  | `Disable, `Enable -> Some (`Enable, 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)
+let last_loc (type a) (fg : a list item) (l : a list) =
+  let open Option.Monad_infix in
+  match fg with
+  | Structure -> List.last l >>| fun x -> x.pstr_loc
+  | Signature -> List.last l >>| fun x -> x.psig_loc
+  | Use_file -> (
+      List.last l
+      >>= function
+      | Ptop_def x -> List.last x >>| fun x -> x.pstr_loc
+      | Ptop_dir x -> Some x.pdir_loc )
+
+let mk ~attr_loc ~chunk_loc state items = {attr_loc; chunk_loc; state; items}
+
+let mk_tmp ~loc state items = mk ~attr_loc:loc ~chunk_loc:loc state items
+
+(* Build chunks of each disabled/enabled code spans. The [chunk_loc] of each
+   chunk has an unprecise ending position that needs to be set after looking
+   at the following chunk. *)
+let split_with_imprecise_locs fg ~state l =
+  let init = ([], state) in
+  let chunks, _ =
+    List.fold_left l ~init ~f:(fun (acc, state) x ->
+        match is_state_attr fg ~state x with
+        | Some (state, loc) -> (mk_tmp ~loc state [x] :: acc, state)
         | None -> (
           match acc with
-          | [] ->
-              let chunk =
-                if c.opr_opts.disable.v 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 state, lx -> (state, List.rev lx))
+          (* first chunk *)
+          | [] -> (mk_tmp ~loc:init_loc state [x] :: acc, state)
+          | chunk :: t -> ({chunk with items= x :: chunk.items} :: t, state)
+          ) )
+  in
+  List.rev_map chunks ~f:(fun x -> {x with items= List.rev x.items})
+
+(* Extend the [chunk_loc] to make it span until the start of [last_loc]. *)
+let extend_end_loc ~last_loc chunk =
+  let loc_end = last_loc.Location.loc_start in
+  let chunk_loc = {chunk.chunk_loc with loc_end} in
+  {chunk with chunk_loc}
+
+(* Update the [chunk_loc] of each chunk by using the loc of the following chunk.  *)
+let extend_end_locs fg l =
+  match List.rev l with
+  | [] -> []
+  | h :: t ->
+      (* last chunk *)
+      let init =
+        let last_loc =
+          Option.value (last_loc fg h.items) ~default:h.chunk_loc
+        in
+        let chunk_loc = {h.chunk_loc with loc_end= last_loc.loc_end} in
+        let h = {h with chunk_loc} in
+        (h.attr_loc, [h])
+      in
+      let _, chunks =
+        List.fold_left t ~init ~f:(fun (last_loc, acc) chunk ->
+            let chunk = extend_end_loc ~last_loc chunk in
+            (chunk.attr_loc, chunk :: acc) )
+      in
+      chunks
+
+let split ~state fg l =
+  extend_end_locs fg @@ split_with_imprecise_locs fg ~state l

lib/Chunk.mli

@@ -9,11 +9,16 @@
 (*                                                                        *)
 (**************************************************************************)
 
-type state = Enable | Disable of Location.t
+type 'a item =
+  | Structure : Extended_ast.structure item
+  | Signature : Extended_ast.signature item
+  | Use_file : Extended_ast.use_file item
 
 type 'a t =
-  | Structure : Extended_ast.structure t
-  | Signature : Extended_ast.signature t
-  | Use_file : Extended_ast.use_file t
+  { attr_loc: Location.t
+  ; chunk_loc: Location.t
+  ; state: [`Enable | `Disable]
+  ; items: 'a list }
 
-val split : 'a list t -> Conf.t -> 'a list -> (state * 'a list) list
+val split :
+  state:[`Enable | `Disable] -> 'a list item -> 'a list -> 'a t list

lib/Conf.ml

@@ -1523,29 +1523,37 @@ let parse_line config ?(version_check = config.opr_opts.version_check.v)
 
 open Parsetree
 
-let update ?(quiet = false) c {attr_name= {txt; loc}; attr_payload; _} =
-  let result =
-    match txt with
-    | "ocamlformat" -> (
-      match attr_payload with
-      | PStr
-          [ { pstr_desc=
-                Pstr_eval
-                  ( { pexp_desc=
-                        Pexp_constant
-                          {pconst_desc= Pconst_string (str, strloc, None); _}
-                    ; pexp_attributes= []
-                    ; _ }
-                  , [] )
-            ; _ } ] ->
-          parse_line ~from:(`Attribute strloc) c str
-          |> Result.map_error ~f:Error.to_string
-      | _ -> Error "Invalid format: String expected" )
-    | _ when String.is_prefix ~prefix:"ocamlformat." txt ->
-        Error
+let parse_attr {attr_name= {txt; loc= _}; attr_payload; _} =
+  match txt with
+  | "ocamlformat" -> (
+    match attr_payload with
+    | PStr
+        [ { pstr_desc=
+              Pstr_eval
+                ( { pexp_desc=
+                      Pexp_constant
+                        {pconst_desc= Pconst_string (str, strloc, None); _}
+                  ; pexp_attributes= []
+                  ; _ }
+                , [] )
+          ; _ } ] ->
+        Ok (str, strloc)
+    | _ -> Error (`Msg "Invalid format: String expected") )
+  | _ when String.is_prefix ~prefix:"ocamlformat." txt ->
+      Error
+        (`Msg
           (Format.sprintf "Invalid format: Unknown suffix %S"
-             (String.chop_prefix_exn ~prefix:"ocamlformat." txt) )
-    | _ -> Ok c
+             (String.chop_prefix_exn ~prefix:"ocamlformat." txt) ) )
+  | _ -> Error `Ignore
+
+let update ?(quiet = false) c ({attr_name= {txt; loc}; _} as attr) =
+  let result =
+    match parse_attr attr with
+    | Ok (str, strloc) ->
+        parse_line ~from:(`Attribute strloc) c str
+        |> Result.map_error ~f:Error.to_string
+    | Error (`Msg msg) -> Error msg
+    | Error `Ignore -> Ok c
   in
   match result with
   | Ok conf -> conf
@@ -1565,6 +1573,12 @@ let update_state c state =
   in
   {c with opr_opts}
 
+let parse_state_attr attr =
+  match parse_attr attr with
+  | Ok ("enable", _) -> Some `Enable
+  | Ok ("disable", _) -> Some `Disable
+  | _ -> None
+
 let print_config = Decl.print_config options
 
 let term = Decl.Store.to_term options

lib/Conf.mli

@@ -23,6 +23,8 @@ val update_value : t -> name:string -> value:string -> (t, Error.t) Result.t
 
 val update_state : t -> [`Enable | `Disable] -> t
 
+val parse_state_attr : Parsetree.attribute -> [`Enable | `Disable] option
+
 val parse_line :
      t
   -> ?version_check:bool

lib/Fmt_ast.ml

@@ -4324,43 +4324,46 @@ let fmt_repl_file c _ itms =
 module Chunk = struct
   open Chunk
 
-  let fmt_item (type a) (fg : a list t) : c -> Ast.t -> a list -> Fmt.t =
+  let fmt_item (type a) (fg : a list item) : 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
+    List.foldi chunks ~init:(c, noop, [])
+      ~f:(fun i (c, output, locs) chunk ->
+        let c = update_conf c chunk.state in
+        let output, locs =
+          match chunk.state with
+          | `Disable ->
+              let output =
+                output
+                $ Cmts.fmt_before c chunk.attr_loc ~eol:(fmt "\n@;<1000 0>")
+                $ fmt_if (i > 0) "\n@;<1000 0>"
+                $ str
+                    (String.strip
+                       (Source.string_at c.source chunk.chunk_loc) )
+              in
+              (output, chunk.chunk_loc :: locs)
+          | `Enable ->
+              let output =
+                output
+                $ fmt_if (i > 0) "@;<1000 0>"
+                $ fmt_item fg c ctx chunk.items
+              in
+              (output, locs)
+        in
+        (c, output, locs) )
+    |> fun ((_ : c), output, locs) ->
+    List.iter locs ~f:(Cmts.drop_inside c.cmts) ;
+    output
+
+  let split_and_fmt fg c ctx l =
+    let state = if c.conf.opr_opts.disable.v then `Disable else `Enable in
+    fmt fg c ctx @@ split fg l ~state
 end
 
 let fmt_file (type a) ~ctx ~fmt_code ~debug (fragment : a Extended_ast.t)

test/passing/tests/comments_around_disabled.ml

@@ -5,4 +5,13 @@ let () =
   ()
 [@@@ocamlformat "enable"]
 
+[@@@ocamlformat "disable"]
+ (* x *)
+ (* y *)
+let x =
+  x
+(* z *)
+
+[@@@ocamlformat "enable"]
+
 (* cmts *)

test/passing/tests/comments_around_disabled.ml.ref

@@ -5,4 +5,12 @@ let () =
   ()
 [@@@ocamlformat "enable"]
 
+[@@@ocamlformat "disable"]
+ (* x *)
+ (* y *)
+let x =
+  x
+(* z *)
+[@@@ocamlformat "enable"]
+
 (* cmts *)