Correctly traverse patterns when itering for docs #1572

voodoos · web-flow · commit e27873475393 · 2023-02-24T07:22:16.000-08:00
from voodoos/iter-on-whole-pattern-to-find-uids
diff --git a/CHANGES.md b/CHANGES.md
@@ -12,6 +12,7 @@ unreleased
       jump-to-definition'q behavior (#1563)
     - Improve locate's behavior in various ill-typed expressions (#1546, fixes
       #1567 and partially #1543)
+    - Correctly traverse patterns when looking for docs in the typedtree (#1572)
   + test suite
     - Add multiple tests for locate over ill-typed expressions (#1546)
 
diff --git a/src/analysis/locate.ml b/src/analysis/locate.ml
@@ -796,10 +796,13 @@ let from_string ~config ~env ~local_defs ~pos ?namespaces switch path =
     a uid-based search and return the attached comment in the attributes.
     This is a more sound way to get documentation than resorting on the
     [Ocamldoc.associate_comment] heuristic *)
-let doc_from_uid ~config ~comp_unit uid =
-  let exception Found of Typedtree.attributes in
+(* In a future release of OCaml the cmt's uid_to_loc table will contain
+   fragments of the typedtree that might be used to get the docstrings without
+   relying on this iteration *)
+let find_doc_attributes_in_typedtree ~config ~comp_unit uid =
+  let exception Found_attributes of Typedtree.attributes in
   let test elt_uid attributes =
-    if Shape.Uid.equal uid elt_uid then raise (Found attributes)
+    if Shape.Uid.equal uid elt_uid then raise (Found_attributes attributes)
   in
   let iterator =
     let first_item = ref true in
@@ -813,14 +816,14 @@ let doc_from_uid ~config ~comp_unit uid =
          The module docstring must be the first signature or structure item *)
       signature_item = (fun sub ({ sig_desc; _} as si) ->
         begin match sig_desc, !first_item, uid_is_comp_unit with
-        | Tsig_attribute attr, true, true -> raise (Found [attr])
+        | Tsig_attribute attr, true, true -> raise (Found_attributes [attr])
         | _, false, true -> raise Not_found
         | _, _, _ -> first_item := false end;
         Tast_iterator.default_iterator.signature_item sub si);
 
       structure_item = (fun sub ({ str_desc; _} as sti) ->
         begin match str_desc, !first_item, uid_is_comp_unit with
-        | Tstr_attribute attr, true, true -> raise (Found [attr])
+        | Tstr_attribute attr, true, true -> raise (Found_attributes [attr])
         | _, false, true -> raise Not_found
         | _, _, _ -> first_item := false end;
         Tast_iterator.default_iterator.structure_item sub sti);
@@ -834,25 +837,34 @@ let doc_from_uid ~config ~comp_unit uid =
         Tast_iterator.default_iterator.type_declaration sub td);
 
       value_binding = (fun sub ({ vb_pat; vb_attributes; _ } as vb) ->
-        begin match vb_pat.pat_desc with
-        | Tpat_var (id, _) ->
-            begin try
-              let vd = Env.find_value (Pident id) env in
-              test vd.val_uid vb_attributes
-            with Not_found -> () end
-        | _ -> () end;
+        let pat_var_iter ~f pat =
+          let rec aux pat =
+            let open Typedtree in
+            match pat.pat_desc with
+            | Tpat_var (id, _) -> f id
+            | Tpat_alias (pat, _, _)
+            | Tpat_variant (_, Some pat, _)
+            | Tpat_lazy pat
+            | Tpat_or (pat, _, _) ->
+                aux pat
+            | Tpat_tuple pats
+            | Tpat_construct (_, _, pats, _)
+            | Tpat_array pats ->
+                List.iter ~f:aux pats
+            | Tpat_record (pats, _) ->
+                List.iter ~f:(fun (_, _, pat) -> aux pat) pats
+            | _ -> ()
+          in
+          aux pat
+        in
+        pat_var_iter vb_pat ~f:(fun id ->
+          try
+            let vd = Env.find_value (Pident id) env in
+            test vd.val_uid vb_attributes
+          with Not_found -> ());
         Tast_iterator.default_iterator.value_binding sub vb)
     }
   in
-  let parse_attributes attrs =
-    let open Parsetree in
-    try Some (List.find_map attrs ~f:(fun attr ->
-      if List.exists ["ocaml.doc"; "ocaml.text"]
-        ~f:(String.equal attr.attr_name.txt)
-      then Ast_helper.extract_str_payload attr.attr_payload
-      else None))
-    with Not_found -> None
-  in
   let typedtree =
     log ~title:"doc_from_uid" "Loading the cmt for unit %S" comp_unit;
     match load_cmt ~config comp_unit `MLI with
@@ -879,38 +891,86 @@ let doc_from_uid ~config ~comp_unit uid =
     | _ -> () end;
     `No_documentation
   with
