occ: improve handling of paths in current buffer

src/analysis/occurrences.ml

@@ -25,7 +25,7 @@ let decl_of_path_or_lid env namespace path lid =
     end
   | _ -> Env_lookup.by_path path namespace env
 
-let index_buffer_ ~current_buffer_path ~local_defs () =
+let index_buffer_ ~config ~source ~current_buffer_path ~local_defs () =
   let {Logger. log} = Logger.for_section "index" in
   let defs = Hashtbl.create 64 in
   let module Shape_reduce =
@@ -66,38 +66,51 @@ let index_buffer_ ~current_buffer_path ~local_defs () =
         let result =  Shape_reduce.reduce_for_uid env path_shape in
         begin match Locate.uid_of_result ~traverse_aliases:false result with
         | Some uid, false ->
-          log ~title:"index_buffer" "Found %s (%a) wiht uid %a"
-            (Longident.head lid.txt)
+          log ~title:"index_buffer" "Found %a (%a) wiht uid %a"
+            Logger.fmt (Fun.flip Pprintast.longident lid.txt)
             Logger.fmt (Fun.flip Location.print_loc lid.loc)
             Logger.fmt (Fun.flip Shape.Uid.print uid);
-            Index_format.(add defs uid (LidSet.singleton lid))
-          | Some uid, true ->
-            log ~title:"index_buffer" "Shape is approximative, found uid: %a"
-              Logger.fmt (Fun.flip Shape.Uid.print uid);
-            index_decl ()
-          | None, _ ->
-            log ~title:"index_buffer" "Reduction failed: missing uid";
-            index_decl ()
+          Index_format.(add defs uid (LidSet.singleton lid))
+        | Some uid, true ->
+          log ~title:"index_buffer" "Shape is approximative, found uid: %a"
+            Logger.fmt (Fun.flip Shape.Uid.print uid);
+          index_decl ()
+        | None, _ ->
+          log ~title:"index_buffer" "Reduction failed: missing uid";
+          index_decl ()
         end
   in
   let f ~namespace env path (lid : Longident.t Location.loc)  =
     (* The compiler lacks sufficient location information to precisely hihglight
        modules in paths. This function hacks around that issue when looking for
        occurrences in the current buffer only. *)
-    let rec iter_on_path ~namespace path ({Location.txt; loc} as lid) =
-      let () = f ~namespace env path lid in
-      match path, txt with
-      | Pdot (path, _), Ldot (lid, s) ->
-        let length_with_dot = String.length s + 1 in
-        let lid =
-          { Location.txt = lid; loc = { loc with loc_end = {loc.loc_end with
-            pos_cnum = loc.loc_end.pos_cnum - length_with_dot}} }
-        in
-        iter_on_path ~namespace:Module path lid
-      | Papply _, _ -> ()
-      | _, _ -> ()
+    (* We rely on a custom re-parsing of the longidents that provide us with
+       location information and match these with the real path and longident. *)
+    let rec iter_on_path ~namespace (path : Path.t) lid reparsed =
+      log ~title:"iter_on_path" "Path %a, lid: %a"
+        Logger.fmt (Fun.flip Path.print path)
+        Logger.fmt (Fun.flip Pprintast.longident lid.Location.txt);
+      match path, lid.txt, reparsed with
+      | Pdot (path', n), (Ldot (_, s) | Lident s), _
+          when n <> s && String.lowercase_ascii n = n ->
+          iter_on_path ~namespace path' lid reparsed
+      | (Pdot _ | Pident _), Lident s, [{ Location.txt = name; loc}]
+        when name = s ->
+          log ~title:"iter_on_path" "Last %a, lid: %a"
+            Logger.fmt (Fun.flip Path.print path)
+            Logger.fmt (Fun.flip Pprintast.longident lid.Location.txt);
+          f ~namespace env path { lid with loc = loc }
+      | Pdot (path', _), Ldot (lid', s), { txt = name; loc} :: tl when name = s ->
+          let () = f ~namespace env path { lid with loc = loc } in
+          iter_on_path ~namespace:Module path' { lid with txt = lid' } tl
+      | Papply _, _, _ -> f ~namespace env path lid
+      | _, _, _ -> f ~namespace env path lid
+    in
+    let reparsed_lid =
+      Misc_utils.parse_identifier (config, source) lid.loc.loc_end
+      |> List.rev
     in
-    iter_on_path ~namespace path lid
+    iter_on_path ~namespace path lid reparsed_lid
   in
   Ast_iterators.iter_on_usages ~f local_defs;
   defs
@@ -106,7 +119,7 @@ let index_buffer =
   (* Right now, we only cache the last used index. We could do better by caching
      the index for every known buffer. *)
   let cache = ref None in
