track local rewrites and fix ast equality check

Signed-off-by: alanechang <alanechang@janestreet.com>

lib/Conf.ml

@@ -260,8 +260,7 @@ let default =
       ; ocaml_version= elt Ocaml_version.Releases.v4_04_0
       ; quiet= elt false
       ; disable_conf_attrs= elt false
-      ; version_check= elt true
-      ; rewrite_old_style_jane_street_local_annotations= elt false } }
+      ; version_check= elt true } }
 
 module V = struct
   let v0_12 = Version.make ~major:0 ~minor:12 ~patch:None
@@ -1454,23 +1453,6 @@ module Operational = struct
       (fun conf elt -> update conf ~f:(fun f -> {f with version_check= elt}))
       (fun conf -> conf.opr_opts.version_check)
 
-  let rewrite_old_style_jane_street_local_annotations =
-    let doc =
-      "Rewrite all Jane Street annotations for use with the local mode, \
-       such as \"[%local]\" or \"[@ocaml.global]\", into their \
-       pretty-printed syntactic form, such as \"local_\" or \"global_\".  \
-       THIS OPTION WILL CHANGE THE RESULTING AST."
-    in
-    Decl.flag ~default
-      ~names:["rewrite-old-style-jane-street-local-annotations"]
-      ~doc ~kind
-      (fun conf elt ->
-        update conf ~f:(fun f ->
-            {f with rewrite_old_style_jane_street_local_annotations= elt} )
-        )
-      (fun conf ->
-        conf.opr_opts.rewrite_old_style_jane_street_local_annotations )
-
   let options : Store.t =
     Store.
       [ elt comment_check
@@ -1481,8 +1463,7 @@ module Operational = struct
       ; elt ocaml_version
       ; elt quiet
       ; elt disable_conf_attrs
-      ; elt version_check
-      ; elt rewrite_old_style_jane_street_local_annotations ]
+      ; elt version_check ]
 end
 
 let options = Operational.options @ Formatting.options @ options
