Respect Context.IncludeDeclaration in handler for "textDocument/findR…

…eferences"
razzmatazz · Nov 15, 2024 · 364c876 · 364c876
1 parent bd7f06c
commit 364c876
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 12 deletions.
diff --git a/src/CSharpLanguageServer/Handlers/CodeLens.fs b/src/CSharpLanguageServer/Handlers/CodeLens.fs
@@ -162,13 +162,12 @@ module CodeLens =
         match! context.FindSymbol lensData.DocumentUri lensData.Position with
         | None -> return p |> success
         | Some symbol ->
-            let! refs = context.FindReferences symbol
+            let! locations = context.FindReferences symbol false
             // FIXME: refNum is wrong. There are lots of false positive even if we distinct locations by
-            // (l.Location.SourceTree.FilePath, l.Location.SourceSpan)
+            // (l.SourceTree.FilePath, l.SourceSpan)
             let refNum =
-                refs
-                |> Seq.collect (fun r -> r.Locations)
-                |> Seq.distinctBy (fun l -> (l.Location.SourceTree.FilePath, l.Location.SourceSpan))
+                locations
+                |> Seq.distinctBy (fun l -> (l.SourceTree.FilePath, l.SourceSpan))
                 |> Seq.length
 
             let title = sprintf "%d Reference(s)" refNum

diff --git a/src/CSharpLanguageServer/Handlers/References.fs b/src/CSharpLanguageServer/Handlers/References.fs
@@ -42,11 +42,11 @@ module References =
         match! context.FindSymbol p.TextDocument.Uri p.Position with
         | None -> return None |> success
         | Some symbol ->
-            let! refs = context.FindReferences symbol
+            let! locations = context.FindReferences symbol p.Context.IncludeDeclaration
+
             return
-                refs
-                |> Seq.collect (fun r -> r.Locations)
-                |> Seq.map (fun rl -> Location.fromRoslynLocation rl.Location)
+                locations
+                |> Seq.map Location.fromRoslynLocation
                 |> Seq.distinct
                 |> Seq.toArray
                 |> Some

diff --git a/src/CSharpLanguageServer/State/ServerRequestContext.fs b/src/CSharpLanguageServer/State/ServerRequestContext.fs
@@ -193,7 +193,7 @@ type ServerRequestContext (requestId: int, state: ServerState, emitServerEvent)
         return aggregatedLspLocations
     }
 
-    member this.FindSymbols (pattern: string option): Async<ISymbol seq> = async {
+    member this.FindSymbols (pattern: string option): Async<Microsoft.CodeAnalysis.ISymbol seq> = async {
         let findTask ct =
                 match pattern with
                 | Some pat ->
@@ -209,14 +209,24 @@ type ServerRequestContext (requestId: int, state: ServerState, emitServerEvent)
             return! findTask ct solution |> Async.AwaitTask
     }
 
-    member this.FindReferences (symbol: ISymbol): Async<ReferencedSymbol seq> = async {
+    member this.FindReferences (symbol: ISymbol) (withDefinition: bool): Async<Microsoft.CodeAnalysis.Location seq> = async {
         match this.State.Solution with
         | None -> return []
         | Some solution ->
             let! ct = Async.CancellationToken
-            return!
+
+            let locationsFromReferencedSym (r: ReferencedSymbol) =
+                let locations = r.Locations |> Seq.map (fun rl -> rl.Location)
+
+                match withDefinition with
+                | true -> locations |> Seq.append r.Definition.Locations
+                | false -> locations
+
+            let! refs =
                 SymbolFinder.FindReferencesAsync(symbol, solution, cancellationToken=ct)
                 |> Async.AwaitTask
+
+            return refs |> Seq.collect locationsFromReferencedSym
     }
 
     member this.GetDocumentVersion (uri: DocumentUri): int option =

diff --git a/tests/CSharpLanguageServer.Tests/CSharpLanguageServer.Tests.fsproj b/tests/CSharpLanguageServer.Tests/CSharpLanguageServer.Tests.fsproj
@@ -8,11 +8,13 @@
 
   <ItemGroup>
     <Compile Include="Tooling.fs" />
+    <!--
     <Compile Include="CodeActionTests.fs" />
     <Compile Include="DiagnosticTests.fs" />
     <Compile Include="DocumentationTests.fs" />
     <Compile Include="HoverTests.fs" />
     <Compile Include="InitializationTests.fs" />
+    -->
     <Compile Include="ReferenceTests.fs" />
   </ItemGroup>
 

diff --git a/tests/CSharpLanguageServer.Tests/ReferenceTests.fs b/tests/CSharpLanguageServer.Tests/ReferenceTests.fs
@@ -50,3 +50,36 @@ let testReferenceWorks() =
          |]
 
     Assert.AreEqual(expectedLocations1, locations1.Value)
+
+    //
+    // try references request at MethodA declaration on line 2
+    // (with IncludeDeclaration=true)
+    //
+    let referenceParams2: ReferenceParams =
+        { TextDocument = { Uri = classFile.Uri }
+          Position = { Line = 2u; Character = 16u }
+          WorkDoneToken = None
+          PartialResultToken = None
+          Context = { IncludeDeclaration = true }
+        }
+
+    let locations2: Location[] option = classFile.Request("textDocument/references", referenceParams2)
+
+    let expectedLocations2: Location array =
+        [|
+            { Uri = classFile.Uri.Substring(7)
+              Range = {
+                  Start = { Line = 2u; Character = 16u }
+                  End = { Line = 2u; Character = 23u }
+              }
+            }
+
+            { Uri = classFile.Uri.Substring(7)
+              Range = {
+                  Start = { Line = 8u; Character = 8u }
+                  End = { Line = 8u; Character = 15u }
+              }
+            }
+         |]
+
+    Assert.AreEqual(expectedLocations2, locations2.Value)