A more comprehensive fix for comments on local syntax (#52)

* exclave and global should not drop comments Signed-off-by: alanechang <[email protected]> * keep attrs during local rewrite Signed-off-by: alanechang <[email protected]> --------- Signed-off-by: alanechang <[email protected]>
janestreet · Jan 2, 2024 · 98ce3cd · 98ce3cd
1 parent 3fb8b2b
commit 98ce3cd
Show file tree

Hide file tree

Showing 10 changed files with 248 additions and 53 deletions.
diff --git a/lib/Cmts.ml b/lib/Cmts.ml
@@ -81,6 +81,15 @@ let copy
   ; remaining
   ; layout_cache }
 
+(** returns comments that have not been formatted *)
+let remaining_comments t =
+  List.concat_map ~f:Multimap.to_list
+    [t.cmts_before; t.cmts_within; t.cmts_after]
+
+let remaining_before t loc = Map.find_multi t.cmts_before loc
+
+let remaining_locs t = Set.to_list t.remaining
+
 let restore src ~into =
   into.cmts_before <- src.cmts_before ;
   into.cmts_after <- src.cmts_after ;
@@ -321,6 +330,31 @@ let relocate (t : t) ~src ~before ~after =
         let s = Set.add s after in
         Set.add s before )
 
+let relocate_all_to_after (t : t) ~src ~after =
+  if t.debug then
+    Format.eprintf "relocate %a all to %a@\n%!" Location.fmt src Location.fmt
+      after ;
+  let merge_and_sort x y =
+    List.rev_append x y
+    |> List.sort ~compare:(Comparable.lift Location.compare_start ~f:Cmt.loc)
+  in
+  t.cmts_after <-
+    Map.change t.cmts_after after ~f:(fun r ->
+        let cmts = remaining_before t src in
+        match (r, cmts) with
+        | Some data, _ -> Some (merge_and_sort data cmts)
+        | None, _ :: _ -> Some cmts
+        | None, [] -> None ) ;
+  t.cmts_before <- Map.remove t.cmts_before src ;
+  update_cmts t `After
+    ~f:(Multimap.update_multi ~src ~dst:after ~f:merge_and_sort) ;
+  update_cmts t `Within
+    ~f:(Multimap.update_multi ~src ~dst:after ~f:merge_and_sort) ;
+  if t.debug then
+    update_remaining t ~f:(fun s ->
+        let s = Set.remove s src in
+        Set.add s after )
+
 let relocate_cmts_before (t : t) ~src ~sep ~dst =
   let f map =
     Multimap.partition_multi map ~src ~dst ~f:(fun cmt ->
@@ -719,12 +753,3 @@ let has_within t loc = pop_if_debug t loc ; Map.mem t.cmts_within loc
 let has_after t loc =
   pop_if_debug t loc ;
   Map.mem t.cmts_within loc || Map.mem t.cmts_after loc
-
-(** returns comments that have not been formatted *)
-let remaining_comments t =
-  List.concat_map ~f:Multimap.to_list
-    [t.cmts_before; t.cmts_within; t.cmts_after]
-
-let remaining_before t loc = Map.find_multi t.cmts_before loc
-
-let remaining_locs t = Set.to_list t.remaining
diff --git a/lib/Cmts.mli b/lib/Cmts.mli
@@ -37,6 +37,10 @@ val relocate :
     locations) comments before [src] to [before] and comments after [src] to
     [after]. *)
 