-    | Found attrs ->
-       log ~title:"doc_from_uid" "Found attributes for this uid";
+    | Found_attributes attrs ->
+        log ~title:"doc_from_uid" "Found attributes for this uid";
+        let parse_attributes attrs =
+          let open Parsetree in
+          try Some (List.find_map attrs ~f:(fun attr ->
+            if List.exists ["ocaml.doc"; "ocaml.text"]
+              ~f:(String.equal attr.attr_name.txt)
+            then Ast_helper.extract_str_payload attr.attr_payload
+            else None))
+          with Not_found -> None
+        in
         begin match parse_attributes attrs with
         | Some (doc, _) -> `Found (doc |> String.trim)
         | None -> `No_documentation end
     | Not_found -> `No_documentation
 
+let doc_from_uid ~config ~loc uid =
+  begin match uid with
+  | Some (Shape.Uid.Item { comp_unit; _ } as uid)
+  | Some (Shape.Uid.Compilation_unit comp_unit as uid)
+      when Env.get_unit_name () <> comp_unit ->
+        log ~title:"get_doc" "the doc (%a) you're looking for is in another
+          compilation unit (%s)"
+          Logger.fmt (fun fmt -> Shape.Uid.print fmt uid) comp_unit;
+        (match find_doc_attributes_in_typedtree ~config ~comp_unit uid with
+        | `Found doc -> `Found_doc doc
+        | `No_documentation ->
+            (* We fallback on the legacy heuristic to handle some unproper
+               doc placement. See test [unattached-comment.t] *)
+            `Found_loc loc)
+  | _ ->
+    (* Uid based search doesn't works in the current CU since Merlin's parser
+       does not attach doc comments to the typedtree *)
+    `Found_loc loc
+  end
+
+let doc_from_comment_list ~local_defs ~buffer_comments loc =
+  (* When the doc we look for is in the current buffer or if search by uid
+    has failed we use an alternative heuristic since Merlin's pure parser
+    does not poulates doc attributes in the typedtree. *)
+  let comments =
+    match File_switching.where_am_i () with
+    | None ->
+      log ~title:"get_doc" "Using reader's comment (current buffer)";
+      buffer_comments
+    | Some cmt_path ->
+      log ~title:"get_doc" "File switching: actually in %s" cmt_path;
+      let {Cmt_cache. cmt_infos; _ } = Cmt_cache.read cmt_path in
+      cmt_infos.Cmt_format.cmt_comments
+  in
+  log ~title:"get_doc" "%a" Logger.fmt (fun fmt ->
+      Format.fprintf fmt "looking around %a inside: [\n"
+        Location.print_loc !last_location;
+      List.iter comments ~f:(fun (c, l) ->
+          Format.fprintf fmt "  (%S, %a);\n" c
+            Location.print_loc l);
+      Format.fprintf fmt "]\n"
+    );
+  let browse = Mbrowse.of_typedtree local_defs in
+  let (_, deepest_before) =
+    Mbrowse.(leaf_node @@ deepest_before loc.Location.loc_start [browse])
+  in
+  (* based on https://v2.ocaml.org/manual/doccomments.html#ss:label-comments: *)
+  let after_only = begin match deepest_before with
+    | Browse_raw.Constructor_declaration _ -> true
+    (* The remaining `true` cases are currently not reachable *)
+    | Label_declaration _ | Record_field _ | Row_field _  -> true
+    | _ -> false
+  end in
+  match
+    Ocamldoc.associate_comment ~after_only comments loc !last_location
+  with
+  | None, _     -> `No_documentation
+  | Some doc, _ -> `Found doc
+
 let get_doc ~config ~env ~local_defs ~comments ~pos =
   File_switching.reset ();
-  let from_uid ~loc uid =
-    begin match uid with
-    | Some (Shape.Uid.Item { comp_unit; _ } as uid)
-    | Some (Shape.Uid.Compilation_unit comp_unit as uid)
-        when Env.get_unit_name () <> comp_unit ->
-          log ~title:"get_doc" "the doc (%a) you're looking for is in another
-            compilation unit (%s)"
-            Logger.fmt (fun fmt -> Shape.Uid.print fmt uid) comp_unit;
-          (match doc_from_uid ~config ~comp_unit uid with
-          | `Found doc -> `Found_doc doc
-          | `No_documentation ->
-              (* We fallback on the legacy heuristic to handle some unproper
-                 doc placement. See test [unattached-comment.t] *)
-              `Found loc)
-    | _ ->
-      (* Uid based search doesn't works in the current CU since Merlin's parser
-         does not attach doc comments to the typedtree *)
-      `Found loc
-    end
-  in
   fun path ->
   let_ref last_location Location.none @@ fun () ->
-  match
+  let doc_from_uid_result =
     match path with
     | `Completion_entry (namespace, path, _loc) ->
       log ~title:"get_doc" "completion: looking for the doc of '%a'"
@@ -921,7 +981,7 @@ let get_doc ~config ~env ~local_defs ~comments ~pos =
         let loc : Location.t =
           { loc_start = pos; loc_end = pos; loc_ghost = true }
         in
-        from_uid ~loc uid
+        doc_from_uid ~config ~loc uid
       | (`Builtin |`Not_in_env _|`File_not_found _|`Not_found _)
         as otherwise -> otherwise
       end
@@ -932,53 +992,17 @@ let get_doc ~config ~env ~local_defs ~comments ~pos =
         let loc : Location.t =
           { loc_start = pos; loc_end = pos; loc_ghost = true }
         in
-        from_uid ~loc uid
+        doc_from_uid ~config ~loc uid
       | `At_origin | `Missing_labels_namespace -> `No_documentation
       | `Builtin _ -> `Builtin
       | (`Not_in_env _ | `Not_found _ |`File_not_found _ )
         as otherwise -> otherwise
       end
