Classify each Find All References line while the search runs

docs/release-notes/.VisualStudio/18.vNext.md

@@ -5,6 +5,8 @@
 
 ### Fixed
 
+* Find All References classifies the line of each F# reference while it searches, as it does for C# and Visual Basic, instead of asking the F# classification service for every line again after the search. ([PR #20533](https://github.com/dotnet/fsharp/pull/20533))
+
 * Improve Find All References performance by throttling parallel typechecks. ([PR #20128](https://github.com/dotnet/fsharp/pull/20128))
 * Fixed Rename incorrectly renaming `get` and `set` keywords for properties with explicit accessors. ([Issue #18270](https://github.com/dotnet/fsharp/issues/18270), [PR #19252](https://github.com/dotnet/fsharp/pull/19252))
 * Fixed Find All References crash when F# project contains non-F# files like `.cshtml`. ([Issue #16394](https://github.com/dotnet/fsharp/issues/16394), [PR #19252](https://github.com/dotnet/fsharp/pull/19252))

vsintegration/src/FSharp.Editor/FSharp.Editor.fsproj

@@ -92,6 +92,7 @@
     <Compile Include="Navigation\GoToDefinitionService.fs" />
     <Compile Include="Navigation\NavigationBarItemService.fs" />
     <Compile Include="Navigation\NavigateToSearchService.fs" />
+    <Compile Include="Navigation\ClassifiedReferenceLine.fs" />
     <Compile Include="Navigation\FindUsagesService.fs" />
     <Compile Include="Navigation\FindDefinitionService.fs" />
     <Compile Include="QuickInfo\WpfFactories.fs" />

vsintegration/src/FSharp.Editor/Navigation/ClassifiedReferenceLine.fs

@@ -0,0 +1,100 @@
+// Copyright (c) Microsoft Corporation.  All Rights Reserved.  See License.txt in the project root for license information.
+
+/// The line a Find All References entry shows, classified while the search runs, as the C# and VB searches classify it
+/// with `ClassifiedSpansAndHighlightSpanFactory`, instead of by the window one entry at a time after the search.
+module internal Microsoft.VisualStudio.FSharp.Editor.ClassifiedReferenceLine
+
+open System
+open System.Collections.Generic
+open System.Collections.Immutable
+
+open Microsoft.CodeAnalysis
+open Microsoft.CodeAnalysis.Classification
+open Microsoft.CodeAnalysis.ExternalAccess.FSharp.Classification
+open Microsoft.CodeAnalysis.Text
+
+open CancellableTasks
+
+/// From the first non-whitespace character of the reference's line to the end of the line.
+let lineSpanOf (sourceText: SourceText) (referenceSpan: TextSpan) =
+    let line = sourceText.Lines.GetLineFromPosition referenceSpan.Start
+    let mutable start = line.Start
+
+    while start < line.End && Char.IsWhiteSpace sourceText[start] do
+        start <- start + 1
+
+    let firstNonWhitespace = if start = line.End then line.Start else start
+    TextSpan.FromBounds(min firstNonWhitespace referenceSpan.Start, line.End)
+
+let private byStart =
+    Comparison<ClassifiedSpan>(fun left right -> left.TextSpan.Start - right.TextSpan.Start)
+
+/// Sorted and clipped to the line. A span that starts before the one kept ahead of it ends is emptied, as
+/// `ClassifierHelper.AdjustSpans` empties it.
+let private adjusted (lineSpan: TextSpan) (spans: List<ClassifiedSpan>) =
+    spans.Sort byStart
+    let result = List<ClassifiedSpan>(spans.Count)
+
+    for span in spans do
+        let intersection = span.TextSpan.Intersection lineSpan
+
+        let kept =
+            intersection.HasValue
+            && (result.Count = 0
+                || result[result.Count - 1].TextSpan.End <= intersection.Value.Start)
+
+        result.Add(ClassifiedSpan(span.ClassificationType, (if kept then intersection.Value else TextSpan())))
+
+    result
+
+/// Every semantic span, and what no semantic span covers of each syntactic one.
+let private merged (syntactic: List<ClassifiedSpan>) (semantic: List<ClassifiedSpan>) =
+    let semanticParts = semantic.FindAll(fun span -> not span.TextSpan.IsEmpty)
+    let result = List<ClassifiedSpan>(semanticParts)
+
+    for part in syntactic do
+        let mutable start = part.TextSpan.Start
+
+        for covering in semanticParts do
+            if covering.TextSpan.Start < part.TextSpan.End && covering.TextSpan.End > start then
+                if covering.TextSpan.Start > start then
+                    result.Add(ClassifiedSpan(part.ClassificationType, TextSpan.FromBounds(start, covering.TextSpan.Start)))
+
+                start <- max start covering.TextSpan.End
+
+        if start < part.TextSpan.End then
+            result.Add(ClassifiedSpan(part.ClassificationType, TextSpan.FromBounds(start, part.TextSpan.End)))
+
+    result.Sort byStart
+    result
+
+let private withGapsFilled (lineStart: int) (spans: List<ClassifiedSpan>) =
+    let result = ImmutableArray.CreateBuilder<ClassifiedSpan>(spans.Count)
+    let mutable position = lineStart
+
+    for span in spans do
+        if not span.TextSpan.IsEmpty then
+            if position < span.TextSpan.Start then
+                result.Add(ClassifiedSpan(ClassificationTypeNames.Text, TextSpan.FromBounds(position, span.TextSpan.Start)))
+
+            result.Add span
+            position <- span.TextSpan.End
+
+    result.ToImmutable()
+
+/// The classified line of the reference, and the reference's place in it: what `ClassifiedSpansAndHighlightSpan` holds.
+let classifyAsync (classifier: IFSharpClassificationService) (document: Document) (sourceText: SourceText) (referenceSpan: TextSpan) =
+    cancellableTask {
+        let! ct = CancellableTask.getCancellationToken ()
+        let lineSpan = lineSpanOf sourceText referenceSpan
+        let syntactic = List<ClassifiedSpan>()
+        let semantic = List<ClassifiedSpan>()
+        do! classifier.AddSyntacticClassificationsAsync(document, lineSpan, syntactic, ct)
+        do! classifier.AddSemanticClassificationsAsync(document, lineSpan, semantic, ct)
+
+        let classifiedSpans =
+            merged (adjusted lineSpan syntactic) (adjusted lineSpan semantic)
+            |> withGapsFilled lineSpan.Start
+
+        return struct (classifiedSpans, TextSpan(referenceSpan.Start - lineSpan.Start, referenceSpan.Length))
+    }

vsintegration/src/FSharp.Editor/Navigation/FindUsagesService.fs

@@ -8,6 +8,7 @@ open System.Threading.Tasks
 
 open Microsoft.CodeAnalysis
 open Microsoft.CodeAnalysis.ExternalAccess.FSharp
+open Microsoft.CodeAnalysis.ExternalAccess.FSharp.Classification
 open Microsoft.CodeAnalysis.ExternalAccess.FSharp.FindUsages
 open Microsoft.CodeAnalysis.ExternalAccess.FSharp.Editor.FindUsages
 
@@ -48,11 +49,21 @@ module FSharpFindUsagesService =
                             |> Option.map (fun (definitionItem, _) -> definitionItem)
                             |> Option.defaultValue externalDefinitionItem
 
-                    let referenceItem =
-                        FSharpSourceReferenceItem(definitionItem, FSharpDocumentSpan(doc, fixedSpan))
+                    let classifier = FSharpClassificationService() :> IFSharpClassificationService
                     // REVIEW: OnReferenceFoundAsync is throwing inside Roslyn, putting a try/with so find-all refs doesn't fail.
                     try
-                        do! onReferenceFoundAsync referenceItem
+                        let! struct (classifiedSpans, highlightSpan) =
+                            ClassifiedReferenceLine.classifyAsync classifier doc sourceText fixedSpan
+
+                        do!
+                            onReferenceFoundAsync (
+                                FSharpSourceReferenceItem(
+                                    definitionItem,
+                                    FSharpDocumentSpan(doc, fixedSpan),
+                                    classifiedSpans,
+                                    highlightSpan
+                                )
+                            )
                     with _ ->
                         ()
                 | _ -> ()

vsintegration/tests/FSharp.Editor.Tests/ClassifiedReferenceLineTests.fs

@@ -0,0 +1,57 @@
+// Copyright (c) Microsoft Corporation.  All Rights Reserved.  See License.txt in the project root for license information.
+
+module FSharp.Editor.Tests.ClassifiedReferenceLineTests
+
+open Xunit
+
+open Microsoft.CodeAnalysis.Classification
+open Microsoft.CodeAnalysis.ExternalAccess.FSharp.Classification
+open Microsoft.CodeAnalysis.Text
+open Microsoft.VisualStudio.FSharp.Editor
+
+open FSharp.Editor.Tests.Helpers
+
+let private source =
+    """
+module M
+
+let add x y = x + y
+
+let result =
+    add 1 2
+"""
+
+[<Fact>]
+let ``A reference line is classified from its first non-whitespace character with every character covered`` () =
+    let document =
+        RoslynTestHelpers.CreateSolution(source).Projects
+        |> Seq.exactlyOne
+        |> _.Documents
+        |> Seq.exactlyOne
+
+    let sourceText = document.GetTextAsync().Result
+
+    let referenceSpan =
+        TextSpan(source.LastIndexOf("add", System.StringComparison.Ordinal), 3)
+
+    let classifier = FSharpClassificationService() :> IFSharpClassificationService
+
+    let struct (classifiedSpans, highlightSpan) =
+        (ClassifiedReferenceLine.classifyAsync classifier document sourceText referenceSpan) System.Threading.CancellationToken.None
+        |> _.Result
+
+    let line = sourceText.Lines.GetLineFromPosition referenceSpan.Start
+    Assert.Equal("add 1 2", sourceText.ToString(TextSpan.FromBounds(classifiedSpans[0].TextSpan.Start, line.End)))
+    Assert.Equal(line.End, classifiedSpans[classifiedSpans.Length - 1].TextSpan.End)
+
+    for previous, next in Seq.pairwise classifiedSpans do
+        Assert.Equal(previous.TextSpan.End, next.TextSpan.Start)
+
+    Assert.Equal(TextSpan(0, 3), highlightSpan)
+
+    let reference =
+        classifiedSpans |> Seq.find (fun span -> span.TextSpan = referenceSpan)
+
+    Assert.NotEqual<string>(ClassificationTypeNames.Text, reference.ClassificationType)
+
+    Assert.Contains(classifiedSpans, (fun span -> span.ClassificationType = ClassificationTypeNames.NumericLiteral))

vsintegration/tests/FSharp.Editor.Tests/FSharp.Editor.Tests.fsproj

@@ -27,6 +27,7 @@
     <Compile Include="IndentationServiceTests.fs" />
     <Compile Include="CompletionProviderTests.fs" />
     <Compile Include="FindReferencesTests.fs" />
+    <Compile Include="ClassifiedReferenceLineTests.fs" />
     <Compile Include="GoToDefinitionServiceTests.fs" />
     <Compile Include="HelpContextServiceTests.fs" />
     <Compile Include="QuickInfoTests.fs" />