fix: remove duplicate from location list

OCamlPro · Aug 8, 2024 · 893ec55 · 893ec55
1 parent 48c2947
commit 893ec55
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git a/src/lsp/cobol_lsp/lsp_request.ml b/src/lsp/cobol_lsp/lsp_request.ml
@@ -625,6 +625,23 @@ let codelens_positions ~uri group =
     end group (None, PosSet.empty)
   |> snd
 
+
+module Locations = Set.Make(struct
+    type t = Location.t
+    let compare (l1: t) (l2: t) =
+      let c = String.compare
+          (DocumentUri.to_path l1.uri) (DocumentUri.to_path l2.uri) in
+      if c <> 0 then c else
+        let r1, r2 = l1.range, l2.range in
+        let c = Int.compare r1.start.line r2.start.line in
+        if c <> 0 then c else
+          let c = Int.compare r1.start.character r2.start.character in
+          if c <> 0 then c else
+            let c = Int.compare r1.end_.line r2.end_.line in
+            if c <> 0 then c else
+              Int.compare r1.end_.character r2.end_.character
+  end)
+
 let handle_codelens registry ({ textDocument; _ }: CodeLensParams.t) =
   try_with_main_document_data registry textDocument
     ~f:begin fun ~doc checked_doc ->
@@ -638,7 +655,9 @@ let handle_codelens registry ({ textDocument; _ }: CodeLensParams.t) =
           ReferenceParams.create ~context ~position ~textDocument () in
         let ref_count =
           lookup_references_in_doc ~rootdir params checked_doc
-          |> Option.fold ~none:0 ~some:List.length in
+          |> Option.value ~default:[]
+          |> Locations.of_list
+          |> Locations.cardinal in
         let title = string_of_int ref_count
                     ^ " reference"
                     ^ if ref_count > 1 then "s" else "" in

diff --git a/test/lsp/lsp_codelens.ml b/test/lsp/lsp_codelens.ml
@@ -196,7 +196,7 @@ let%expect_test "codelens-procedure" =
     ----               ^
        6               BB.
        7               PERFORM BB.
-    4 references
+    2 references
     __rootdir__/prog.cob:6.12:
        3           PROGRAM-ID. prog.
        4           PROCEDURE DIVISION.
@@ -205,7 +205,7 @@ let%expect_test "codelens-procedure" =
     ----               ^
        7               PERFORM BB.
        8               CC.
-    4 references
+    3 references
     __rootdir__/prog.cob:8.12:
        5               AA SECTION.
        6               BB.
@@ -214,7 +214,7 @@ let%expect_test "codelens-procedure" =
     ----               ^
        9               DD SECTION.
       10               PERFORM AA.
-    2 references
+    1 reference
     __rootdir__/prog.cob:9.12:
        6               BB.
        7               PERFORM BB.
@@ -223,4 +223,4 @@ let%expect_test "codelens-procedure" =
     ----               ^
       10               PERFORM AA.
       11               PERFORM BB.
-    4 references |}];;
+    2 references |}];;