Fix occurrences with let punning

Author: liam923

src/analysis/index_occurrences.ml

@@ -25,11 +25,12 @@ let decl_of_path_or_lid env namespace path lid =
   end
   | _ -> Env_lookup.by_path path namespace env
 
+let should_ignore_lid (lid : Longident.t Location.loc) = Location.is_none lid.loc
+
 let iterator ~current_buffer_path ~index ~stamp ~reduce_for_uid =
   let add uid loc = Stamped_hashtable.add index ~stamp (uid, loc) () in
   let f ~namespace env path (lid : Longident.t Location.loc) =
     log ~title:"index_buffer" "Path: %a" Logger.fmt (Fun.flip Path.print path);
-    let not_ghost { Location.loc = { loc_ghost; _ }; _ } = not loc_ghost in
     let lid = { lid with loc = set_fname ~file:current_buffer_path lid.loc } in
     let index_decl () =
       begin
@@ -42,7 +43,7 @@ let iterator ~current_buffer_path ~index ~stamp ~reduce_for_uid =
           add decl.uid lid
       end
     in
-    if not_ghost lid then
+    if not (should_ignore_lid lid) then
       match Env.shape_of_path ~namespace env path with
       | exception Not_found -> ()
       | path_shape ->

src/analysis/occurrences.ml

@@ -200,18 +200,25 @@ let locs_of ~config ~env ~typer_result ~pos ~scope path =
             Option.map external_locs ~f:(fun (index, locs) ->
                 let stats = Stat_check.create ~cache_size:128 index in
                 ( Lid_set.filter
-                    (fun { loc; _ } ->
-                      (* We ignore external results that concern the current buffer *)
-                      let file = loc.Location.loc_start.Lexing.pos_fname in
-                      let file, buf =
-                        match config.merlin.source_root with
-                        | Some root ->
-                          (Filename.concat root file, current_buffer_path)
-                        | None -> (file, config.query.filename)
+                    (fun ({ loc; _ } as lid) ->
+                      let is_current_buffer =
+                        (* We filter external results that concern the current buffer *)
+                        let file = loc.Location.loc_start.Lexing.pos_fname in
+                        let file, buf =
+                          match config.merlin.source_root with
+                          | Some root ->
+                            (Filename.concat root file, current_buffer_path)
+                          | None -> (file, config.query.filename)
+                        in
+                        let file = Misc.canonicalize_filename file in
+                        let buf = Misc.canonicalize_filename buf in
+                        String.equal file buf
                       in
-                      let file = Misc.canonicalize_filename file in
-                      let buf = Misc.canonicalize_filename buf in
-                      if String.equal file buf then false
+                      let should_be_ignored =
+                        (* We ignore results that don't have a location *)
+                        Index_occurrences.should_ignore_lid lid
+                      in
+                      if is_current_buffer || should_be_ignored then false
                       else begin
                         (* We ignore external results if their source was modified *)
                         let check = Stat_check.check stats ~file in

tests/test-dirs/let-punning.t/run.t

@@ -248,8 +248,6 @@ Test that finding occurrences of a variable includes usages in a punned let. i.e
 finding occurrences of x on line 1 returns the definition on line 1 and the usage on
 line 2.
 
-TODO: fix these tests
-
 let*
   $ occurrences 12:8
   Occurrences of:
@@ -258,6 +256,9 @@ let*
   Occurrence at 12:8-9:
       let a = return 1 in
           ^
+  Occurrence at 13:9-10:
+      let* a in
+           ^
 
 parallel let*
   $ occurrences 18:8
@@ -267,13 +268,19 @@ parallel let*
   Occurrence at 18:8-9:
       let a = return 1 in
           ^
+  Occurrence at 20:9-10:
+      let* a and* b in
+           ^
   $ occurrences 19:8
   Occurrences of:
       let b = return 1 in
           ^
   Occurrence at 19:8-9:
       let b = return 1 in
           ^
+  Occurrence at 20:16-17:
+      let* a and* b in
+                  ^
 
 sequential let*
   $ occurrences 25:8
@@ -283,10 +290,16 @@ sequential let*
   Occurrence at 25:8-9:
       let a = return 1 in
           ^
+  Occurrence at 27:9-10:
+      let* a in
+           ^
   $ occurrences 26:8
   Occurrences of:
       let b = return 1 in
           ^
   Occurrence at 26:8-9:
       let b = return 1 in
           ^
+  Occurrence at 28:9-10:
+      let* b in
+           ^

tests/test-dirs/occurrences/punning.t/run.t

@@ -9,9 +9,9 @@ Convenience function for querying occurrences
 
 Get occurrences of an identifier that is used as the expression part of a punned let
 expression
-FIXME: this should also include the occurrence on line 5
   $ occurrences 4:6
   4:6-7
+  5:7-8
 
 Get occurrences, with the cursor pointing at the identifier in a punned let.
 Merlin returns the occurrences of the new variable bound in that let, rather than the