Merge pull request #34762 from Therzok/allocation-diagnostics

Reduce allocations in DiagnosticIncrementalAnalyzer

src/Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.StateSet.cs

@@ -2,6 +2,7 @@
 
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -57,14 +58,12 @@ public StateSet(string language, DiagnosticAnalyzer analyzer, string errorSource
             public DiagnosticAnalyzer Analyzer => _analyzer;
             public VersionStamp AnalyzerVersion => _analyzerVersion;
 
+            [PerformanceSensitive("https://github.com/dotnet/roslyn/issues/34761", AllowCaptures = false, AllowGenericEnumeration = false)]
             public bool ContainsAnyDocumentOrProjectDiagnostics(ProjectId projectId)
             {
-                foreach (var state in GetActiveFileStates(projectId))
+                if (ContainsAnyDocumentDiagnostics(_activeFileStates, projectId))
                 {
-                    if (!state.IsEmpty)
-                    {
-                        return true;
-                    }
+                    return true;
                 }
 
                 if (!_projectStates.TryGetValue(projectId, out var projectState))
@@ -73,6 +72,24 @@ public bool ContainsAnyDocumentOrProjectDiagnostics(ProjectId projectId)
                 }
 
                 return !projectState.IsEmpty();
+
+                static bool ContainsAnyDocumentDiagnostics(ConcurrentDictionary<DocumentId, ActiveFileState> activeFileStates, ProjectId projectId)
+                {
+                    foreach (var kvp in activeFileStates)
+                    {
+                        var documentId = kvp.Key;
+                        if (documentId.ProjectId == projectId)
+                        {
+                            var state = kvp.Value;
+                            if (!state.IsEmpty)
+                            {
+                                return true;
+                            }
+                        }
+                    }
+
+                    return false;
+                }
             }
 
             public IEnumerable<ProjectId> GetProjectsWithDiagnostics()
@@ -103,29 +120,24 @@ public IEnumerable<ProjectId> GetProjectsWithDiagnostics()
                                                            .Select(kv => kv.Key)));
             }
 
-            public IEnumerable<DocumentId> GetDocumentsWithDiagnostics(ProjectId projectId)
+            [PerformanceSensitive("https://github.com/dotnet/roslyn/issues/34761", AllowCaptures = false, AllowGenericEnumeration = false)]
+            public void CollectDocumentsWithDiagnostics(ProjectId projectId, HashSet<DocumentId> set)
             {
-                HashSet<DocumentId> set = null;
-                foreach (var state in GetActiveFileStates(projectId))
+                Debug.Assert(set != null);
+
+                // Collect active documents with diagnostics
+                foreach (var kvp in _activeFileStates)
                 {
-                    set ??= new HashSet<DocumentId>();
-                    set.Add(state.DocumentId);
+                    if (kvp.Key.ProjectId == projectId)
+                    {
+                        set.Add(kvp.Value.DocumentId);
+                    }
                 }
 
-                if (!_projectStates.TryGetValue(projectId, out var projectState) || projectState.IsEmpty())
+                if (_projectStates.TryGetValue(projectId, out var projectState) && !projectState.IsEmpty())
                 {
-                    return set ?? SpecializedCollections.EmptyEnumerable<DocumentId>();
+                    set.UnionWith(projectState.GetDocumentsWithDiagnostics());
                 }
-
-                set ??= new HashSet<DocumentId>();
-                set.UnionWith(projectState.GetDocumentsWithDiagnostics());
-
-                return set;
-            }
-
-            private IEnumerable<ActiveFileState> GetActiveFileStates(ProjectId projectId)
-            {
-                return _activeFileStates.Where(kv => kv.Key.ProjectId == projectId).Select(kv => kv.Value);
             }
 
             public bool IsActiveFile(DocumentId documentId)

src/Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.cs

@@ -9,6 +9,7 @@
 using Microsoft.CodeAnalysis.CodeStyle;
 using Microsoft.CodeAnalysis.Diagnostics.Log;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Options;
 using Microsoft.CodeAnalysis.Simplification;
 using Microsoft.CodeAnalysis.SolutionCrawler;
@@ -115,14 +116,20 @@ public void Shutdown()
             Owner.RaiseBulkDiagnosticsUpdated(raiseEvents =>
             {
                 var handleActiveFile = true;
+                var documentSet = PooledHashSet<DocumentId>.GetInstance();
+
                 foreach (var stateSet in stateSets)
                 {
                     var projectIds = stateSet.GetProjectsWithDiagnostics();
                     foreach (var projectId in projectIds)
                     {
-                        RaiseProjectDiagnosticsRemoved(stateSet, projectId, stateSet.GetDocumentsWithDiagnostics(projectId), handleActiveFile, raiseEvents);
+                        stateSet.CollectDocumentsWithDiagnostics(projectId, documentSet);
+                        RaiseProjectDiagnosticsRemoved(stateSet, projectId, documentSet, handleActiveFile, raiseEvents);
+                        documentSet.Clear();
                     }
                 }
+
+                documentSet.Free();
             });
         }
 
@@ -131,6 +138,8 @@ private void ClearAllDiagnostics(ImmutableArray<StateSet> stateSets, ProjectId p
             Owner.RaiseBulkDiagnosticsUpdated(raiseEvents =>
             {
                 var handleActiveFile = true;
+                var documentSet = PooledHashSet<DocumentId>.GetInstance();
+
                 foreach (var stateSet in stateSets)
                 {
                     // PERF: don't fire events for ones that we dont have any diagnostics on
@@ -139,8 +148,12 @@ private void ClearAllDiagnostics(ImmutableArray<StateSet> stateSets, ProjectId p
                         continue;
                     }
 
-                    RaiseProjectDiagnosticsRemoved(stateSet, projectId, stateSet.GetDocumentsWithDiagnostics(projectId), handleActiveFile, raiseEvents);
+                    stateSet.CollectDocumentsWithDiagnostics(projectId, documentSet);
+                    RaiseProjectDiagnosticsRemoved(stateSet, projectId, documentSet, handleActiveFile, raiseEvents);
+                    documentSet.Clear();
                 }
+
+                documentSet.Free();
             });
         }