-  fun ~current_buffer_path ~stamp ~local_defs () ->
+  fun ~config ~source ~current_buffer_path ~stamp ~local_defs () ->
     let {Logger. log} = Logger.for_section "index" in
     match !cache with
     | Some (path, stamp', value) when
@@ -117,7 +130,9 @@ let index_buffer =
       value
     | _ ->
       log ~title:"index_cache" "No valid cache found, reindexing.";
-      let result = index_buffer_ ~current_buffer_path ~local_defs () in
+      let result =
+        index_buffer_ ~config ~source ~current_buffer_path ~local_defs ()
+      in
       cache := Some (current_buffer_path, stamp, result);
       result
 
@@ -174,7 +189,7 @@ let comp_unit_of_uid = function
   | Item { comp_unit; _ } -> Some comp_unit
   | Internal | Predef _ -> None
 
-let locs_of ~config ~env ~typer_result ~pos path =
+let locs_of ~config ~source ~env ~typer_result ~pos path =
   log ~title:"occurrences" "Looking for occurences of %s (pos: %s)"
     path
     (Lexing.print_position () pos);
@@ -218,7 +233,7 @@ let locs_of ~config ~env ~typer_result ~pos path =
     log ~title:"locs_of" "Indexing current buffer";
     let buffer_index =
       let stamp = Mtyper.get_stamp typer_result in
-      index_buffer ~current_buffer_path ~stamp ~local_defs ()
+      index_buffer ~config ~source ~current_buffer_path ~stamp ~local_defs ()
     in
     let buffer_locs = Hashtbl.find_opt buffer_index def_uid in
     let locs = Option.value ~default:LidSet.empty buffer_locs in

src/analysis/occurrences.mli

@@ -1,5 +1,6 @@
 val locs_of
   : config:Mconfig.t
+  -> source:Msource.t
   -> env:Env.t
   -> typer_result:Mtyper.result
   -> pos:Lexing.position

src/frontend/query_commands.ml

@@ -797,6 +797,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a =
 
   | Occurrences (`Ident_at pos, _) ->
     let config = Mpipeline.final_config pipeline in
+    let source = Mpipeline.raw_source pipeline in
     let typer_result = Mpipeline.typer_result pipeline in
     let pos = Mpipeline.get_lexing_pos pipeline pos in
     let env, _node = Mbrowse.leaf_node (Mtyper.node_at typer_result pos) in
@@ -809,7 +810,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a =
       path
     in
     let locs =
-      Occurrences.locs_of ~config ~env ~typer_result ~pos path
+      Occurrences.locs_of ~config ~source ~env ~typer_result ~pos path
       |> Result.value ~default:[]
     in
     let loc_start l = l.Location.loc_start in

tests/test-dirs/occurrences/mod-in-path-2.t

@@ -0,0 +1,78 @@
+  $ cat >test.ml <<'EOF'
+  > module Mod = struct
+  >   type t = A 
+  > end
+  > let () = 
+  >   match Mod.A with
+  >   | Mod.A -> ()
+  > EOF
+
+  $ $MERLIN single occurrences -identifier-at 1:8 -filename test.ml <test.ml | 
+  > jq '.value'
+  [
+    {
+      "start": {
+        "line": 1,
+        "col": 7
+      },
+      "end": {
+        "line": 1,
+        "col": 10
+      }
+    },
+    {
+      "start": {
+        "line": 5,
+        "col": 8
+      },
+      "end": {
+        "line": 5,
+        "col": 11
+      }
+    },
+    {
+      "start": {
+        "line": 6,
+        "col": 4
+      },
+      "end": {
+        "line": 6,
+        "col": 7
+      }
+    }
+  ]
+
+  $ $MERLIN single occurrences -identifier-at 2:11 -filename test.ml <test.ml | 
+  > jq '.value'
+  [
+    {
+      "start": {
+        "line": 2,
+        "col": 11
+      },
+      "end": {
+        "line": 2,
+        "col": 12
+      }
+    },
+    {
+      "start": {
+        "line": 5,
+        "col": 12
+      },
+      "end": {
+        "line": 5,
+        "col": 13
+      }
+    },
+    {
+      "start": {
+        "line": 6,
+        "col": 8
+      },
+      "end": {
+        "line": 6,
+        "col": 9
+      }
+    }
+  ]

tests/test-dirs/occurrences/mod-in-path-3.t

@@ -0,0 +1,114 @@
+  $ cat >test.ml <<'EOF'
+  > module Mod = struct
+  >   type t = A of { lbl : int } 
+  > end
+  > let x = Mod.A { lbl = 42 }
+  > let _ =
+  >   match x with
+  >   | Mod.A r -> r.lbl
+  > EOF
+
+  $ $MERLIN single occurrences -identifier-at 4:9 -filename test.ml <test.ml | 
+  > jq '.value'
+  [
+    {
+      "start": {
+        "line": 1,
+        "col": 7
+      },
+      "end": {
+        "line": 1,
+        "col": 10
+      }
+    },
+    {
+      "start": {
+        "line": 4,
+        "col": 8
+      },
+      "end": {
+        "line": 4,
+        "col": 11
+      }
+    },
+    {
+      "start": {
+        "line": 7,
+        "col": 4
+      },
+      "end": {
+        "line": 7,
+        "col": 7
+      }
+    }
+  ]
+
+  $ $MERLIN single occurrences -identifier-at 4:12 -filename test.ml <test.ml | 
+  > jq '.value'
+  [
+    {
+      "start": {
+        "line": 2,
+        "col": 11
+      },
+      "end": {
+        "line": 2,
+        "col": 12
+      }
+    },
+    {
+      "start": {
+        "line": 4,
+        "col": 12
+      },
+      "end": {
+        "line": 4,
+        "col": 13
+      }
+    },
+    {
+      "start": {
+        "line": 7,
+        "col": 8
+      },
+      "end": {
+        "line": 7,
+        "col": 9
+      }
+    }
+  ]
+
+  $ $MERLIN single occurrences -identifier-at 4:18 -filename test.ml <test.ml | 
+  > jq '.value'
+  [
+    {
+      "start": {
+        "line": 2,
+        "col": 18
+      },
+      "end": {
+        "line": 2,
+        "col": 21
+      }
+    },
+    {
+      "start": {
+        "line": 4,
+        "col": 16
+      },
+      "end": {
+        "line": 4,
+        "col": 19
+      }
+    },
+    {
+      "start": {
+        "line": 7,
+        "col": 17
+      },
+      "end": {
+        "line": 7,
+        "col": 20
+      }
+    }
+  ]