Merge pull request #1983 from sharwell/fixall-performance

Improve DocumentBasedFixAllProvider performance for multiple documents

Author: sharwell

StyleCop.Analyzers/StyleCop.Analyzers.CodeFixes/DocumentationRules/SA1626CodeFixProvider.cs

@@ -68,9 +68,8 @@ private class FixAll : DocumentBasedFixAllProvider
             protected override string CodeActionTitle =>
                 DocumentationResources.SA1626CodeFix;
 
-            protected override async Task<SyntaxNode> FixAllInDocumentAsync(FixAllContext fixAllContext, Document document)
+            protected override async Task<SyntaxNode> FixAllInDocumentAsync(FixAllContext fixAllContext, Document document, ImmutableArray<Diagnostic> diagnostics)
             {
-                var diagnostics = await fixAllContext.GetDocumentDiagnosticsAsync(document).ConfigureAwait(false);
                 if (diagnostics.IsEmpty)
                 {
                     return null;

StyleCop.Analyzers/StyleCop.Analyzers.CodeFixes/Helpers/CustomBatchFixAllProvider.cs

@@ -210,93 +210,14 @@ public virtual string GetFixAllTitle(FixAllContext fixAllContext)
             }
         }
 
-        public virtual async Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(FixAllContext fixAllContext)
+        public virtual Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(FixAllContext fixAllContext)
         {
-            var allDiagnostics = ImmutableArray<Diagnostic>.Empty;
-            var projectsToFix = ImmutableArray<Project>.Empty;
-
-            var document = fixAllContext.Document;
-            var project = fixAllContext.Project;
-
-            switch (fixAllContext.Scope)
-            {
-            case FixAllScope.Document:
-                if (document != null)
-                {
-                    var documentDiagnostics = await fixAllContext.GetDocumentDiagnosticsAsync(document).ConfigureAwait(false);
-                    return ImmutableDictionary<Document, ImmutableArray<Diagnostic>>.Empty.SetItem(document, documentDiagnostics);
-                }
-
-                break;
-
-            case FixAllScope.Project:
-                projectsToFix = ImmutableArray.Create(project);
-                allDiagnostics = await fixAllContext.GetAllDiagnosticsAsync(project).ConfigureAwait(false);
-                break;
-
-            case FixAllScope.Solution:
-                projectsToFix = project.Solution.Projects
-                    .Where(p => p.Language == project.Language)
-                    .ToImmutableArray();
-
-                var diagnostics = new ConcurrentDictionary<ProjectId, ImmutableArray<Diagnostic>>();
-                var tasks = new Task[projectsToFix.Length];
-                for (int i = 0; i < projectsToFix.Length; i++)
-                {
-                    fixAllContext.CancellationToken.ThrowIfCancellationRequested();
-                    var projectToFix = projectsToFix[i];
-                    tasks[i] = Task.Run(
-                        async () =>
-                        {
-                            var projectDiagnostics = await fixAllContext.GetAllDiagnosticsAsync(projectToFix).ConfigureAwait(false);
-                            diagnostics.TryAdd(projectToFix.Id, projectDiagnostics);
-                        }, fixAllContext.CancellationToken);
-                }
-
-                await Task.WhenAll(tasks).ConfigureAwait(false);
-                allDiagnostics = allDiagnostics.AddRange(diagnostics.SelectMany(i => i.Value));
-                break;
-            }
-
-            if (allDiagnostics.IsEmpty)
-            {
-                return ImmutableDictionary<Document, ImmutableArray<Diagnostic>>.Empty;
-            }
-
-            return await GetDocumentDiagnosticsToFixAsync(allDiagnostics, projectsToFix, fixAllContext.CancellationToken).ConfigureAwait(false);
+            return FixAllContextHelper.GetDocumentDiagnosticsToFixAsync(fixAllContext);
         }
 
-        public virtual async Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(FixAllContext fixAllContext)
+        public virtual Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(FixAllContext fixAllContext)
         {
-            var project = fixAllContext.Project;
-            if (project != null)
-            {
-                switch (fixAllContext.Scope)
-                {
-                case FixAllScope.Project:
-                    var diagnostics = await fixAllContext.GetProjectDiagnosticsAsync(project).ConfigureAwait(false);
-                    return ImmutableDictionary<Project, ImmutableArray<Diagnostic>>.Empty.SetItem(project, diagnostics);
-
-                case FixAllScope.Solution:
-                    var projectsAndDiagnostics = new ConcurrentDictionary<Project, ImmutableArray<Diagnostic>>();
-                    var options = new ParallelOptions() { CancellationToken = fixAllContext.CancellationToken };
-                    Parallel.ForEach(project.Solution.Projects, options, proj =>
-                    {
-                        fixAllContext.CancellationToken.ThrowIfCancellationRequested();
-                        var projectDiagnosticsTask = fixAllContext.GetProjectDiagnosticsAsync(proj);
-                        projectDiagnosticsTask.Wait(fixAllContext.CancellationToken);
-                        var projectDiagnostics = projectDiagnosticsTask.Result;
-                        if (projectDiagnostics.Any())
-                        {
-                            projectsAndDiagnostics.TryAdd(proj, projectDiagnostics);
-                        }
-                    });
-
-                    return projectsAndDiagnostics.ToImmutableDictionary();
-                }
-            }
-
-            return ImmutableDictionary<Project, ImmutableArray<Diagnostic>>.Empty;
+            return FixAllContextHelper.GetProjectDiagnosticsToFixAsync(fixAllContext);
         }
 
         public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEnumerable<CodeAction> codeActions, CancellationToken cancellationToken)
