Stricter AST checks (#2557)

* Normalize docstrings using the extended AST

This use the extended AST normalization for docstrings within the std
AST. This ensures that repl phrases are handled correctly and removes
some code.

* Stricter AST checks

The AST check was previously trigerring if both the std AST and the
extended AST were different.

The std AST is intended to be the most trustworthy and is intended to
catch changes in the code that the extended AST couldn't.

Author: Julow

lib/Normalize_extended_ast.ml

@@ -174,6 +174,10 @@ let normalize_cmt (conf : Conf.t) =
       normalize_code ~normalize_cmt conf mapper c
   end
 
+let normalize_code (conf : Conf.t) code =
+  let n = normalize_cmt conf in
+  n#code code
+
 let ast fragment ~ignore_doc_comments c =
   let normalize_cmt = normalize_cmt c in
   map fragment

lib/Normalize_extended_ast.mli

@@ -20,3 +20,6 @@ val equal : 'a t -> ignore_doc_comments:bool -> Conf.t -> 'a -> 'a -> bool
 val diff_cmts :
   Conf.t -> Cmt.t list -> Cmt.t list -> (unit, Cmt.error list) Result.t
 (** Difference between two lists of comments. *)
+
+val normalize_code : Conf.t -> string -> string
+(** Normalize a code block in docstrings. *)

lib/Normalize_std_ast.ml

@@ -18,49 +18,10 @@ let is_doc = function
   | {attr_name= {Location.txt= "ocaml.doc" | "ocaml.text"; _}; _} -> true
   | _ -> false
 
-let dedup_cmts fragment ast comments =
-  let of_ast ast =
-    let docs = ref (Set.empty (module Cmt)) in
-    let attribute m atr =
-      match atr with
-      | { attr_payload=
-            PStr
-              [ { pstr_desc=
-                    Pstr_eval
-                      ( { pexp_desc=
-                            Pexp_constant (Pconst_string (doc, _, None))
-                        ; pexp_loc
-                        ; _ }
-                      , [] )
-                ; _ } ]
-        ; _ }
-        when is_doc atr ->
-          docs := Set.add !docs (Cmt.create_docstring doc pexp_loc) ;
-          atr
-      | _ -> Ast_mapper.default_mapper.attribute m atr
-    in
-    map fragment {Ast_mapper.default_mapper with attribute} ast |> ignore ;
-    !docs
-  in
-  Set.(to_list (diff (of_list (module Cmt) comments) (of_ast ast)))
-
-let normalize_code conf (m : Ast_mapper.mapper) txt =
-  let input_name = "<output>" in
-  match
-    Parse_with_comments.parse Parse.ast Structure conf ~input_name
-      ~source:txt
-  with
-  | {ast; comments; _} ->
-      let comments = dedup_cmts Structure ast comments in
-      let print_comments fmt (l : Cmt.t list) =
-        List.sort l ~compare:(fun a b ->
-            Migrate_ast.Location.compare (Cmt.loc a) (Cmt.loc b) )
-        |> List.iter ~f:(fun cmt -> Format.fprintf fmt "%s," (Cmt.txt cmt))
-      in
-      let ast = m.structure m ast in
-      Format.asprintf "AST,%a,COMMENTS,[%a]" Printast.implementation ast
-        print_comments comments
-  | exception _ -> txt
+let normalize_code conf txt =
+  (* Normalize code blocks in docstrings using the extended AST. This
+     correctly handles repl phrases. *)
+  Normalize_extended_ast.normalize_code conf txt
 
 let docstring (c : Conf.t) =
   Docstring.normalize ~parse_docstrings:c.fmt_opts.parse_docstrings.v
@@ -83,7 +44,7 @@ let make_mapper conf ~ignore_doc_comments =
                   , [] )
             ; _ } as pstr ) ]
       when is_doc attr ->
-        let normalize_code = normalize_code conf m in
+        let normalize_code = normalize_code conf in
         let doc' = docstring conf ~normalize_code doc in
         Ast_mapper.default_mapper.attribute m
           { attr with
@@ -239,8 +200,7 @@ let docstrings (type a) (fragment : a t) s =
   !docstrings
 
 let docstring conf =
-  let mapper = make_mapper conf ~ignore_doc_comments:false in
-  let normalize_code = normalize_code conf mapper in
+  let normalize_code = normalize_code conf in
   docstring conf ~normalize_code
 
 let moved_docstrings fragment c s1 s2 =

lib/Translation_unit.ml

@@ -322,13 +322,9 @@ let format (type ext std) (ext_fg : ext Extended_ast.t)
       in
       (* Ast not preserved ? *)
       ( if
-          (not
-             (Normalize_std_ast.equal std_fg conf std_t.ast std_t_new.ast
-                ~ignore_doc_comments:(not conf.opr_opts.comment_check.v) ) )
-          && not
-               (Normalize_extended_ast.equal ext_fg conf ext_t.ast
-                  ext_t_new.ast
-                  ~ignore_doc_comments:(not conf.opr_opts.comment_check.v) )
+          not
+            (Normalize_std_ast.equal std_fg conf std_t.ast std_t_new.ast
+               ~ignore_doc_comments:(not conf.opr_opts.comment_check.v) )
         then
           let old_ast =
             dump_ast std_fg ~suffix:".old"