Preserve position of comments located after the semi-colon of the las…

…t element of lists/arrays/records (ocaml-ppx#2032)
antalsz · May 19, 2022 · 7ef6284 · 7ef6284
1 parent c3381ff
commit 7ef6284
Show file tree

Hide file tree

Showing 14 changed files with 275 additions and 23 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -29,6 +29,7 @@
 
 - New syntax `(*= ... *)` for verbatim comments (#2028, @gpetiot)
 - Preserve the begin-end construction in the AST (#1785, @hhugo, @gpetiot)
+- Preserve position of comments located after the semi-colon of the last element of lists/arrays/records (#2032, @gpetiot)
 
 ## 0.21.0 (2022-02-25)
 

diff --git a/lib/Ast.ml b/lib/Ast.ml
@@ -282,6 +282,8 @@ module Ext = struct
 end
 
 module Exp = struct
+  let location x = x.pexp_loc
+
   let test_id ~f = function
     | {pexp_desc= Pexp_ident {txt= i; _}; _} -> f i
     | _ -> false
@@ -356,6 +358,8 @@ module Exp = struct
 end
 
 module Pat = struct
+  let location x = x.ppat_loc
+
   let is_simple {ppat_desc; _} =
     match ppat_desc with
     | Ppat_any | Ppat_constant _ | Ppat_var _

diff --git a/lib/Ast.mli b/lib/Ast.mli
@@ -87,6 +87,8 @@ module Longident : sig
 end
 
 module Exp : sig
+  val location : expression -> Location.t
+
   val is_prefix : expression -> bool
   (** Holds for prefix symbol expressions. *)
 
@@ -145,6 +147,8 @@ type cmt_checker =
   ; cmts_after: Location.t -> bool }
 
 module Pat : sig
+  val location : pattern -> Location.t
+
   val is_simple : pattern -> bool
 end
 

diff --git a/lib/Cmts.ml b/lib/Cmts.ml
@@ -442,11 +442,12 @@ let preserve ~cache_key f t =
 let pop_if_debug t loc =
   if t.debug then update_remaining t ~f:(fun s -> Set.remove s loc)
 
-let find_cmts t pos loc =
+let find_cmts ?(filter = Fn.const true) t pos loc =
   pop_if_debug t loc ;
-  let r = find_at_position t loc pos in
-  update_cmts t pos ~f:(fun m -> Map.remove m loc) ;
-  r
+  Option.map (find_at_position t loc pos) ~f:(fun cmts ->
+      let picked, not_picked = List.partition_tf cmts ~f:filter in
+      update_cmts t pos ~f:(Map.set ~key:loc ~data:not_picked) ;
+      picked )
 
 let break_comment_group source margin {Cmt.loc= a; _} {Cmt.loc= b; _} =
   let vertical_align =
@@ -647,14 +648,15 @@ let fmt_before t conf ~fmt_code ?pro ?(epi = Fmt.break 1 0) ?eol ?adj loc =
   fmt_cmts t conf (find_cmts t `Before loc) ~fmt_code ?pro ~epi ?eol ?adj loc
     Before
 
-let fmt_after t conf ~fmt_code ?(pro = Fmt.break 1 0) ?epi loc =
+let fmt_after t conf ~fmt_code ?(pro = Fmt.break 1 0) ?epi ?filter loc =
   let open Fmt in
   let within =
-    fmt_cmts t conf (find_cmts t `Within loc) ~fmt_code ~pro ?epi loc Within
+    let cmts = find_cmts ?filter t `Within loc in
+    fmt_cmts t conf cmts ~fmt_code ~pro ?epi loc Within
   in
   let after =
-    fmt_cmts t conf (find_cmts t `After loc) ~fmt_code ~pro ?epi ~eol:noop
-      loc After
+    let cmts = find_cmts ?filter t `After loc in
+    fmt_cmts t conf cmts ~fmt_code ~pro ?epi ~eol:noop loc After
   in
   within $ after
 

diff --git a/lib/Cmts.mli b/lib/Cmts.mli
@@ -65,6 +65,7 @@ val fmt_after :
   -> fmt_code:(Conf.t -> Fmt.code_formatter)
   -> ?pro:Fmt.t
   -> ?epi:Fmt.t
+  -> ?filter:(Cmt.t -> bool)
   -> Location.t
   -> Fmt.t
 (** [fmt_after loc] formats the comments associated with [loc] that appear

diff --git a/lib/Fmt_ast.ml b/lib/Fmt_ast.ml
@@ -97,25 +97,40 @@ let protect =
 let update_config ?quiet c l =
   {c with conf= List.fold ~init:c.conf l ~f:(Conf.update ?quiet)}
 
-let fmt_elements_collection ?(first_sep = true) ?(last_sep = true)
-    (p : Params.elements_collection) fmt_x xs =
+(* Preserve the position of comments located after the last element of a
+   list/array (after `;`), otherwise comments are picked up by
+   `fmt_expression` and printed before `;`. *)
+let collection_last_cmt ?pro c (loc : Location.t) locs =
+  let filter = function Parser.SEMI -> true | _ -> false in
+  opt (List.last locs) (fun (last : Location.t) ->
+      match
+        Source.tokens_between c.source last.loc_end loc.loc_end ~filter
+      with
+      | [] -> noop
+      | (_, semicolon_loc) :: _ ->
+          Cmts.fmt_after ?pro c last ~filter:(fun Cmt.{loc; _} ->
+              Location.compare loc semicolon_loc >= 0 ) )
+
+let fmt_elements_collection ?pro ?(first_sep = true) ?(last_sep = true) c
+    (p : Params.elements_collection) f loc fmt_x xs =
   let fmt_one ~first ~last x =
     fmt_if_k (not (first && first_sep)) p.sep_before
     $ fmt_x x
     $ fmt_or_k (last && last_sep) p.sep_after_final p.sep_after_non_final
   in
-  list_fl xs fmt_one
+  list_fl xs fmt_one $ collection_last_cmt ?pro c loc (List.map ~f xs)
 
-let fmt_expressions c width sub_exp exprs fmt_expr p =
+let fmt_expressions c width sub_exp exprs fmt_expr p loc =
   match c.conf.fmt_opts.break_collection_expressions with
-  | `Fit_or_vertical -> fmt_elements_collection p fmt_expr exprs
+  | `Fit_or_vertical ->
+      fmt_elements_collection c p Exp.location loc fmt_expr exprs
   | `Wrap ->
       let is_simple x = is_simple c.conf width (sub_exp x) in
       let break x1 x2 = not (is_simple x1 && is_simple x2) in
       let grps = List.group exprs ~break in
       let fmt_grp ~first:first_grp ~last:last_grp exprs =
-        fmt_elements_collection ~first_sep:first_grp ~last_sep:last_grp p
-          fmt_expr exprs
+        fmt_elements_collection c ~first_sep:first_grp ~last_sep:last_grp p
+          Exp.location loc fmt_expr exprs
       in
       list_fl grps fmt_grp
 
@@ -1069,9 +1084,12 @@ and fmt_pattern ?ext c ?pro ?parens ?(box = false)
       let last_sep, fmt_underscore =
         match closed_flag with
         | Closed -> (true, noop)
-        | Open loc -> (false, Cmts.fmt c loc p2.wildcard)
+        | Open loc -> (false, Cmts.fmt ~pro:(break 1 2) c loc p2.wildcard)
+      in
+      let fmt_fields =
+        fmt_elements_collection c ~last_sep p1 (snd >> Pat.location) ppat_loc
+          fmt_field flds
       in
-      let fmt_fields = fmt_elements_collection ~last_sep p1 fmt_field flds in
       hvbox_if parens 0
         (Params.parens_if parens c.conf
            (p1.box (fmt_fields $ fmt_underscore)) )
@@ -1081,13 +1099,13 @@ and fmt_pattern ?ext c ?pro ?parens ?(box = false)
   | Ppat_array pats ->
       let p = Params.get_array_pat c.conf ~ctx:ctx0 in
       p.box
-        (fmt_elements_collection p
-           (fun pat -> fmt_pattern c (sub_pat ~ctx pat))
+        (fmt_elements_collection c p Pat.location ppat_loc
+           (sub_pat ~ctx >> fmt_pattern c >> hvbox 0)
            pats )
   | Ppat_list pats ->
       let p = Params.get_list_pat c.conf ~ctx:ctx0 in
       p.box
-        (fmt_elements_collection p
+        (fmt_elements_collection c p Pat.location ppat_loc
            (sub_pat ~ctx >> fmt_pattern c >> hvbox 0)
            pats )
   | Ppat_or _ ->
@@ -2039,7 +2057,7 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
         ( p.box
             (fmt_expressions c (expression_width c) (sub_exp ~ctx) e1N
                (sub_exp ~ctx >> fmt_expression c)
-               p )
+               p pexp_loc )
         $ fmt_atrs )
   | Pexp_list e1N ->
       let p = Params.get_list_expr c.conf in
@@ -2054,7 +2072,7 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
                   (fun e ->
                     let fmt_cmts = Cmts.fmt c ~eol:cmt_break e.pexp_loc in
                     fmt_cmts @@ (sub_exp ~ctx >> fmt_expression c) e )
-                  p )
+                  p pexp_loc )
            $ fmt_atrs ) )
   | Pexp_assert e0 ->
       let paren_body, wrap_symbol =
