From 364c876889d96e86007b39cae072af84ca614b23 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Saulius=20Menkevi=C4=8Dius?=
 <sauliusmenkevicius@fastmail.com>
Date: Fri, 15 Nov 2024 09:35:25 +0200
Subject: [PATCH] Respect Context.IncludeDeclaration in handler for
 "textDocument/findReferences"

---
 src/CSharpLanguageServer/Handlers/CodeLens.fs |  9 +++--
 .../Handlers/References.fs                    |  8 ++---
 .../State/ServerRequestContext.fs             | 16 +++++++--
 .../CSharpLanguageServer.Tests.fsproj         |  2 ++
 .../ReferenceTests.fs                         | 33 +++++++++++++++++++
 5 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/src/CSharpLanguageServer/Handlers/CodeLens.fs b/src/CSharpLanguageServer/Handlers/CodeLens.fs
index 10bf462b..05a542af 100644
--- 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
index f1602d8f..81f39ac6 100644
--- 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
index 23bf796d..133395c5 100644
--- 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
index a7f4bac1..49d70498 100644
--- 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
index 70cc1708..05ffc04b 100644
--- 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)