-  with
+  in
+  match doc_from_uid_result with
   | `Found_doc doc -> `Found doc
-  | `Found loc ->
-    (* When the doc we look for is in the current buffer or if search by uid
-      has failed we use an alternative heuristic since Merlin's pure parser
-      does not poulates doc attributes in the typedtree. *)
-    let comments =
-      match File_switching.where_am_i () with
-      | None ->
-        log ~title:"get_doc" "Using reader's comment (current buffer)";
-        comments
-      | Some cmt_path ->
-        log ~title:"get_doc" "File switching: actually in %s" cmt_path;
-        let {Cmt_cache. cmt_infos; _ } = Cmt_cache.read cmt_path in
-        cmt_infos.Cmt_format.cmt_comments
-    in
-    log ~title:"get_doc" "%a" Logger.fmt (fun fmt ->
-        Format.fprintf fmt "looking around %a inside: [\n"
-          Location.print_loc !last_location;
-        List.iter comments ~f:(fun (c, l) ->
-            Format.fprintf fmt "  (%S, %a);\n" c
-              Location.print_loc l);
-        Format.fprintf fmt "]\n"
-      );
-    let browse = Mbrowse.of_typedtree local_defs in
-    let (_, deepest_before) =
-      Mbrowse.(leaf_node @@ deepest_before loc.loc_start [browse])
-    in
-    (* based on https://v2.ocaml.org/manual/doccomments.html#ss:label-comments: *)
-    let after_only = begin match deepest_before with
-      | Browse_raw.Constructor_declaration _ -> true
-      (* The remaining `true` cases are currently not reachable *)
-      | Label_declaration _ | Record_field _ | Row_field _  -> true
-      | _ -> false
-    end in
-    begin match
-      Ocamldoc.associate_comment ~after_only comments loc !last_location
-    with
-    | None, _     -> `No_documentation
-    | Some doc, _ -> `Found doc
-    end
+  | `Found_loc loc ->
+      doc_from_comment_list ~local_defs ~buffer_comments:comments loc
   | `Builtin ->
     begin match path with
     | `User_input path -> `Builtin path
diff --git a/tests/test-dirs/document/external-pattern-comments.t b/tests/test-dirs/document/external-pattern-comments.t
@@ -0,0 +1,23 @@
+  $ cat >main.ml <<EOF
+  > let _ = Lib.y
+  > EOF
+
+  $ cat >lib.ml <<EOF
+  > (** doc for all node *)
+  > let x, y = 2, 3
+  > EOF
+
+  $ $OCAMLC -c -bin-annot lib.ml main.ml
+
+  $ $MERLIN single document -position 1:12 \
+  > -log-file - 2>log \
+  > -filename main.ml <main.ml 
+  {
+    "class": "return",
+    "value": "doc for all node",
+    "notifications": []
+  }
+
+We should not rely on the heuristic to get that comment
+  $ cat log | grep -A 2 "looking around"
+  [1]
