feat: Provide reference lens data for VSCode

Marksman/Lenses.fs

@@ -1,41 +1,58 @@
 module Marksman.Lenses
 
 open Ionide.LanguageServerProtocol.Types
+module LspServer = Ionide.LanguageServerProtocol.Server
 
 open Marksman.Doc
 open Marksman.Folder
+open Marksman.State
 open Marksman.Refs
 
 let findReferencesLens = "marksman.findReferences"
 
+// fsharplint:disable-next-line
+type FindReferencesData = { Uri: DocumentUri; Position: Position; Locations: Location[] }
+
 let private humanRefCount cnt = if cnt = 1 then "1 reference" else $"{cnt} references"
 
-let forDoc (folder: Folder) (doc: Doc) =
-    seq {
-        // Process headers
-        for h in doc.Index.headings do
-            let refCount =
-                Dest.findElementRefs false folder doc (Cst.H h) |> Seq.length
-
-            if refCount > 0 then
-                let command =
-                    { Title = humanRefCount refCount
-                      Command = findReferencesLens
-                      Arguments = None }
-
-                yield { Range = h.range; Command = Some command; Data = None }
-
-        // Process link defs
-        for ld in doc.Index.linkDefs do
-            let refCount =
-                Dest.findElementRefs false folder doc (Cst.MLD ld) |> Seq.length
-
-            if refCount > 0 then
-                let command =
-                    { Title = humanRefCount refCount
-                      Command = findReferencesLens
-                      Arguments = None }
-
-                yield { Range = ld.range; Command = Some command; Data = None }
-    }
-    |> Array.ofSeq
+let buildReferenceLens (client: ClientDescription) (folder: Folder) (doc: Doc) (el: Cst.Element) =
+    let refs = Dest.findElementRefs false folder doc el |> Array.ofSeq
+    let refCount = Array.length refs
+
+    if refCount > 0 then
+        let data =
+            if client.SupportsLensFindReferences then
+                let locations =
+                    [| for doc, el in refs -> { Uri = Doc.uri doc; Range = el.Range } |]
+
+                let data =
+                    { Uri = Doc.uri doc
+                      Position = el.Range.Start
+                      Locations = locations }
+
+                Some [| LspServer.serialize data |]
+            else
+                None
+
+        let command =
+            { Title = humanRefCount refCount
+              Command = findReferencesLens
+              Arguments = data }
+
+
+        Some { Range = el.Range; Command = Some command; Data = None }
+    else
+        None
+
+let forDoc (client: ClientDescription) (folder: Folder) (doc: Doc) =
+    let headingLenses =
+        doc.Index.headings
+        |> Seq.map Cst.H
+        |> Seq.choose (buildReferenceLens client folder doc)
+
+    let linkDefLenses =
+        doc.Index.linkDefs
+        |> Seq.map Cst.MLD
+        |> Seq.choose (buildReferenceLens client folder doc)
+
+    Seq.append headingLenses linkDefLenses |> Array.ofSeq

Marksman/Server.fs

@@ -1022,7 +1022,7 @@ type MarksmanServer(client: MarksmanClient) =
             let lenses =
                 monad' {
                     let! folder, srcDoc = State.tryFindFolderAndDoc docPath state
-                    Lenses.forDoc folder srcDoc
+                    Lenses.forDoc (State.client state) folder srcDoc
                 }
 
             LspResult.success lenses

Marksman/State.fs

@@ -64,6 +64,11 @@ type ClientDescription =
         | None -> false
         | Some exp -> exp.Value<bool>("statusNotification")
 
+    member this.SupportsLensFindReferences: bool =
+        match this.caps.Experimental with
+        | None -> false
+        | Some exp -> exp.Value<bool>("codeLensFindReferences")
+
     member this.SupportsHierarchy: bool =
         monad' {
             let! textDoc = this.caps.TextDocument

Tests/LensesTests.fs

@@ -1,7 +1,12 @@
 module Marksman.LensesTests
 
+open Newtonsoft.Json.Linq
 open Xunit
 
+module LspServer = Ionide.LanguageServerProtocol.Server
+
+open Marksman.State
+
 open Marksman.Helpers
 
 [<Fact>]
@@ -18,7 +23,7 @@ let basicHeaderLenses () =
     let f = FakeFolder.Mk([ d1; d2 ])
 
     let lenses =
-        Lenses.forDoc f d1
+        Lenses.forDoc ClientDescription.empty f d1
         |> Array.map (fun lens -> $"{lens.Command.Value}, {lens.Range}")
 
     checkInlineSnapshot
@@ -31,6 +36,41 @@ let basicHeaderLenses () =
           "  Command = \"marksman.findReferences\""
           "  Arguments = None }, (1,0)-(1,6)" ]
 
+[<Fact>]
+let basicHeaderLenses_withCommandArguments () =
+    let d1 =
+        FakeDoc.Mk(
+            path = "d1.md",
+            contentLines = [| "# Doc 1"; "## Sub"; "[[#Sub]]"; "## No ref" |]
+        )
+
+    let d2 =
+        FakeDoc.Mk(path = "d2.md", contentLines = [| "# Doc 2"; "[[Doc 1#Sub]]" |])
+
+    let f = FakeFolder.Mk([ d1; d2 ])
+
+    let client =
+        { ClientDescription.empty with
+            caps =
+                { ClientDescription.empty.caps with
+                    Experimental = Some(JToken.Parse("""{"codeLensFindReferences": true}""")) } }
+
+    let lensesData =
+        Lenses.forDoc client f d1
+        |> Array.map (fun lens ->
+            let data =
+                lens.Command.Value.Arguments.Value[0]
+                |> LspServer.deserialize<Lenses.FindReferencesData>
+
+            {| Title = lens.Command.Value.Title; Data = data |})
+
+    Assert.Equal(2, lensesData.Length)
+    Assert.Equal("1 reference", lensesData[0].Title)
+    Assert.Equal(1, lensesData[0].Data.Locations.Length)
+    Assert.Equal("2 references", lensesData[1].Title)
+    Assert.Equal(2, lensesData[1].Data.Locations.Length)
+
+
 [<Fact>]
 let basicLinkDefLenses () =
     let d1 =
@@ -42,7 +82,7 @@ let basicLinkDefLenses () =
     let f = FakeFolder.Mk([ d1 ])
 
     let lenses =
-        Lenses.forDoc f d1
+        Lenses.forDoc ClientDescription.empty f d1
         |> Array.map (fun lens -> $"{lens.Command.Value}, {lens.Range}")
 
     checkInlineSnapshot