feat: Provide reference lens data for VSCode

When client's experimental.codeLensFindReferences capability is set (only VSCode plugin now), Marksman will output extra data for the code lens command that can be used to implement client-side behavior to show references when clicking on the lens.

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