Move GetDocumentsWithDiagnostics to a caller allocated pooled, reused hashset

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;
@@ -120,34 +121,22 @@ public IEnumerable<ProjectId> GetProjectsWithDiagnostics()
             }
 
             [PerformanceSensitive("https://github.com/dotnet/roslyn/issues/34761", AllowCaptures = false, AllowGenericEnumeration = false)]
-            public IEnumerable<DocumentId> GetDocumentsWithDiagnostics(ProjectId projectId)
+            public void CollectDocumentsWithDiagnostics(ProjectId projectId, HashSet<DocumentId> set)
             {
-                var set = GetActiveDocumentsForProject(_activeFileStates, projectId);
+                Debug.Assert(set != null);
 
-                if (!_projectStates.TryGetValue(projectId, out var projectState) || projectState.IsEmpty())
+                // Collect active documents with diagnostics
+                foreach (var kvp in _activeFileStates)
                 {
-                    return set ?? SpecializedCollections.EmptyEnumerable<DocumentId>();
-                }
-
-                set ??= new HashSet<DocumentId>();
-                set.UnionWith(projectState.GetDocumentsWithDiagnostics());
-
-                return set;
-
-                static HashSet<DocumentId> GetActiveDocumentsForProject(ConcurrentDictionary<DocumentId, ActiveFileState> activeFileStates, ProjectId projectId)
-                {
-                    HashSet<DocumentId> result = null;
-
-                    foreach (var kvp in activeFileStates)
+                    if (kvp.Key.ProjectId == projectId)
                     {
-                        if (kvp.Key.ProjectId == projectId)
-                        {
-                            result = result ?? new HashSet<DocumentId>();
-                            result.Add(kvp.Value.DocumentId);
-                        }
+                        set.Add(kvp.Value.DocumentId);
                     }
+                }
 
-                    return result;
+                if (_projectStates.TryGetValue(projectId, out var projectState) && !projectState.IsEmpty())
+                {
+                    set.UnionWith(projectState.GetDocumentsWithDiagnostics());
                 }
             }
 

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();
             });
         }