Share one whole-file classification between requests for a document version

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

@@ -15,6 +15,8 @@
 * Fix doubled F# diagnostics in tooltips. ([Issue #16360](https://github.com/dotnet/fsharp/issues/16360))
 * Fix `NotSupportedException` in the memory-mapped-file optimization when copying `ReadOnlyMemory` into `MemoryMappedFileViewStream`. ([Issue #20263](https://github.com/dotnet/fsharp/issues/20263))
 * Reduce allocations in the VS project options reactor: the command-line options and project options caches and the mailbox reply payloads now hold struct tuples, and `IProjectSite.CompilationBinOutputPath` returns `string voption` picked with a new `Array.tryPickV`. ([PR #20413](https://github.com/dotnet/fsharp/pull/20413))
+* Fix semantic colours of open documents disappearing: their classification was stored in the unopened-documents cache and only covered the requested span, so scrolling re-ran the checker or got nothing back, and a check that could not complete answered with no classifications, which clears the colours already shown. ([Issue #20445](https://github.com/dotnet/fsharp/issues/20445), [PR #20450](https://github.com/dotnet/fsharp/pull/20450))
+* Requests for one version of an open document share a single whole-file semantic classification, kept for the life of the document instead of two seconds, so scrolling a large file after a pause no longer re-walks every symbol in it. ([PR #20452](https://github.com/dotnet/fsharp/pull/20452))
 
 ### Changed
 

vsintegration/src/FSharp.Editor/Classification/ClassificationService.fs

@@ -4,10 +4,13 @@ namespace Microsoft.VisualStudio.FSharp.Editor
 
 open System
 open System.Composition
+open System.Collections.Concurrent
 open System.Collections.Generic
 open System.Collections.Immutable
 open System.Threading
+open System.Threading.Tasks
 open System.Runtime.Caching
+open System.Runtime.CompilerServices
 
 open Microsoft.CodeAnalysis
 open Microsoft.CodeAnalysis.Classification
@@ -30,6 +33,63 @@ open Microsoft.VisualStudio.FSharp.Editor.Telemetry
 type SemanticClassificationData = SemanticClassificationView
 type SemanticClassificationLookup = IReadOnlyDictionary<int, ResizeArray<SemanticClassificationItem>>
 
+/// What a classification answers for: the text it colours and the semantics it was checked against.
+/// An edit elsewhere in the project reclassifies this document without touching its text.
+[<Struct>]
+type internal ClassificationVersion =
+    {
+        TextVersion: VersionStamp
+        SemanticVersion: VersionStamp
+    }
+
+/// The whole-file semantic classification of one version of an open document.
+type internal OpenDocumentClassification =
+    {
+        Version: ClassificationVersion
+        Text: SourceText
+        Lookup: SemanticClassificationLookup
+    }
+
+type internal ClassifyWholeFile = CancellationToken -> Task<OpenDocumentClassification voption>
+
+/// One classification of a document version, shared by every request that arrives while it runs.
+/// The computation starts with the first waiter and is cancelled only when the last one leaves.
+[<Sealed>]
+type internal InFlightClassification(version: ClassificationVersion, compute: ClassifyWholeFile) =
+    let cts = new CancellationTokenSource()
+    let job = lazy (compute cts.Token)
+    let mutable waiters = 0
+
+    member _.Version = version
+
+    member _.IsCancelled = cts.IsCancellationRequested
+
+    member _.IsCompleted = job.IsValueCreated && job.Value.IsCompleted
+
+    member _.Join(cancellationToken: CancellationToken) : Task<OpenDocumentClassification voption> =
+        Interlocked.Increment &waiters |> ignore
+        let job = job.Value
+        let left = ref 0
+
+        let leave () =
+            if
+                Interlocked.Exchange(left, 1) = 0
+                && Interlocked.Decrement &waiters = 0
+                && not job.IsCompleted
+            then
+                cts.Cancel()
+
+        task {
+            use _ = cancellationToken.Register(fun () -> leave ())
+
+            try
+                let! _ = Task.WhenAny(job, Task.Delay(Timeout.Infinite, cancellationToken))
+                cancellationToken.ThrowIfCancellationRequested()
+                return! job
+            finally
+                leave ()
+        }
+
 [<Export(typeof<IFSharpClassificationService>)>]
 type internal FSharpClassificationService [<ImportingConstructor>] () =
 
@@ -146,8 +206,96 @@ type internal FSharpClassificationService [<ImportingConstructor>] () =
     static let unopenedDocumentsSemanticClassificationCache =
         new DocumentCache<SemanticClassificationLookup>("fsharp-unopened-documents-semantic-classification-cache", 5.)
 
-    static let openedDocumentsSemanticClassificationCache =
-        new DocumentCache<SemanticClassificationLookup>("fsharp-opened-documents-semantic-classification-cache", 2.)
+    // The classification of an open document's latest checked version. Roslyn asks for it span by span
+    // for as long as that version is on screen, and replaces a span's semantic tags with whatever comes
+    // back - so it is kept for the life of the document rather than expiring, and when the checker
+    // cannot answer (project loading or reloading, a superseded check) it is re-emitted rather than
+    // answering "nothing", which would strip the colours the user already sees.
+    static let openDocumentClassifications =
+        ConditionalWeakTable<DocumentId, OpenDocumentClassification>()
+
+    static let inFlightClassifications =
+        ConcurrentDictionary<DocumentId, InFlightClassification>()
+
+    static let remember (documentId: DocumentId) (classification: OpenDocumentClassification) =
+        // net472 has no ConditionalWeakTable.AddOrUpdate.
+        lock openDocumentClassifications (fun () ->
+            openDocumentClassifications.Remove documentId |> ignore
+            openDocumentClassifications.Add(documentId, classification))
+
+    // Only for the text it was computed from: the lookup names positions, so against edited text it
+    // would colour the wrong characters.
+    static let addLastGood (documentId: DocumentId) (sourceText: SourceText) (targetSpan: TextSpan) (result: List<ClassifiedSpan>) =
+        match openDocumentClassifications.TryGetValue documentId with
+        | true, classification when classification.Text.ContentEquals sourceText ->
+            addSemanticClassificationByLookup sourceText targetSpan classification.Lookup result
+        | _ -> ()
+
+    static let addLastGoodForCurrentText (document: Document) (targetSpan: TextSpan) (result: List<ClassifiedSpan>) =
+        match document.TryGetText() with
+        | true, sourceText -> addLastGood document.Id sourceText targetSpan result
+        | _ -> ()
+
+    static let classifyWholeFile (document: Document) (version: ClassificationVersion) (sourceText: SourceText) =
+        cancellableTask {
+            match! document.TryGetFSharpParseAndCheckResultsAsync(nameof (IFSharpClassificationService)) with
+            | ValueNone -> return ValueNone
+            | ValueSome(struct (_, checkResults)) ->
+                let classificationData =
+                    checkResults.GetSemanticClassification(None, RelatedSymbolUseKind.All)
+
+                // Every checked file resolves at least its enclosing module, so nothing here means
+                // the classification itself failed (SemanticClassification.fs recovers with an empty
+                // array). Remembering that would pin the version to no colours.
+                if classificationData.Length = 0 then
+                    return ValueNone
+                else
+                    let classification =
+                        {
+                            Version = version
+                            Text = sourceText
+                            Lookup = itemToSemanticClassificationLookup classificationData
+                        }
+
+                    remember document.Id classification
+                    return ValueSome classification
+        }
+
+    // Requests for the same version that overlap - split views, the taggers above and below the
+    // viewport - share one classification instead of each walking the whole file.
+    static let classifyOpenDocument (document: Document) (version: ClassificationVersion) (sourceText: SourceText) =
+        cancellableTask {
+            let! cancellationToken = CancellableTask.getCancellationToken ()
+
+            let start () =
+                InFlightClassification(version, classifyWholeFile document version sourceText)
+
+            let inFlight =
+                inFlightClassifications.AddOrUpdate(
+                    document.Id,
+                    (fun _ -> start ()),
+                    fun _ running ->
+                        if running.Version = version && not running.IsCancelled then
+                            running
+                        else
+                            start ()
+                )
+
+            try
+                return! inFlight.Join cancellationToken
+            finally
+                // A waiter that leaves early keeps the entry for those still waiting.
+                if inFlight.IsCompleted || inFlight.IsCancelled then
+                    (inFlightClassifications :> ICollection<KeyValuePair<_, _>>).Remove(KeyValuePair(document.Id, inFlight))
+                    |> ignore
+        }
+
+    // Which store a document lands in is not observable from its classifications - a miss only costs
+    // a recheck - so tests reach them directly to tell the branches apart.
+    static member internal OpenDocumentClassifications = openDocumentClassifications
+
+    static member internal UnopenedDocumentsSemanticClassificationCache =
+        unopenedDocumentsSemanticClassificationCache
 
     interface IFSharpClassificationService with
         // Do not perform classification if we don't have project options (#defines matter)
@@ -252,15 +400,25 @@ type internal FSharpClassificationService [<ImportingConstructor>] () =
                             use _eventDuration =
                                 TelemetryReporter.ReportSingleEventWithDuration(TelemetryEvents.AddSemanticClassifications, eventProps)
 
-                            let! classificationData = document.GetFSharpSemanticClassificationAsync(nameof (FSharpClassificationService))
-
-                            let classificationDataLookup = toSemanticClassificationLookup classificationData
-                            do! unopenedDocumentsSemanticClassificationCache.SetAsync(document, classificationDataLookup)
-                            addSemanticClassificationByLookup sourceText textSpan classificationDataLookup result
+                            match! document.TryGetFSharpSemanticClassificationAsync(nameof (FSharpClassificationService)) with
+                            | ValueNone -> ()
+                            | ValueSome classificationData ->
+                                let classificationDataLookup = toSemanticClassificationLookup classificationData
+                                do! unopenedDocumentsSemanticClassificationCache.SetAsync(document, classificationDataLookup)
+                                addSemanticClassificationByLookup sourceText textSpan classificationDataLookup result
                     else
 
-                        match! openedDocumentsSemanticClassificationCache.TryGetValueAsync document with
-                        | ValueSome classificationDataLookup ->
+                        let! textVersion = document.GetTextVersionAsync(cancellationToken)
+                        let! semanticVersion = document.Project.GetDependentSemanticVersionAsync(cancellationToken)
+
+                        let version =
+                            {
+                                TextVersion = textVersion
+                                SemanticVersion = semanticVersion
+                            }
+
+                        match openDocumentClassifications.TryGetValue document.Id with
+                        | true, classification when classification.Version = version ->
                             let eventProps: (string * obj) array =
                                 [|
                                     "context.document.project.id", document.Project.Id.Id.ToString()
@@ -273,8 +431,8 @@ type internal FSharpClassificationService [<ImportingConstructor>] () =
                             use _eventDuration =
                                 TelemetryReporter.ReportSingleEventWithDuration(TelemetryEvents.AddSemanticClassifications, eventProps)
 
-                            addSemanticClassificationByLookup sourceText textSpan classificationDataLookup result
-                        | ValueNone ->
+                            addSemanticClassificationByLookup sourceText textSpan classification.Lookup result
+                        | _ ->
 
                             let eventProps: (string * obj) array =
                                 [|
@@ -288,21 +446,13 @@ type internal FSharpClassificationService [<ImportingConstructor>] () =
                             use _eventDuration =
                                 TelemetryReporter.ReportSingleEventWithDuration(TelemetryEvents.AddSemanticClassifications, eventProps)
 
-                            let! _, checkResults = document.GetFSharpParseAndCheckResultsAsync(nameof (IFSharpClassificationService))
-
-                            let targetRange =
-                                RoslynHelpers.TextSpanToFSharpRange(document.FilePath, textSpan, sourceText)
-
-                            let classificationData =
-                                checkResults.GetSemanticClassification(Some targetRange, RelatedSymbolUseKind.All)
-
-                            if classificationData.Length > 0 then
-                                let classificationDataLookup = itemToSemanticClassificationLookup classificationData
-                                do! unopenedDocumentsSemanticClassificationCache.SetAsync(document, classificationDataLookup)
-
-                                addSemanticClassification sourceText textSpan classificationData result
+                            match! classifyOpenDocument document version sourceText with
+                            | ValueSome classification -> addSemanticClassificationByLookup sourceText textSpan classification.Lookup result
+                            | ValueNone -> addLastGood document.Id sourceText textSpan result
                 }
-                |> CancellableTask.ifCanceledReturn ()
+                // A cancellation that is not Roslyn's own (a superseded or aborted check surfaces as one)
+                // must not turn into an empty answer, which Roslyn would paint as "no colours".
+                |> CancellableTask.ifCanceledThen (fun () -> addLastGoodForCurrentText document textSpan result)
                 |> CancellableTask.startAsTask cancellationToken
 
         // Do not perform classification if we don't have project options (#defines matter)

vsintegration/src/FSharp.Editor/Common/CancellableTasks.fs

@@ -1160,6 +1160,17 @@ module CancellableTasks =
                     return value
             }
 
+        /// If this CancellableTask gets canceled for another reason than the token being canceled, run the fallback instead.
+        let inline ifCanceledThen ([<InlineIfLambda>] fallback: unit -> unit) (ctask: CancellableTask<unit>) =
+            cancellableTask {
+                let! ct = getCancellationToken ()
+
+                try
+                    return! ctask
+                with :? OperationCanceledException when ct.IsCancellationRequested = false ->
+                    return fallback ()
+            }
+
     /// <exclude />
     [<AutoOpen>]
     module MergeSourcesExtensions =

vsintegration/src/FSharp.Editor/Common/DocumentCache.fs

@@ -18,17 +18,26 @@ type DocumentCache<'Value when 'Value: not struct>(name: string, ?cacheItemPolic
     let policy =
         defaultArg cacheItemPolicy (CacheItemPolicy(SlidingExpiration = (TimeSpan.FromSeconds defaultSlidingExpiration)))
 
+    // A document's own text is not the whole story: an edit elsewhere in the project can change what
+    // its names mean while its text stands still, so the key has to cover both versions.
+    static let currentVersion (doc: Document) (ct: CancellationToken) =
+        task {
+            let! textVersion = doc.GetTextVersionAsync ct
+            let! semanticVersion = doc.Project.GetDependentSemanticVersionAsync ct
+            return textVersion, semanticVersion
+        }
+
     static let tryGetCachedValueAsync (doc: Document, cache: MemoryCache, ct: CancellationToken) =
         if ct.IsCancellationRequested then
             Task.FromCanceled<'Value voption>(ct)
         else
             task {
-                let! currentVersion = doc.GetTextVersionAsync ct
+                let! version = currentVersion doc ct
 
                 match cache.Get(doc.Id.ToString()) with
                 | null -> return ValueNone
-                | :? (VersionStamp * 'Value) as value ->
-                    if fst value = currentVersion then
+                | :? ((VersionStamp * VersionStamp) * 'Value) as value ->
+                    if fst value = version then
                         return ValueSome(snd value)
                     else
                         return ValueNone
@@ -40,8 +49,8 @@ type DocumentCache<'Value when 'Value: not struct>(name: string, ?cacheItemPolic
             Task.FromCanceled<unit>(ct)
         else
             task {
-                let! currentVersion = doc.GetTextVersionAsync ct
-                do cache.Set(doc.Id.ToString(), (currentVersion, value), policy)
+                let! version = currentVersion doc ct
+                do cache.Set(doc.Id.ToString(), (version, value), policy)
             }
 
     new(name: string, slidingExpirationSeconds: float) =