+val relocate_all_to_after : t -> src:Location.t -> after:Location.t -> unit
+(** [relocate_all_to_after src after] moves (changes the association with
+    locations) comments before and after [src] all to after [after]. *)
+
 val relocate_wrongfully_attached_cmts :
   t -> Source.t -> Extended_ast.expression -> unit
 (** [relocate_wrongfully_attached_cmts] relocates wrongfully attached

diff --git a/lib/Fmt_ast.ml b/lib/Fmt_ast.ml
@@ -891,7 +891,9 @@ and fmt_core_type c ?(box = true) ?pro ?(pro_space = true) ?constraint_ctx
   | Ptyp_arrow (args, ret_typ) ->
       Cmts.relocate c.cmts ~src:ptyp_loc
         ~before:(List.hd_exn args).pap_type.ptyp_loc ~after:ret_typ.ptyp_loc ;
-      let args, ret_typ, ctx = Sugar.decompose_arrow ctx args ret_typ in
+      let args, ret_typ, ctx =
+        Sugar.decompose_arrow c.cmts ctx args ret_typ
+      in
       let indent =
         match pro with
         | Some pro when c.conf.fmt_opts.ocp_indent_compat.v ->
@@ -1397,7 +1399,9 @@ and fmt_pattern_extension ~ext:_ c ~pro:_ ~parens:_ ~box:_ ~ctx0 ~ctx
 
 and fmt_fun_args c args =
   let fmt_fun_arg (a : function_param) =
-    let a = {a with pparam_desc= Sugar.remove_local_attrs a.pparam_desc} in
+    let a =
+      {a with pparam_desc= Sugar.remove_local_attrs c.cmts a.pparam_desc}
+    in
     let ctx = Fp a in
     Cmts.fmt c a.pparam_loc
     @@
@@ -3713,7 +3717,7 @@ and fmt_label_declaration c ctx ?(last = false) decl =
   let global_attr_opt, atrs = split_global_flags_from_attrs atrs in
   ( match global_attr_opt with
   | Some attr ->
-      Cmts.relocate c.cmts ~src:attr.attr_loc ~before:pld_type.ptyp_loc
+      Cmts.relocate_all_to_after c.cmts ~src:attr.attr_loc
         ~after:pld_type.ptyp_loc
   | None -> () ) ;
   hovbox 0
@@ -3776,7 +3780,7 @@ and fmt_core_type_gf c ctx typ =
   let global_attr_opt, _ = split_global_flags_from_attrs ptyp_attributes in
   ( match global_attr_opt with
   | Some attr ->
-      Cmts.relocate c.cmts ~src:attr.attr_loc ~before:typ.ptyp_loc
+      Cmts.relocate_all_to_after c.cmts ~src:attr.attr_loc
         ~after:typ.ptyp_loc
   | None -> () ) ;
   fmt_if (Option.is_some global_attr_opt) "global_ "

diff --git a/lib/Sugar.ml b/lib/Sugar.ml
@@ -18,13 +18,15 @@ let mk_function_param {Location.loc_start; _} {Location.loc_end; _} p =
   let pparam_loc = {Location.loc_start; loc_end; loc_ghost= true} in
   {pparam_desc= p; pparam_loc}
 
-let check_local_attr attrs =
+let check_local_attr_and_reloc_cmts cmts attrs loc =
   match
     List.partition_tf attrs ~f:(fun attr ->
         Conf.is_jane_street_local_annotation "local" ~test:attr.attr_name.txt )
   with
-  | [], _ -> (attrs, false)
-  | _ :: _, rest -> (rest, true)
+  | [local_attr], rest ->
+      Cmts.relocate_all_to_after cmts ~src:local_attr.attr_loc ~after:loc ;
+      (rest, true)
+  | _, _ -> (attrs, false)
 
 (* This function pulls apart an arrow type, pulling out local attributes into
    bools and producing a context without those attributes. This addresses the
@@ -33,16 +35,19 @@ let check_local_attr attrs =
    this to pass some internal ocamlformat sanity checks. It's not the
    cleanest solution in a vacuum, but is perhaps the one that will cause the
    fewest merge conflicts in the future. *)
-let decompose_arrow ctx ctl ct2 =
+let decompose_arrow cmts ctx ctl ct2 =
   let pull_out_local ap =
     let ptyp_attributes, local =
-      check_local_attr ap.pap_type.ptyp_attributes
+      check_local_attr_and_reloc_cmts cmts ap.pap_type.ptyp_attributes
+        ap.pap_type.ptyp_loc
     in
     ({ap with pap_type= {ap.pap_type with ptyp_attributes}}, local)
   in
   let args = List.map ~f:pull_out_local ctl in
   let ((res_ap, _) as res) =
-    let ptyp_attributes, local = check_local_attr ct2.ptyp_attributes in
+    let ptyp_attributes, local =
+      check_local_attr_and_reloc_cmts cmts ct2.ptyp_attributes ct2.ptyp_loc
+    in
     let ap =
       { pap_label= Nolabel
       ; pap_loc= ct2.ptyp_loc
@@ -102,7 +107,10 @@ let cl_fun ?(will_keep_first_ast_node = true) cmts xexp =
             Cmts.relocate cmts ~src:pcl_loc ~before ~after ;
           let xargs, xbody = fun_ (sub_cl ~ctx body) in
           let islocal, pattern =
-            match check_local_attr pattern.ppat_attributes with
+            match
+              check_local_attr_and_reloc_cmts cmts pattern.ppat_attributes
+                pattern.ppat_loc
+            with
             | _, false -> (false, pattern)
             | ppat_attributes, true -> (true, {pattern with ppat_attributes})
           in
@@ -113,12 +121,13 @@ let cl_fun ?(will_keep_first_ast_node = true) cmts xexp =
   in
   fun_ ~will_keep_first_ast_node xexp
 
-let remove_local_attrs param =
+let remove_local_attrs cmts param =
   match param with
   | Pparam_newtype _ -> param
   | Pparam_val (_, label, default, pattern) ->
       let ppat_attributes, is_local =
-        check_local_attr pattern.ppat_attributes
+        check_local_attr_and_reloc_cmts cmts pattern.ppat_attributes
+          pattern.ppat_loc
       in
       Pparam_val (is_local, label, default, {pattern with ppat_attributes})
 
@@ -329,7 +338,10 @@ module Let_binding = struct
       *)
   let local_pattern_can_be_sugared pvb_pat pvb_constraint exp_loc cmts =
     (* If the original code was sugared, preserve that always. *)
-    let _, already_sugared = check_local_attr pvb_pat.ppat_attributes in
+    let _, already_sugared =
+      check_local_attr_and_reloc_cmts cmts pvb_pat.ppat_attributes
+        pvb_pat.ppat_loc
+    in
     (* Don't wipe away comments before [local_]. *)
     let comment_before = Cmts.has_before cmts exp_loc in
     already_sugared
@@ -369,11 +381,17 @@ module Let_binding = struct
               local_pattern_can_be_sugared pvb_pat pvb_constraint
                 pvb_expr.pexp_loc cmts
             then
-              let sattrs, _ = check_local_attr sbody.pexp_attributes in
+              let sattrs, _ =
+                check_local_attr_and_reloc_cmts cmts sbody.pexp_attributes
+                  sbody.pexp_loc
+              in
               (true, {sbody with pexp_attributes= sattrs})
             else (false, pvb_expr)
           in
-          let pattrs, _ = check_local_attr pvb_pat.ppat_attributes in
+          let pattrs, _ =
+            check_local_attr_and_reloc_cmts cmts pvb_pat.ppat_attributes
+              pvb_pat.ppat_loc
+          in
           let pat = {pvb_pat with ppat_attributes= pattrs} in
           let fake_ctx =
             Lb

diff --git a/lib/Sugar.mli b/lib/Sugar.mli
@@ -14,7 +14,8 @@ open Asttypes
 open Extended_ast
 
 val decompose_arrow :
-     Ast.t
+     Cmts.t
+  -> Ast.t
   -> arrow_param list
   -> core_type
   -> (arrow_param * bool) list * (arrow_param * bool) * Ast.t
@@ -42,7 +43,7 @@ val cl_fun :
     and the body of the function [exp]. [will_keep_first_ast_node] is set by
     default, otherwise the [exp] is returned without modification. *)
 
-val remove_local_attrs : function_param_desc -> function_param_desc
+val remove_local_attrs : Cmts.t -> function_param_desc -> function_param_desc
 
 module Exp : sig
   val infix :

diff --git a/test/passing/tests/local.ml b/test/passing/tests/local.ml
@@ -115,22 +115,30 @@ let[@local always] upstream_local_attr_always_short x = x
 
 let[@local maybe] upstream_local_attr_maybe_short x = x
 
-let f x = (* a *) local_
+let f x = (* a *) (* b *) local_ (* c *) (* d *)
   let y = 1 in
   x + y
 
-let f x = (* a *) exclave_
+let f x = (* a *) (* b *) exclave_ (* c *) (* d *)
   let y = 1 in
   x + y
 
-let x = (* a *) local_
+let x = (* a *) (* b *) local_ (* c *) (* d *)
   let y = 1 in
   y
 
-let x = (* a *) exclave_
+let x = (* a *) (* b *) exclave_ (* c *) (* d *)
   let y = 1 in
   y
 
 module type S = S -> S -> S
 (* this is here to make sure we pass the AST equality checks even when the
    extended AST is different *)
+
+let f ((* a *) (* b *)local_ (* c *) (* d *)a) ~foo:((* e *) (* f *)local_(* g *) (* h *) b) ?foo:(local_ c = 1) ~(local_ d) = ()
+type 'a r = {mutable a: 'a; b: 'a; (* a *) (* b *)global_ (* c *) (* d *)c: 'a}
+
+type 'a r =
+  | Foo of (* a *) (* b *)global_(* c *) (* d *) 'a
+  | Bar of 'a * (* e *) (* f *)global_ (* g *) (* h *)'a
+  | Baz of global_ int * string * (* i *) (* j *) global_ (* k *) (* l *)'a
diff --git a/test/passing/tests/local.ml.ref b/test/passing/tests/local.ml.ref
@@ -124,28 +124,85 @@ let[@local maybe] upstream_local_attr_maybe_short x = x
 
 let f x =
   (* a *)
