Cache semantic classifications and display from the cache while a solution is loading.

src/EditorFeatures/Core/Implementation/Classification/SemanticClassificationUtilities.cs

@@ -3,14 +3,19 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
+using System.Runtime.CompilerServices;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Classification;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Editor.Tagging;
 using Microsoft.CodeAnalysis.ErrorReporting;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.LanguageServices;
+using Microsoft.CodeAnalysis.PersistentStorage;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.CodeAnalysis.Text.Shared.Extensions;
@@ -22,6 +27,17 @@ namespace Microsoft.CodeAnalysis.Editor.Implementation.Classification
 {
     internal static class SemanticClassificationUtilities
     {
+        /// <summary>
+        /// Mapping from workspaces to a task representing when they are fully loaded.  While this task is not complete,
+        /// the workspace is still loading.  Once complete the workspace is loaded.  We store the task around, instead
+        /// of just awaiting <see cref="IWorkspaceStatusService.IsFullyLoadedAsync"/> as actually awaiting that call
+        /// takes non-neglible time (upwards of several hundred ms, to a second), whereas we can actually just use the
+        /// status of the <see cref="IWorkspaceStatusService.WaitUntilFullyLoadedAsync"/> task to know when we have
+        /// actually transitioned to a loaded state.
+        /// </summary>
+        private static readonly ConditionalWeakTable<Workspace, Task> s_workspaceToFullyLoadedStateTask =
+            new ConditionalWeakTable<Workspace, Task>();
+
         public static async Task ProduceTagsAsync(
             TaggerContext<IClassificationTag> context,
             DocumentSnapshotSpan spanToTag,
@@ -131,8 +147,8 @@ private static async Task ClassifySpansAsync(
                 {
                     var classifiedSpans = ClassificationUtilities.GetOrCreateClassifiedSpanList();
 
-                    await classificationService.AddSemanticClassificationsAsync(
-                        document, snapshotSpan.Span.ToTextSpan(), classifiedSpans, cancellationToken: cancellationToken).ConfigureAwait(false);
+                    await AddSemanticClassificationsAsync(
+                        document, snapshotSpan.Span.ToTextSpan(), classificationService, classifiedSpans, cancellationToken: cancellationToken).ConfigureAwait(false);
 
                     ClassificationUtilities.Convert(typeMap, snapshotSpan.Snapshot, classifiedSpans, context.AddTag);
                     ClassificationUtilities.ReturnClassifiedSpanList(classifiedSpans);
@@ -149,5 +165,57 @@ await classificationService.AddSemanticClassificationsAsync(
                 throw ExceptionUtilities.Unreachable;
             }
         }
+
+        private static async Task AddSemanticClassificationsAsync(
+            Document document,
+            TextSpan textSpan,
+            IClassificationService classificationService,
+            List<ClassifiedSpan> classifiedSpans,
+            CancellationToken cancellationToken)
+        {
+            var workspace = document.Project.Solution.Workspace;
+            var fullyLoadedStateTask = s_workspaceToFullyLoadedStateTask.GetValue(
+                workspace, w =>
+                {
+                    var workspaceLoadedService = workspace.Services.GetRequiredService<IWorkspaceStatusService>();
+                    return workspaceLoadedService.WaitUntilFullyLoadedAsync(CancellationToken.None);
+                });
+
+            // If we're not fully loaded try to read from the cache instead so that classifications appear up to date.
+            // New code will not be semantically classified, but will eventually when the project fully loads.
+            var isFullyLoaded = fullyLoadedStateTask.IsCompleted;
+            if (await TryAddSemanticClassificationsFromCacheAsync(document, textSpan, classifiedSpans, isFullyLoaded, cancellationToken).ConfigureAwait(false))
+                return;
+
+            await classificationService.AddSemanticClassificationsAsync(
+                document, textSpan, classifiedSpans, cancellationToken).ConfigureAwait(false);
+        }
+
+        private static async Task<bool> TryAddSemanticClassificationsFromCacheAsync(
+            Document document,
+            TextSpan textSpan,
+            List<ClassifiedSpan> classifiedSpans,
+            bool isFullyLoaded,
+            CancellationToken cancellationToken)
+        {
+            // Don't use the cache if we're fully loaded.  We should just compute values normally.
+            if (isFullyLoaded)
+                return false;
+
+            var semanticCacheService = document.Project.Solution.Workspace.Services.GetService<ISemanticClassificationCacheService>();
+            if (semanticCacheService == null)
+                return false;
+
+            var checksums = await document.State.GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
+            var checksum = checksums.Text;
+
+            var result = await semanticCacheService.GetCachedSemanticClassificationsAsync(
+                (DocumentKey)document, textSpan, checksum, cancellationToken).ConfigureAwait(false);
+            if (result.IsDefault)
+                return false;
+
+            classifiedSpans.AddRange(result);
+            return true;
+        }
     }
 }

src/EditorFeatures/Core/Tagging/AbstractAsynchronousTaggerProvider.TagSource.cs

@@ -89,16 +89,6 @@ private sealed partial class TagSource : ForegroundThreadAffinitizedObject
             /// </summary>
             private readonly CancellationTokenSource _initialComputationCancellationTokenSource = new CancellationTokenSource();
 
-            /// <summary>
-            /// Whether or not we've gotten any change notifications from our <see cref="ITaggerEventSource"/>.
-            /// The first time we hear about changes, we fast track getting tags and reporting 
-            /// them to the UI.
-            /// 
-            /// We use an int so we can use <see cref="Interlocked.CompareExchange(ref int, int, int)"/> 
-            /// to read/set this.
-            /// </summary>
-            private int _seenEventSourceChanged;
-
             public TaggerDelay AddedTagNotificationDelay => _dataSource.AddedTagNotificationDelay;
             public TaggerDelay RemovedTagNotificationDelay => _dataSource.RemovedTagNotificationDelay;
 

src/EditorFeatures/Core/Tagging/AbstractAsynchronousTaggerProvider.TagSource_ProduceTags.cs

@@ -13,6 +13,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Tagging;
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.Text;
 using Microsoft.VisualStudio.Text.Editor;
@@ -59,24 +60,13 @@ private void OnUIUpdatesResumed(object sender, EventArgs e)
 
             private void OnEventSourceChanged(object sender, TaggerEventArgs e)
             {
-                var result = Interlocked.CompareExchange(ref _seenEventSourceChanged, value: 1, comparand: 0);
-                if (result == 0)
-                {
-                    // this is the first time we're hearing about changes from our event-source.
-                    // Don't have any delay here.  We want to just compute the tags and display
-                    // them as soon as we possibly can.
-                    ComputeInitialTags();
-                }
-                else
-                {
-                    // First, cancel any previous requests (either still queued, or started).  We no longer
-                    // want to continue it if new changes have come in.
-                    _workQueue.CancelCurrentWork();
-                    RegisterNotification(
-                        () => RecomputeTagsForeground(initialTags: false),
-                        (int)e.Delay.ComputeTimeDelay().TotalMilliseconds,
-                        GetCancellationToken(initialTags: false));
-                }
+                // First, cancel any previous requests (either still queued, or started).  We no longer
+                // want to continue it if new changes have come in.
+                _workQueue.CancelCurrentWork();
+                RegisterNotification(
+                    () => RecomputeTagsForeground(initialTags: false),
+                    (int)e.Delay.ComputeTimeDelay().TotalMilliseconds,
+                    GetCancellationToken(initialTags: false));
             }
 
             private void OnCaretPositionChanged(object sender, CaretPositionChangedEventArgs e)
@@ -560,10 +550,28 @@ private void ProcessNewTagTrees(
                     // If we're on the foreground already, we can just update our internal state directly.
                     UpdateStateAndReportChanges(newTagTrees, bufferToChanges, newState, initialTags);
                 }
+                else if (initialTags)
+                {
+                    // If this is the initial set of tags, we fast-track a notification about whatever initial tags we
+                    // computed.  This way the UI is updated quickly for that initial set, and we don't have to wait a
+                    // potentially very long time as the foreground-thread-queue makes it way to our notification.
+                    //
+                    // Do this in a fire and forget manner, but ensure we notify the test harness of this so that it
+                    // doesn't try to acquire tag results prior to this work finishing.
+                    var asyncToken = this._asyncListener.BeginAsyncOperation(nameof(ProcessNewTagTrees));
+                    Task.Run(async () =>
+                    {
+                        await this.ThreadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+                        UpdateStateAndReportChanges(newTagTrees, bufferToChanges, newState, initialTags);
+                    }).CompletesAsyncOperation(asyncToken);
+                }
                 else
                 {
-                    // Otherwise report back on the foreground asap to update the state and let our 
-                    // clients know about the change.
+                    // Otherwise report back on the foreground to update the state and let our clients know about the
+                    // change.  This will go to the end of the foreground processing queue.  This will normally process
+                    // quickly once VS is loaded, but it may take some time initially when VS is loading and the UI
+                    // thread is highly occupied.  This helps ensure that we don't oversaturate the UI during a very
+                    // contended period of time.
                     RegisterNotification(() => UpdateStateAndReportChanges(
                         newTagTrees, bufferToChanges, newState, initialTags),
                         delay: 0,

src/VisualStudio/Core/Def/Implementation/SemanticClassificationCache/SemanticClassificationCacheIncrementalAnalyzerProvider.cs

@@ -0,0 +1,70 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#nullable enable
+
+using System;
+using System.Composition;
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.Classification;
+using Microsoft.CodeAnalysis.Host;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.Remote;
+using Microsoft.CodeAnalysis.Shared.Extensions;
+using Microsoft.CodeAnalysis.SolutionCrawler;
+
+namespace Microsoft.VisualStudio.LanguageServices.Implementation.SemanticClassificationCache
+{
+    [ExportIncrementalAnalyzerProvider(nameof(SemanticClassificationCacheIncrementalAnalyzerProvider), new[] { WorkspaceKind.Host }), Shared]
+    internal class SemanticClassificationCacheIncrementalAnalyzerProvider : IIncrementalAnalyzerProvider
+    {
+        [ImportingConstructor]
+        [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+        public SemanticClassificationCacheIncrementalAnalyzerProvider()
+        {
+        }
+
+        public IIncrementalAnalyzer? CreateIncrementalAnalyzer(Workspace workspace)
+        {
+            if (workspace is not VisualStudioWorkspace)
+                return null;
+
+            return new SemanticClassificationCacheIncrementalAnalyzer();
+        }
+
+        private class SemanticClassificationCacheIncrementalAnalyzer : IncrementalAnalyzerBase
+        {
+            public override async Task AnalyzeDocumentAsync(Document document, SyntaxNode bodyOpt, InvocationReasons reasons, CancellationToken cancellationToken)
+            {
+                if (!document.IsOpen())
+                    return;
+
+                var solution = document.Project.Solution;
+                var client = await RemoteHostClient.TryGetClientAsync(solution.Workspace, cancellationToken).ConfigureAwait(false);
+                if (client == null)
+                {
+                    // We don't do anything if we fail to get the external process.  That's the case when something has gone
+                    // wrong, or the user is explicitly choosing to run inproc only.   In neither of those cases do we want
+                    // to bog down the VS process with the work to semantically classify files.
+                    return;
+                }
+
+                var statusService = document.Project.Solution.Workspace.Services.GetRequiredService<IWorkspaceStatusService>();
+                var isFullyLoaded = await statusService.IsFullyLoadedAsync(cancellationToken).ConfigureAwait(false);
+                Debug.Assert(isFullyLoaded, "We should only be called by the incremental analyzer once the solution is fully loaded.");
+
+                await client.RunRemoteAsync(
+                    WellKnownServiceHubService.CodeAnalysis,
+                    nameof(IRemoteSemanticClassificationCacheService.CacheSemanticClassificationsAsync),
+                    document.Project.Solution,
+                    arguments: new object[] { document.Id, isFullyLoaded },
+                    callbackTarget: null,
+                    cancellationToken).ConfigureAwait(false);
+            }
+        }
+    }
+}

src/VisualStudio/Core/Def/Implementation/SemanticClassificationCache/VisualStudioSemanticClassificationCacheService.cs

@@ -0,0 +1,73 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#nullable enable
+
+using System;
+using System.Collections.Immutable;
+using System.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.Classification;
+using Microsoft.CodeAnalysis.Editor.Implementation.Classification;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.PersistentStorage;
+using Microsoft.CodeAnalysis.Remote;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.LanguageServices.Implementation.ProjectSystem;
+
+namespace Microsoft.VisualStudio.LanguageServices.Implementation.SemanticClassificationCache
+{
+    [ExportWorkspaceService(typeof(ISemanticClassificationCacheService), ServiceLayer.Host), Shared]
+    internal class VisualStudioSemanticClassificationCacheService
+        : ForegroundThreadAffinitizedObject, ISemanticClassificationCacheService
+    {
+        private readonly VisualStudioWorkspaceImpl _workspace;
+
+        [ImportingConstructor]
+        [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+        public VisualStudioSemanticClassificationCacheService(
+            VisualStudioWorkspaceImpl workspace,
+            IThreadingContext threadingContext)
+            : base(threadingContext)
+        {
+            _workspace = workspace;
+        }
+
+        public async Task<ImmutableArray<ClassifiedSpan>> GetCachedSemanticClassificationsAsync(
+            DocumentKey documentKey,
+            TextSpan textSpan,
+            Checksum checksum,
+            CancellationToken cancellationToken)
+        {
+            var client = await RemoteHostClient.TryGetClientAsync(_workspace, cancellationToken).ConfigureAwait(false);
+            if (client == null)
+            {
+                // We don't do anything if we fail to get the external process.  That's the case when something has gone
+                // wrong, or the user is explicitly choosing to run inproc only.   In neither of those cases do we want
+                // to bog down the VS process with the work to semantically classify files.
+                return default;
+            }
+
+            var classifiedSpans = await client.RunRemoteAsync<SerializableClassifiedSpans>(
+                WellKnownServiceHubService.CodeAnalysis,
+                nameof(IRemoteSemanticClassificationCacheService.GetCachedSemanticClassificationsAsync),
+                solution: null,
+                arguments: new object[] { documentKey.Dehydrate(), textSpan, checksum },
+                callbackTarget: null,
+                cancellationToken).ConfigureAwait(false);
+            if (classifiedSpans == null)
+                return default;
+
+            var list = ClassificationUtilities.GetOrCreateClassifiedSpanList();
+            classifiedSpans.Rehydrate(list);
+
+            var result = list.ToImmutableArray();
+            ClassificationUtilities.ReturnClassifiedSpanList(list);
+            return result;
+        }
+    }
+}