@@ -2453,7 +2471,10 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
               @@ fmt_record_field c ~rhs:(fmt_rhs f) lid1 )
       in
       let p1, p2 = Params.get_record_expr c.conf in
-      let fmt_fields = fmt_elements_collection p1 fmt_field flds in
+      let fmt_fields =
+        fmt_elements_collection c p1 (snd >> Exp.location) pexp_loc fmt_field
+          flds ~pro:(break 1 2)
+      in
       hvbox_if has_attr 0
         ( p1.box
             ( opt default (fun d ->

diff --git a/test/passing/dune.inc b/test/passing/dune.inc
@@ -863,6 +863,60 @@
  (package ocamlformat)
  (action (diff tests/coerce.ml.err coerce.ml.stderr)))
 
+(rule
+ (deps tests/.ocamlformat )
+ (package ocamlformat)
+ (action
+  (with-stdout-to collections-conventional.ml.stdout
+   (with-stderr-to collections-conventional.ml.stderr
+     (run %{bin:ocamlformat} --margin-check --profile=conventional --max-iters=3 %{dep:tests/collections.ml})))))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/collections-conventional.ml.ref collections-conventional.ml.stdout)))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/collections-conventional.ml.err collections-conventional.ml.stderr)))
+
+(rule
+ (deps tests/.ocamlformat )
+ (package ocamlformat)
+ (action
+  (with-stdout-to collections-janestreet.ml.stdout
+   (with-stderr-to collections-janestreet.ml.stderr
+     (run %{bin:ocamlformat} --margin-check --profile=janestreet --max-iters=3 %{dep:tests/collections.ml})))))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/collections-janestreet.ml.ref collections-janestreet.ml.stdout)))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/collections-janestreet.ml.err collections-janestreet.ml.stderr)))
+
+(rule
+ (deps tests/.ocamlformat )
+ (package ocamlformat)
+ (action
+  (with-stdout-to collections.ml.stdout
+   (with-stderr-to collections.ml.stderr
+     (run %{bin:ocamlformat} --margin-check --profile=ocamlformat --max-iters=3 %{dep:tests/collections.ml})))))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/collections.ml.ref collections.ml.stdout)))
+
+(rule
+ (alias runtest)
+ (package ocamlformat)
+ (action (diff tests/collections.ml.err collections.ml.stderr)))
+
 (rule
  (deps tests/.ocamlformat )
  (package ocamlformat)

diff --git a/test/passing/tests/collections-conventional.ml.opts b/test/passing/tests/collections-conventional.ml.opts
@@ -0,0 +1,2 @@
+--profile=conventional
+--max-iters=3
diff --git a/test/passing/tests/collections-conventional.ml.ref b/test/passing/tests/collections-conventional.ml.ref
@@ -0,0 +1,48 @@
+let _ =
+  [
+    a;
+    b
+    (* foooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo *);
+  ]
+
+let [
+      fooooooooooooooooooooooooooooooo;
+      fooooooooooooooooooooooooooooooo;
+      fooooooooooooooooooooooooooooooo;
+      (* before end of the list *)
+    ] =
+  [
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    (* after all elements *)
+    (* after all elements as well *)
+  ]
+
+let [|
+      fooooooooooooooooooooooooooooooo;
+      fooooooooooooooooooooooooooooooo;
+      fooooooooooooooooooooooooooooooo;
+      (* before end of the array *)
+    |] =
+  [|
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    (* after all elements *)
+    (* after all elements as well *)
+  |]
+
+let {
+      fooooooooooooooooooooooooooooooo;
+      fooooooooooooooooooooooooooooooo;
+      fooooooooooooooooooooooooooooooo;
+      _;
+      (* xxx *)
+    } =
+  {
+    fooooooooooooooooooooooooooooooo = x;
+    fooooooooooooooooooooooooooooooo = y;
+    fooooooooooooooooooooooooooooooo = z;
+    (* after all fields *)
+  }
diff --git a/test/passing/tests/collections-janestreet.ml.opts b/test/passing/tests/collections-janestreet.ml.opts
@@ -0,0 +1,2 @@
+--profile=janestreet
+--max-iters=3
diff --git a/test/passing/tests/collections-janestreet.ml.ref b/test/passing/tests/collections-janestreet.ml.ref
@@ -0,0 +1,42 @@
+let _ =
+  [ a; b (* foooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo *) ]
+;;
+
+let [ fooooooooooooooooooooooooooooooo
+    ; fooooooooooooooooooooooooooooooo
+    ; fooooooooooooooooooooooooooooooo (* before end of the list *)
+    ]
+  =
+  [ fooooooooooooooooooooooooooooooo
+  ; fooooooooooooooooooooooooooooooo
+  ; fooooooooooooooooooooooooooooooo
+    (* after all elements *)
+    (* after all elements as well *)
+  ]
+;;
+
+let [| fooooooooooooooooooooooooooooooo
+     ; fooooooooooooooooooooooooooooooo
+     ; fooooooooooooooooooooooooooooooo (* before end of the array *)
+    |]
+  =
+  [| fooooooooooooooooooooooooooooooo
+   ; fooooooooooooooooooooooooooooooo
+   ; fooooooooooooooooooooooooooooooo
+     (* after all elements *)
+     (* after all elements as well *)
+  |]
+;;
+
+let { fooooooooooooooooooooooooooooooo
+    ; fooooooooooooooooooooooooooooooo
+    ; fooooooooooooooooooooooooooooooo
+    ; _
+      (* xxx *)
+    }
+  =
+  { fooooooooooooooooooooooooooooooo = x
+  ; fooooooooooooooooooooooooooooooo = y
+  ; fooooooooooooooooooooooooooooooo = z (* after all fields *)
+  }
+;;
diff --git a/test/passing/tests/collections.ml b/test/passing/tests/collections.ml
@@ -0,0 +1,37 @@
+let _ =
+  [
+    a;
+    b (* foooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo *)
+  ]
+
+let
+  [ fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+  (* before end of the list *) ] =
+  [ fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+  (* after all elements *)
+  (* after all elements as well *) ]
+
+let
+  [| fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+  (* before end of the array *) |] =
+  [| fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+  (* after all elements *)
+  (* after all elements as well *) |]
+
+let { fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    fooooooooooooooooooooooooooooooo;
+    _
+    (* xxx *) } =
+  { fooooooooooooooooooooooooooooooo= x;
+    fooooooooooooooooooooooooooooooo= y;
+    fooooooooooooooooooooooooooooooo= z;
+  (* after all fields *) }
diff --git a/test/passing/tests/collections.ml.opts b/test/passing/tests/collections.ml.opts
@@ -0,0 +1,2 @@
+--profile=ocamlformat
+--max-iters=3