@@ -429,57 +350,6 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
             return currentSolution;
         }
 
-        private static async Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(
-            ImmutableArray<Diagnostic> diagnostics,
-            ImmutableArray<Project> projects,
-            CancellationToken cancellationToken)
-        {
-            var treeToDocumentMap = await GetTreeToDocumentMapAsync(projects, cancellationToken).ConfigureAwait(false);
-
-            var builder = ImmutableDictionary.CreateBuilder<Document, ImmutableArray<Diagnostic>>();
-            foreach (var documentAndDiagnostics in diagnostics.GroupBy(d => GetReportedDocument(d, treeToDocumentMap)))
-            {
-                cancellationToken.ThrowIfCancellationRequested();
-                var document = documentAndDiagnostics.Key;
-                var diagnosticsForDocument = documentAndDiagnostics.ToImmutableArray();
-                builder.Add(document, diagnosticsForDocument);
-            }
-
-            return builder.ToImmutable();
-        }
-
-        private static async Task<ImmutableDictionary<SyntaxTree, Document>> GetTreeToDocumentMapAsync(ImmutableArray<Project> projects, CancellationToken cancellationToken)
-        {
-            var builder = ImmutableDictionary.CreateBuilder<SyntaxTree, Document>();
-            foreach (var project in projects)
-            {
-                cancellationToken.ThrowIfCancellationRequested();
-                foreach (var document in project.Documents)
-                {
-                    cancellationToken.ThrowIfCancellationRequested();
-                    var tree = await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
-                    builder.Add(tree, document);
-                }
-            }
-
-            return builder.ToImmutable();
-        }
-
-        private static Document GetReportedDocument(Diagnostic diagnostic, ImmutableDictionary<SyntaxTree, Document> treeToDocumentsMap)
-        {
-            var tree = diagnostic.Location.SourceTree;
-            if (tree != null)
-            {
-                Document document;
-                if (treeToDocumentsMap.TryGetValue(tree, out document))
-                {
-                    return document;
-                }
-            }
-
-            return null;
-        }
-
         /// <summary>
         /// Try to merge the changes between <paramref name="newDocument"/> and <paramref name="oldDocument"/> into <paramref name="cumulativeChanges"/>.
         /// If there is any conflicting change in <paramref name="newDocument"/> with existing <paramref name="cumulativeChanges"/>, then the original <paramref name="cumulativeChanges"/> are returned.

StyleCop.Analyzers/StyleCop.Analyzers.CodeFixes/Helpers/DocumentBasedFixAllProvider.cs

@@ -58,16 +58,24 @@ public override Task<CodeAction> GetFixAsync(FixAllContext fixAllContext)
         /// </summary>
         /// <param name="fixAllContext">The context for the Fix All operation.</param>
         /// <param name="document">The document to fix.</param>
+        /// <param name="diagnostics">The diagnostics to fix in the document.</param>
         /// <returns>
         /// <para>The new <see cref="SyntaxNode"/> representing the root of the fixed document.</para>
         /// <para>-or-</para>
         /// <para><see langword="null"/>, if no changes were made to the document.</para>
         /// </returns>
-        protected abstract Task<SyntaxNode> FixAllInDocumentAsync(FixAllContext fixAllContext, Document document);
+        protected abstract Task<SyntaxNode> FixAllInDocumentAsync(FixAllContext fixAllContext, Document document, ImmutableArray<Diagnostic> diagnostics);
 
         private async Task<Document> GetDocumentFixesAsync(FixAllContext fixAllContext)
         {
-            var newRoot = await this.FixAllInDocumentAsync(fixAllContext, fixAllContext.Document).ConfigureAwait(false);
+            var documentDiagnosticsToFix = await FixAllContextHelper.GetDocumentDiagnosticsToFixAsync(fixAllContext).ConfigureAwait(false);
+            ImmutableArray<Diagnostic> diagnostics;
+            if (!documentDiagnosticsToFix.TryGetValue(fixAllContext.Document, out diagnostics))
+            {
+                return fixAllContext.Document;
+            }
+
+            var newRoot = await this.FixAllInDocumentAsync(fixAllContext, fixAllContext.Document, diagnostics).ConfigureAwait(false);
             if (newRoot == null)
             {
                 return fixAllContext.Document;
@@ -78,11 +86,20 @@ private async Task<Document> GetDocumentFixesAsync(FixAllContext fixAllContext)
 
         private async Task<Solution> GetSolutionFixesAsync(FixAllContext fixAllContext, ImmutableArray<Document> documents)
         {
+            var documentDiagnosticsToFix = await FixAllContextHelper.GetDocumentDiagnosticsToFixAsync(fixAllContext).ConfigureAwait(false);
+
             Solution solution = fixAllContext.Solution;
             List<Task<SyntaxNode>> newDocuments = new List<Task<SyntaxNode>>(documents.Length);
             foreach (var document in documents)
             {
-                newDocuments.Add(this.FixAllInDocumentAsync(fixAllContext, document));
+                ImmutableArray<Diagnostic> diagnostics;
+                if (!documentDiagnosticsToFix.TryGetValue(document, out diagnostics))
+                {
+                    newDocuments.Add(document.GetSyntaxRootAsync(fixAllContext.CancellationToken));
+                    continue;
+                }
+
+                newDocuments.Add(this.FixAllInDocumentAsync(fixAllContext, document, diagnostics));
             }
 
             for (int i = 0; i < documents.Length; i++)