@@ -1590,12 +1571,12 @@ let parse_state_attr attr =
   | Ok ("disable", _) -> Some `Disable
   | _ -> None
 
-let is_jane_street_local_annotation config name ~test =
-  String.equal test ("extension." ^ name)
-  ||
-  if config.opr_opts.rewrite_old_style_jane_street_local_annotations.v then
+let is_jane_street_local_annotation _config name ~test =
+  let is_legacy =
     String.equal test name || String.equal test ("ocaml." ^ name)
-  else false
+  in
+  if is_legacy then Erase_jane_syntax.set_local_rewrite_occurred true ;
+  is_legacy || String.equal test ("extension." ^ name)
 
 let print_config = Decl.print_config options
 

lib/Conf.mli

@@ -26,6 +26,10 @@ val update_state : t -> [`Enable | `Disable] -> t
 val parse_state_attr : Parsetree.attribute -> [`Enable | `Disable] option
 
 val is_jane_street_local_annotation : t -> string -> test:string -> bool
+(** [is_jane_street_local_annotation c n ~test] checks if [test] refers to
+    a piece of jane street local annotation with the name [c]. This function
+    contains side effects: when a legacy piece of local annotation is encountered,
+    it sets [Erase_jane_syntax.local_rewrite_occurred ()] to [true]. *)
 
 val parse_line :
      t

lib/Fmt_ast.ml

@@ -751,6 +751,7 @@ and type_constr_and_body c xbody =
 
 and fmt_arrow_param c ctx
     ({pap_label= lI; pap_loc= locI; pap_type= tI}, localI) =
+  let localI = localI && not (Erase_jane_syntax.should_erase ()) in
   let arg_label lbl =
     match lbl with
     | Nolabel -> if localI then Some (str "local_ ") else None
@@ -1478,7 +1479,7 @@ and fmt_body c ?ext ({ast= body; _} as xbody) =
            ~test:extension_local
          (* Don't wipe away comments before [local_]. *)
          && not (Cmts.has_before c.cmts pexp_loc) ->
-      ( fmt " local_"
+      ( fmt_if (not (Erase_jane_syntax.should_erase ())) " local_"
       , fmt_expression c ~eol:(fmt "@;<1000 0>") (sub_exp c.conf ~ctx sbody)
       )
   | { pexp_desc=
@@ -1492,7 +1493,7 @@ and fmt_body c ?ext ({ast= body; _} as xbody) =
            ~test:extension_exclave
          (* Don't wipe away comments before [exclave_]. *)
          && not (Cmts.has_before c.cmts pexp_loc) ->
-      ( fmt " exclave_"
+      ( fmt_if (not (Erase_jane_syntax.should_erase ())) " exclave_"
       , fmt_expression c ~eol:(fmt "@;<1000 0>") (sub_exp c.conf ~ctx sbody)
       )
   | _ -> (noop, fmt_expression c ~eol:(fmt "@;<1000 0>") xbody)
@@ -2121,15 +2122,17 @@ and fmt_expression c ?(box = true) ?(pro = noop) ?eol ?parens
            ~test:extension_local ->
       pro
       $ Params.parens_if parens c.conf
-          (fmt "local_@ " $ fmt_expression c (sub_exp c.conf ~ctx sbody))
+          ( fmt_if (not (Erase_jane_syntax.should_erase ())) "local_@ "
+          $ fmt_expression c (sub_exp c.conf ~ctx sbody) )
   | Pexp_apply
       ( {pexp_desc= Pexp_extension ({txt= extension_exclave; _}, PStr []); _}
       , [(Nolabel, sbody)] )
     when Conf.is_jane_street_local_annotation c.conf "exclave"
            ~test:extension_exclave ->
       pro
       $ Params.parens_if parens c.conf
-          (fmt "exclave_@ " $ fmt_expression c (sub_exp c.conf ~ctx sbody))
+          ( fmt_if (not (Erase_jane_syntax.should_erase ())) "exclave_@ "
+          $ fmt_expression c (sub_exp c.conf ~ctx sbody) )
   | Pexp_prefix (op, e) ->
       let has_cmts = Cmts.has_before c.cmts e.pexp_loc in
       pro
@@ -3666,6 +3669,7 @@ and fmt_label_declaration c ctx ?(last = false) decl =
           (str ";")
   in
   let is_global, atrs = split_global_flags_from_attrs c.conf atrs in
+  let is_global = is_global && not (Erase_jane_syntax.should_erase ()) in
   hovbox 0
     ( Cmts.fmt_before c pld_loc
     $ hvbox
@@ -3724,6 +3728,7 @@ and fmt_constructor_declaration c ctx ~first ~last:_ cstr_decl =
 and fmt_core_type_gf c ctx typ =
   let {ptyp_attributes; _} = typ in
   let is_global, _ = split_global_flags_from_attrs c.conf ptyp_attributes in
+  let is_global = is_global && not (Erase_jane_syntax.should_erase ()) in
   fmt_if is_global "global_ " $ fmt_core_type c (sub_typ ~ctx typ)
 
 and fmt_constructor_arguments ?vars c ctx ~pre = function
@@ -4802,7 +4807,12 @@ and fmt_value_binding c ~rec_flag ?ext ?in_ ?epi
                                   $ fmt_extension_suffix c ext
                                   $ fmt_attributes c at_attrs
                                   $ fmt_if rec_flag " rec"
-                                  $ fmt_if lb_local " local_"
+                                  $ fmt_if
+                                      ( lb_local
+                                      && not
+                                           (Erase_jane_syntax.should_erase ())
+                                      )
+                                      " local_"
                                   $ fmt_or pat_has_cmt "@ " " "
                                   $ fmt_pattern c lb_pat )
                               $ fmt_if_k

lib/Normalize_std_ast.ml

@@ -12,13 +12,27 @@
 open Parser_standard
 open Std_ast
 
+let erase_legacy_jane_street_local_annotations c : attributes -> attributes =
+  List.filter ~f:(fun attr ->
+      match attr with
+      | {attr_name= {txt= old_name; _}; attr_payload= PStr []; _} ->
+          not
+            ( Conf.is_jane_street_local_annotation c "local" ~test:old_name
+            || Conf.is_jane_street_local_annotation c "global" ~test:old_name
+            || Conf.is_jane_street_local_annotation c "exclave"
+                 ~test:old_name )
+      | _ -> true )
+
 let is_doc = function
   | {attr_name= {Location.txt= "ocaml.doc" | "ocaml.text"; _}; _} -> true
   | _ -> false
 
 let is_erasable_jane_syntax attr =
   String.is_prefix ~prefix:"jane.erasable." attr.attr_name.txt
 
+let is_local_jane_syntax attr =
+  String.is_prefix ~prefix:"jane.erasable.local" attr.attr_name.txt
+
 let dedup_cmts fragment ast comments =
   let of_ast ast =
     let docs = ref (Set.empty (module Cmt)) in
@@ -69,7 +83,8 @@ let docstring (c : Conf.t) =
 let sort_attributes : attributes -> attributes =
   List.sort ~compare:Poly.compare
 
-let make_mapper conf ~ignore_doc_comments ~erase_jane_syntax =
+let make_mapper conf ~ignore_doc_comments ~erase_jane_syntax
+    ~local_rewrite_occurred =
   let open Ast_helper in
   (* remove locations *)
   let location _ _ = Location.none in
@@ -108,6 +123,11 @@ let make_mapper conf ~ignore_doc_comments ~erase_jane_syntax =
         List.filter atrs ~f:(fun a -> not (is_erasable_jane_syntax a))
       else atrs
     in
+    let atrs =
+      if local_rewrite_occurred then
+        List.filter atrs ~f:(fun a -> not (is_local_jane_syntax a))
+      else atrs
+    in
     let atrs =
       if ignore_doc_comments then
         List.filter atrs ~f:(fun a -> not (is_doc a))
@@ -132,10 +152,27 @@ let make_mapper conf ~ignore_doc_comments ~erase_jane_syntax =
           (Exp.sequence ~loc:loc1 ~attrs:attrs1
              (Exp.sequence ~loc:loc2 ~attrs:attrs2 exp1 exp2)
              exp3 )
+    | Pexp_apply
+        ( {pexp_desc= Pexp_extension ({txt= extension_name; _}, PStr []); _}
+        , [(Nolabel, sbody)] )
+      when (local_rewrite_occurred || erase_jane_syntax)
+           && ( Conf.is_jane_street_local_annotation conf "local"
+                  ~test:extension_name
+              || Conf.is_jane_street_local_annotation conf "exclave"
+                   ~test:extension_name ) ->
+        m.expr m sbody
     | _ -> Ast_mapper.default_mapper.expr m exp
   in
   let pat (m : Ast_mapper.mapper) pat =
     let pat = {pat with ppat_loc_stack= []} in
+    let pat =
+      if local_rewrite_occurred || erase_jane_syntax then
+        { pat with
+          ppat_attributes=
+            erase_legacy_jane_street_local_annotations conf
+              pat.ppat_attributes }
+      else pat
+    in
     let {ppat_desc; ppat_loc= loc1; ppat_attributes= attrs1; _} = pat in
     (* normalize nested or patterns *)
     match ppat_desc with
@@ -158,6 +195,14 @@ let make_mapper conf ~ignore_doc_comments ~erase_jane_syntax =
   in
   let typ (m : Ast_mapper.mapper) typ =
     let typ = {typ with ptyp_loc_stack= []} in
+    let typ =
+      if local_rewrite_occurred || erase_jane_syntax then
+        { typ with
+          ptyp_attributes=
+            erase_legacy_jane_street_local_annotations conf
+              typ.ptyp_attributes }
+      else typ
+    in
     Ast_mapper.default_mapper.typ m typ
   in
   let structure =
@@ -196,6 +241,17 @@ let make_mapper conf ~ignore_doc_comments ~erase_jane_syntax =
       Ast_mapper.default_mapper.class_signature m {x with pcsig_fields}
     else Ast_mapper.default_mapper.class_signature
   in
+  let label_declaration (m : Ast_mapper.mapper) ld =
+    let ld =
+      if local_rewrite_occurred || erase_jane_syntax then
+        { ld with
+          pld_attributes=
+            erase_legacy_jane_street_local_annotations conf ld.pld_attributes
+        }
+      else ld
+    in
+    Ast_mapper.default_mapper.label_declaration m ld
+  in
   { Ast_mapper.default_mapper with
     location
   ; attribute
@@ -206,17 +262,21 @@ let make_mapper conf ~ignore_doc_comments ~erase_jane_syntax =
   ; class_structure
   ; expr
   ; pat
-  ; typ }
+  ; typ
+  ; label_declaration }
 
-let ast fragment ~ignore_doc_comments ~erase_jane_syntax c =
-  map fragment (make_mapper c ~ignore_doc_comments ~erase_jane_syntax)
+let ast fragment ~ignore_doc_comments ~erase_jane_syntax
+    ~local_rewrite_occurred c =
+  map fragment
+    (make_mapper c ~ignore_doc_comments ~erase_jane_syntax
+       ~local_rewrite_occurred )
 
-let equal fragment ~ignore_doc_comments ~erase_jane_syntax c ~old:ast1
-    ~new_:ast2 =
+let equal fragment ~ignore_doc_comments ~erase_jane_syntax
+    ~local_rewrite_occurred c ~old:ast1 ~new_:ast2 =
   let map = ast fragment c ~ignore_doc_comments in
   equal fragment
-    (map ~erase_jane_syntax ast1)
-    (map ~erase_jane_syntax:false ast2)
+    (map ~erase_jane_syntax ~local_rewrite_occurred ast1)
+    (map ~erase_jane_syntax:false ~local_rewrite_occurred ast2)
 
 let ast = ast ~ignore_doc_comments:false
 
@@ -247,20 +307,22 @@ let docstrings (type a) (fragment : a t) s =
   let (_ : a) = map fragment (make_docstring_mapper docstrings) s in
   !docstrings
 
-let docstring conf ~erase_jane_syntax =
+let docstring conf ~erase_jane_syntax ~local_rewrite_occurred =
   let mapper =
     make_mapper conf ~ignore_doc_comments:false ~erase_jane_syntax
+      ~local_rewrite_occurred
   in
   let normalize_code = normalize_code conf mapper in
   docstring conf ~normalize_code
 
-let moved_docstrings fragment ~erase_jane_syntax c ~old:s1 ~new_:s2 =
+let moved_docstrings fragment ~erase_jane_syntax ~local_rewrite_occurred c
+    ~old:s1 ~new_:s2 =
   let d1 = docstrings fragment s1 in
   let d2 = docstrings fragment s2 in
   let equal ~old:(_, x) ~new_:(_, y) =
     String.equal
-      (docstring c x ~erase_jane_syntax)
-      (docstring c y ~erase_jane_syntax:false)
+      (docstring c x ~erase_jane_syntax ~local_rewrite_occurred)
+      (docstring c y ~erase_jane_syntax ~local_rewrite_occurred:false)
   in
   let cmt_kind = `Doc_comment in
   let cmt (loc, x) = Cmt.create_docstring x loc in

lib/Normalize_std_ast.mli

@@ -9,14 +9,21 @@
 (*                                                                        *)
 (**************************************************************************)
 
-val ast : 'a Std_ast.t -> erase_jane_syntax:bool -> Conf.t -> 'a -> 'a
+val ast :
+     'a Std_ast.t
+  -> erase_jane_syntax:bool
+  -> local_rewrite_occurred:bool
+  -> Conf.t
+  -> 'a
+  -> 'a
 (** Normalize an AST fragment. If [erase_jane_syntax] is true, remove all
     [Jane_syntax] attributes signaling erasable syntax. *)
 
 val equal :
      'a Std_ast.t
   -> ignore_doc_comments:bool
   -> erase_jane_syntax:bool
+  -> local_rewrite_occurred:bool
   -> Conf.t
   -> old:'a
   -> new_:'a
@@ -29,6 +36,7 @@ val equal :
 val moved_docstrings :
      'a Std_ast.t
   -> erase_jane_syntax:bool
+  -> local_rewrite_occurred:bool
   -> Conf.t
   -> old:'a
   -> new_:'a