Cache diagnostics produced during workspace pull, to avoid unnecessary computations

src/EditorFeatures/Test/CodeFixes/CodeFixServiceTests.cs

@@ -1084,8 +1084,9 @@ void M()
                 ? root.DescendantNodes().OfType<CodeAnalysis.CSharp.Syntax.VariableDeclarationSyntax>().First().Span
                 : root.DescendantNodes().OfType<CodeAnalysis.CSharp.Syntax.InvocationExpressionSyntax>().First().Span;
 
-            await diagnosticIncrementalAnalyzer.GetDiagnosticsAsync(
-                sourceDocument.Project.Solution, sourceDocument.Project.Id, sourceDocument.Id, includeSuppressedDiagnostics: true, includeNonLocalDocumentDiagnostics: true, CancellationToken.None);
+            await diagnosticIncrementalAnalyzer.GetDiagnosticsForIdsAsync(
+                sourceDocument.Project.Solution, sourceDocument.Project.Id, sourceDocument.Id, diagnosticIds: null, shouldIncludeAnalyzer: null, getDocuments: null,
+                includeSuppressedDiagnostics: true, includeLocalDocumentDiagnostics: true, includeNonLocalDocumentDiagnostics: true, CancellationToken.None);
             await diagnosticIncrementalAnalyzer.GetTestAccessor().TextDocumentOpenAsync(sourceDocument);
 
             var lowPriorityAnalyzers = new ConcurrentSet<DiagnosticAnalyzer>();

src/EditorFeatures/Test/Diagnostics/DiagnosticAnalyzerServiceTests.cs

@@ -12,9 +12,7 @@
 using Microsoft.CodeAnalysis.CSharp.RemoveUnnecessarySuppressions;
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Diagnostics.CSharp;
-using Microsoft.CodeAnalysis.Diagnostics.EngineV2;
 using Microsoft.CodeAnalysis.Editor.Test;
-using Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Remote.Diagnostics;

src/EditorFeatures/TestUtilities/Diagnostics/MockDiagnosticAnalyzerService.cs

@@ -61,7 +61,7 @@ public Task<ImmutableArray<DiagnosticData>> GetCachedDiagnosticsAsync(Workspace
         public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, bool includeSuppressedDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
             => throw new NotImplementedException();
 
-        public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
+        public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, Func<Project, DocumentId?, IReadOnlyList<DocumentId>>? getDocuments, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
             => throw new NotImplementedException();
 
         public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForSpanAsync(TextDocument document, TextSpan? range, Func<string, bool>? shouldIncludeDiagnostic, bool includeCompilerDiagnostics, bool includeSuppressedDiagnostics, ICodeActionRequestPriorityProvider priorityProvider, Func<string, IDisposable?>? addOperationScope, DiagnosticKind diagnosticKind, bool isExplicit, CancellationToken cancellationToken)

src/Features/Core/Portable/Diagnostics/IDiagnosticAnalyzerService.cs

@@ -9,6 +9,7 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Diagnostics;
@@ -94,11 +95,13 @@ internal interface IDiagnosticAnalyzerService
     /// complete set of non-local document diagnostics.
     /// </param>
     /// <param name="cancellationToken">Cancellation token.</param>
-    Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken);
+    Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, Func<Project, DocumentId?, IReadOnlyList<DocumentId>>? getDocumentIds, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken);
 
     /// <summary>
-    /// Get project diagnostics (diagnostics with no source location) of the given diagnostic ids and/or analyzers from the given solution. all diagnostics returned should be up-to-date with respect to the given solution.
-    /// Note that this method doesn't return any document diagnostics. Use <see cref="GetDiagnosticsForIdsAsync(Solution, ProjectId, DocumentId, ImmutableHashSet{string}, Func{DiagnosticAnalyzer, bool}?, bool, bool, bool, CancellationToken)"/> to also fetch those.
+    /// Get project diagnostics (diagnostics with no source location) of the given diagnostic ids and/or analyzers from
+    /// the given solution. all diagnostics returned should be up-to-date with respect to the given solution. Note that
+    /// this method doesn't return any document diagnostics. Use <see cref="GetDiagnosticsForIdsAsync"/> to also fetch
+    /// those.
     /// </summary>
     /// <param name="solution">Solution to fetch the diagnostics for.</param>
     /// <param name="projectId">Optional project to scope the returned diagnostics.</param>
@@ -200,4 +203,12 @@ public static Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForSpanAsync(th
             includeCompilerDiagnostics: true, includeSuppressedDiagnostics, priorityProvider,
             addOperationScope, diagnosticKind, isExplicit, cancellationToken);
     }
+
+    public static Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(
+        this IDiagnosticAnalyzerService service, Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
+    {
+        return service.GetDiagnosticsForIdsAsync(
+            solution, projectId, documentId, diagnosticIds, shouldIncludeAnalyzer, getDocumentIds: null,
+            includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics, cancellationToken);
+    }
 }

src/Features/LanguageServer/Protocol/Features/Diagnostics/DiagnosticAnalyzerService.cs