+  (* b *)
   local_
+  (* c *)
+  (* d *)
   let y = 1 in
   x + y
 
 let f x =
   (* a *)
+  (* b *)
   exclave_
+  (* c *)
+  (* d *)
   let y = 1 in
   x + y
 
 let x =
   (* a *)
+  (* b *)
   local_
+  (* c *)
+  (* d *)
   let y = 1 in
   y
 
 let x =
   (* a *)
+  (* b *)
   exclave_
+  (* c *)
+  (* d *)
   let y = 1 in
   y
 
 module type S = functor (_ : S) (_ : S) -> S
 (* this is here to make sure we pass the AST equality checks even when the
    extended AST is different *)
+
+let f
+    (local_ (* a *)
+            (* b *)
+            (* c *)
+            (* d *)
+      a )
+    ~foo:(local_ (* e *)
+                 (* f *)
+                 (* g *)
+                 (* h *)
+      b ) ?foo:(local_ c = 1) ~(local_ d) =
+  ()
+
+type 'a r =
+  { mutable a: 'a
+  ; b: 'a
+  ; (* a *)
+    (* b *)
+    global_ (* c *)
+            (* d *) c:
+      'a }
+
+type 'a r =
+  | Foo of
+      global_ (* a *)
+              (* b *)
+              (* c *)
+              (* d *)
+      'a
+  | Bar of
+      'a
+      * global_ (* e *)
+                (* f *)
+                (* g *)
+                (* h *)
+      'a
+  | Baz of
+      global_ int
+      * string
+      * global_ (* i *)
+                (* j *)
+                (* k *)
+                (* l *)
+      'a
diff --git a/test/passing/tests/local_rewrite_regressions.ml b/test/passing/tests/local_rewrite_regressions.ml
@@ -6,3 +6,20 @@ module With_length : sig
     ; length : int [@global]
     }
 end = struct end
+
+val find_last : 'a t -> f:(('a -> bool)(* a *)(* b *)[@local](* c *)) -> 'a option
+let find_last : 'a t -> f:(('a -> bool)(* a *)[@local](* b *)) -> 'a option = assert false
+type t = (string[@local]) -> (string(* a *)[@local](* b *))
+
+
+type global_long_attrs =
+  | Foo of { s : string(* a *)(* b *)[@ocaml.global](* c *)(* d *); b: int }
+  | Bar of (string(* e *)(* f *)[@ocaml.global](* g *)(* h *))
+
+let local_long_ext = (* a *)(* b *)[%ocaml.local](* c *)(* d *) ()
+
+let () =
+  let g = (* a *)[%local](* b *) (fun a b c -> 1) in
+  ()
+
+let f (x(* a *)[@local](* b *)) = x