@@ -142,10 +142,10 @@ public async Task ForceAnalyzeProjectAsync(Project project, CancellationToken ca
         }
 
         public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(
-            Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
+            Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, Func<Project, DocumentId?, IReadOnlyList<DocumentId>>? getDocuments, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
         {
             var analyzer = CreateIncrementalAnalyzer(solution.Workspace);
-            return analyzer.GetDiagnosticsForIdsAsync(solution, projectId, documentId, diagnosticIds, shouldIncludeAnalyzer, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics, cancellationToken);
+            return analyzer.GetDiagnosticsForIdsAsync(solution, projectId, documentId, diagnosticIds, shouldIncludeAnalyzer, getDocuments, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics, cancellationToken);
         }
 
         public Task<ImmutableArray<DiagnosticData>> GetProjectDiagnosticsForIdsAsync(

src/Features/LanguageServer/Protocol/Features/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer_GetDiagnostics.cs

@@ -19,13 +19,13 @@ public Task<ImmutableArray<DiagnosticData>> GetCachedDiagnosticsAsync(Solution s
             => new IdeCachedDiagnosticGetter(this, solution, projectId, documentId, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics).GetDiagnosticsAsync(cancellationToken);
 
         public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, bool includeSuppressedDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
-            => new IdeLatestDiagnosticGetter(this, solution, projectId, documentId, diagnosticIds: null, shouldIncludeAnalyzer: null, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics: true, includeNonLocalDocumentDiagnostics).GetDiagnosticsAsync(cancellationToken);
+            => new IdeLatestDiagnosticGetter(this, solution, projectId, documentId, diagnosticIds: null, shouldIncludeAnalyzer: null, getDocuments: null, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics: true, includeNonLocalDocumentDiagnostics).GetDiagnosticsAsync(cancellationToken);
 
-        public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
-            => new IdeLatestDiagnosticGetter(this, solution, projectId, documentId, diagnosticIds, shouldIncludeAnalyzer, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics).GetDiagnosticsAsync(cancellationToken);
+        public Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, Func<Project, DocumentId?, IReadOnlyList<DocumentId>>? getDocuments, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
+            => new IdeLatestDiagnosticGetter(this, solution, projectId, documentId, diagnosticIds, shouldIncludeAnalyzer, getDocuments, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics).GetDiagnosticsAsync(cancellationToken);
 
         public Task<ImmutableArray<DiagnosticData>> GetProjectDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, bool includeSuppressedDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken)
-            => new IdeLatestDiagnosticGetter(this, solution, projectId, documentId: null, diagnosticIds, shouldIncludeAnalyzer, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics: false, includeNonLocalDocumentDiagnostics).GetProjectDiagnosticsAsync(cancellationToken);
+            => new IdeLatestDiagnosticGetter(this, solution, projectId, documentId: null, diagnosticIds, shouldIncludeAnalyzer, getDocuments: null, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics: false, includeNonLocalDocumentDiagnostics).GetProjectDiagnosticsAsync(cancellationToken);
 
         private abstract class DiagnosticGetter
         {
@@ -38,13 +38,16 @@ private abstract class DiagnosticGetter
             protected readonly bool IncludeLocalDocumentDiagnostics;
             protected readonly bool IncludeNonLocalDocumentDiagnostics;
 
+            private readonly Func<Project, DocumentId?, IReadOnlyList<DocumentId>> _getDocuments;
+
             private ImmutableArray<DiagnosticData>.Builder? _lazyDataBuilder;
 
             public DiagnosticGetter(
                 DiagnosticIncrementalAnalyzer owner,
                 Solution solution,
                 ProjectId? projectId,
                 DocumentId? documentId,
+                Func<Project, DocumentId?, IReadOnlyList<DocumentId>>? getDocuments,
                 bool includeSuppressedDiagnostics,
                 bool includeLocalDocumentDiagnostics,
                 bool includeNonLocalDocumentDiagnostics)
@@ -53,6 +56,7 @@ public DiagnosticGetter(
                 Solution = solution;
 
                 DocumentId = documentId;
+                _getDocuments = getDocuments ?? (static (project, documentId) => documentId != null ? [documentId] : project.DocumentIds);
                 ProjectId = projectId ?? documentId?.ProjectId;
 
                 IncludeSuppressedDiagnostics = includeSuppressedDiagnostics;
@@ -79,7 +83,7 @@ public async Task<ImmutableArray<DiagnosticData>> GetDiagnosticsAsync(Cancellati
                         return GetDiagnosticData();
                     }
 
-                    var documentIds = (DocumentId != null) ? SpecializedCollections.SingletonEnumerable(DocumentId) : project.DocumentIds;
+                    var documentIds = _getDocuments(project, DocumentId);
 
                     // return diagnostics specific to one project or document
                     var includeProjectNonLocalResult = DocumentId == null;
@@ -132,7 +136,7 @@ private bool ShouldIncludeSuppressedDiagnostic(DiagnosticData diagnostic)
         private sealed class IdeCachedDiagnosticGetter : DiagnosticGetter
         {
             public IdeCachedDiagnosticGetter(DiagnosticIncrementalAnalyzer owner, Solution solution, ProjectId? projectId, DocumentId? documentId, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics)
-                : base(owner, solution, projectId, documentId, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics)
+                : base(owner, solution, projectId, documentId, getDocuments: null, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics)
             {
             }
 
@@ -232,8 +236,9 @@ private sealed class IdeLatestDiagnosticGetter : DiagnosticGetter
             public IdeLatestDiagnosticGetter(
                 DiagnosticIncrementalAnalyzer owner, Solution solution, ProjectId? projectId,
                 DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer,
-                bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics)
-                : base(owner, solution, projectId, documentId, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics)
+                Func<Project, DocumentId?, IReadOnlyList<DocumentId>>? getDocuments,
+                 bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics)
+                : base(owner, solution, projectId, documentId, getDocuments, includeSuppressedDiagnostics, includeLocalDocumentDiagnostics, includeNonLocalDocumentDiagnostics)
             {
                 _diagnosticIds = diagnosticIds;
                 _shouldIncludeAnalyzer = shouldIncludeAnalyzer;

src/Features/LanguageServer/Protocol/Handler/Diagnostics/DiagnosticSources/AbstractWorkspaceDocumentDiagnosticSource.cs

@@ -3,10 +3,15 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections;
+using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Diagnostics;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics;
 
@@ -21,6 +26,13 @@ public static AbstractWorkspaceDocumentDiagnosticSource CreateForCodeAnalysisDia
     private sealed class FullSolutionAnalysisDiagnosticSource(TextDocument document, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer)
         : AbstractWorkspaceDocumentDiagnosticSource(document)
     {
+        /// <summary>
+        /// Cached mapping between a project instance and all the diagnostics computed for it.  This is used so that
+        /// once we compute the diagnostics once for a particular project, we don't need to recompute them again as we
+        /// walk every document within it.
+        /// </summary>
+        private static readonly ConditionalWeakTable<Project, AsyncLazy<IReadOnlyList<DiagnosticData>>> s_projectToDiagnostics = new();
+
         /// <summary>
         /// This is a normal document source that represents live/fresh diagnostics that should supersede everything else.
         /// </summary>
@@ -40,14 +52,35 @@ public override async Task<ImmutableArray<DiagnosticData>> GetDiagnosticsAsync(
             }
             else
             {
-                // We call GetDiagnosticsForIdsAsync as we want to ensure we get the full set of diagnostics for this document
-                // including those reported as a compilation end diagnostic.  These are not included in document pull (uses GetDiagnosticsForSpan) due to cost.
-                // However we can include them as a part of workspace pull when FSA is on.
-                var documentDiagnostics = await diagnosticAnalyzerService.GetDiagnosticsForIdsAsync(
-                    Document.Project.Solution, Document.Project.Id, Document.Id,
-                    diagnosticIds: null, shouldIncludeAnalyzer, includeSuppressedDiagnostics: false,
-                    includeLocalDocumentDiagnostics: true, includeNonLocalDocumentDiagnostics: true, cancellationToken).ConfigureAwait(false);
-                return documentDiagnostics;
+                var projectDiagnostics = await GetProjectDiagnosticsAsync(diagnosticAnalyzerService, cancellationToken).ConfigureAwait(false);
+                return projectDiagnostics.WhereAsArray(d => d.DocumentId == Document.Id);
+            }
+        }
+
+        private async ValueTask<ImmutableArray<DiagnosticData>> GetProjectDiagnosticsAsync(
+            IDiagnosticAnalyzerService diagnosticAnalyzerService, CancellationToken cancellationToken)
+        {
+            if (!s_projectToDiagnostics.TryGetValue(Document.Project, out var lazyDiagnostics))
+            {
+                // Extracted into local to prevent captures.
+                lazyDiagnostics = GetLazyDiagnostics();
+            }
+
+            var result = await lazyDiagnostics.GetValueAsync(cancellationToken).ConfigureAwait(false);
+            return (ImmutableArray<DiagnosticData>)result;
+
+            AsyncLazy<IReadOnlyList<DiagnosticData>> GetLazyDiagnostics()
+            {
+                return s_projectToDiagnostics.GetValue(
+                    Document.Project,
+                    _ => AsyncLazy.Create<IReadOnlyList<DiagnosticData>>(
+                        async cancellationToken => await diagnosticAnalyzerService.GetDiagnosticsForIdsAsync(
+                            Document.Project.Solution, Document.Project.Id, documentId: null,
+                            diagnosticIds: null, shouldIncludeAnalyzer,
+                            // Ensure we compute and return diagnostics for both the normal docs and the additional docs in this project.
+                            static (project, _) => [.. project.DocumentIds.Concat(project.AdditionalDocumentIds)],
+                            includeSuppressedDiagnostics: false,
+                            includeLocalDocumentDiagnostics: true, includeNonLocalDocumentDiagnostics: true, cancellationToken).ConfigureAwait(false)));
             }
         }
     }