Fix CodeFixService.GetProjectDiagnostics to handle compilation end...

Fix CodeFixService.GetProjectDiagnostics to handle compilation end diagnostics. Currently it was just force computing all document diagnostics, not the project diagnostics, which executes the compilation end actions.
Additionally, also fix the BatchFixAllProvider to correctly handle compilation end diagnostics.

src/Features/Core/CodeFixes/CodeFixService.cs

@@ -296,27 +296,32 @@ private async Task<IEnumerable<Diagnostic>> GetProjectDiagnosticsAsync(Project p
         {
             Contract.ThrowIfNull(project);
 
-            var diagnostics = await _diagnosticService.GetProjectDiagnosticsForIdsAsync(project.Solution, project.Id, diagnosticIds, cancellationToken: cancellationToken).ConfigureAwait(false);
-            Contract.ThrowIfFalse(diagnostics.All(d => d.DocumentId == null));
-            var dxs = diagnostics.Select(d => d.ToDiagnostic(null));
-
             if (includeAllDocumentDiagnostics)
             {
-                var list = new List<Diagnostic>();
-                list.AddRange(dxs);
-
-                foreach (var document in project.Documents)
-                {
-                    var docDiagnsotics = await GetDocumentDiagnosticsAsync(document, diagnosticIds, cancellationToken).ConfigureAwait(false);
-                    list.AddRange(docDiagnsotics);
-                }
-
-                return list;
+                // Get all diagnostics for the entire project, including document diagnostics.
+                var diagnostics = await _diagnosticService.GetDiagnosticsForIdsAsync(project.Solution, project.Id, diagnosticIds: diagnosticIds, cancellationToken: cancellationToken).ConfigureAwait(false);
+                var documentIdsToTreeMap = await GetDocumentIdsToTreeMapAsync(project, cancellationToken).ConfigureAwait(false);
+                return diagnostics.Select(d => d.DocumentId != null ? d.ToDiagnostic(documentIdsToTreeMap[d.DocumentId]) : d.ToDiagnostic(null));
             }
             else
             {
-                return dxs;
+                // Get all no-location diagnostics for the project, doesn't include document diagnostics.
+                var diagnostics = await _diagnosticService.GetProjectDiagnosticsForIdsAsync(project.Solution, project.Id, diagnosticIds, cancellationToken: cancellationToken).ConfigureAwait(false);
+                Contract.ThrowIfFalse(diagnostics.All(d => d.DocumentId == null));
+                return diagnostics.Select(d => d.ToDiagnostic(null));
+            }
+        }
+
+        private static async Task<ImmutableDictionary<DocumentId, SyntaxTree>> GetDocumentIdsToTreeMapAsync(Project project, CancellationToken cancellationToken)
+        {
+            var builder = ImmutableDictionary.CreateBuilder<DocumentId, SyntaxTree>();
+            foreach (var document in project.Documents)
+            {
+                var tree = await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
+                builder.Add(document.Id, tree);
             }
+
+            return builder.ToImmutable();
         }
 
         private async Task<bool> ContainsAnyFix(Document document, DiagnosticData diagnostic, CancellationToken cancellationToken)

src/Features/Core/Diagnostics/EngineV1/DiagnosticIncrementalAnalyzer_GetDiagnostics.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
@@ -35,12 +36,14 @@ public override Task<ImmutableArray<DiagnosticData>> GetDiagnosticsAsync(Solutio
 
         public override Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId projectId, DocumentId documentId, ImmutableHashSet<string> diagnosticIds, CancellationToken cancellationToken)
         {
-            return new IDELatestDiagnosticGetter(this, diagnosticIds).GetDiagnosticsAsync(solution, projectId, documentId, cancellationToken);
+            // Fix all code path, we can make computation concurrent if we are computing diagnostics across a project/solution.
+            return new IDELatestDiagnosticGetter(this, diagnosticIds, concurrent: documentId == null).GetDiagnosticsAsync(solution, projectId, documentId, cancellationToken);
         }
 
         public override Task<ImmutableArray<DiagnosticData>> GetProjectDiagnosticsForIdsAsync(Solution solution, ProjectId projectId, ImmutableHashSet<string> diagnosticIds, CancellationToken cancellationToken)
         {
-            return new IDELatestDiagnosticGetter(this, diagnosticIds).GetProjectDiagnosticsAsync(solution, projectId, cancellationToken);
+            // Fix all code path, we can make computation concurrent if we are computing project diagnostics across solution.
+            return new IDELatestDiagnosticGetter(this, diagnosticIds, concurrent: projectId == null).GetProjectDiagnosticsAsync(solution, projectId, cancellationToken);
         }
 
         private Task ReanalyzeAllDocumentsAsync(Project project, ImmutableHashSet<string> diagnosticIds, CancellationToken cancellationToken)
@@ -64,6 +67,8 @@ protected StateManager StateManager
                 get { return this.Owner._stateManger; }
             }
 
+            protected virtual bool ConcurrentDocumentComputation => false;
+
             protected abstract Task AppendDocumentDiagnosticsOfStateTypeAsync(Document document, StateType stateType, CancellationToken cancellationToken);
             protected abstract Task AppendProjectAndDocumentDiagnosticsAsync(Project project, Func<DiagnosticData, bool> predicate, CancellationToken cancellationToken);
 
@@ -157,9 +162,24 @@ private async Task AppendDiagnosticsAsync(Project project, CancellationToken can
 
                 await AppendProjectAndDocumentDiagnosticsAsync(project, cancellationToken).ConfigureAwait(false);
 
-                foreach (var document in project.Documents)
+                if (!ConcurrentDocumentComputation)
                 {
-                    await AppendDiagnosticsAsync(document, cancellationToken).ConfigureAwait(false);
+                    foreach (var document in project.Documents)
+                    {
+                        await AppendDiagnosticsAsync(document, cancellationToken).ConfigureAwait(false);
+                    }
+                }
+                else
+                {
+                    var documents = project.Documents.ToImmutableArray();
+                    var tasks = new Task[documents.Length];
+                    for (int i = 0; i < documents.Length; i++)
+                    {
+                        var document = documents[i];
+                        tasks[i] = Task.Run(async () => await AppendDiagnosticsAsync(document, cancellationToken).ConfigureAwait(false), cancellationToken);
+                    };
+
+                    await Task.WhenAll(tasks).ConfigureAwait(false);
                 }
             }
 
@@ -179,13 +199,13 @@ protected async Task AppendDiagnosticsAsync(Document document, CancellationToken
                 await AppendDocumentDiagnosticsOfStateTypeAsync(document, StateType.Document, cancellationToken).ConfigureAwait(false);
             }
 
-            protected void AppendDiagnostics(IEnumerable<DiagnosticData> items)
+            protected virtual void AppendDiagnostics(IEnumerable<DiagnosticData> items)
             {
                 _builder = _builder ?? ImmutableArray.CreateBuilder<DiagnosticData>();
                 _builder.AddRange(items);
             }
 
-            protected ImmutableArray<DiagnosticData> GetDiagnosticData()
+            protected virtual ImmutableArray<DiagnosticData> GetDiagnosticData()
             {
                 return _builder != null ? _builder.ToImmutable() : ImmutableArray<DiagnosticData>.Empty;
             }
@@ -312,7 +332,7 @@ protected override async Task AppendProjectAndDocumentDiagnosticsAsync(Project p
 
                     foreach (var document in project.Documents)
                     {
-                        await AppendProjectAndDocumentDiagnosticsAsync(state, project, predicate, cancellationToken).ConfigureAwait(false);
+                        await AppendProjectAndDocumentDiagnosticsAsync(state, document, predicate, cancellationToken).ConfigureAwait(false);
                     }
                 }
             }
@@ -493,12 +513,47 @@ protected override void FilterDiagnostics(AnalysisData analysisData, Func<Diagno
 
         private class IDELatestDiagnosticGetter : LatestDiagnosticsGetter
         {
-            public IDELatestDiagnosticGetter(DiagnosticIncrementalAnalyzer owner) : this(owner, null)
+            private readonly bool _concurrent;
+            private ConcurrentBag<DiagnosticData> _concurrentBag;
+
+            public IDELatestDiagnosticGetter(DiagnosticIncrementalAnalyzer owner, bool concurrent = false) : this(owner, null, concurrent)
+            {
+            }
+
+            public IDELatestDiagnosticGetter(DiagnosticIncrementalAnalyzer owner, ImmutableHashSet<string> diagnosticIds, bool concurrent = false) : base(owner, diagnosticIds)
+            {
+                _concurrent = concurrent;
+            }
+
+            protected override bool ConcurrentDocumentComputation => _concurrent;
+
+            protected override void AppendDiagnostics(IEnumerable<DiagnosticData> items)
             {
+                if (!ConcurrentDocumentComputation)
+                {
+                    base.AppendDiagnostics(items);
+                    return;
+                }
+
+                if (_concurrentBag == null)
+                {
+                    Interlocked.CompareExchange(ref _concurrentBag, new ConcurrentBag<DiagnosticData>(), null);
+                }
+
+                foreach (var item in items)
+                {
+                    _concurrentBag.Add(item);
+                }
             }
 
-            public IDELatestDiagnosticGetter(DiagnosticIncrementalAnalyzer owner, ImmutableHashSet<string> diagnosticIds) : base(owner, diagnosticIds)
+            protected override ImmutableArray<DiagnosticData> GetDiagnosticData()
             {
+                if (!ConcurrentDocumentComputation)
+                {
+                    return base.GetDiagnosticData();
+                }
+
+                return _concurrentBag != null ? _concurrentBag.ToImmutableArray() : ImmutableArray<DiagnosticData>.Empty;
             }
 
             public async Task<ImmutableArray<DiagnosticData>> GetSpecificDiagnosticsAsync(Solution solution, object id, CancellationToken cancellationToken)

src/Features/Core/Diagnostics/IDiagnosticAnalyzerService.cs

@@ -37,7 +37,7 @@ internal interface IDiagnosticAnalyzerService
 
         /// <summary>
         /// get diagnostics of the given diagnostic ids from the given solution. all diagnostics returned should be up-to-date with respect to the given solution.
-        /// Note that for project case, this metHod returns diagnostics from all project documents as well. Use <see cref="GetProjectDiagnosticsForIdsAsync(Solution, ProjectId, ImmutableHashSet{string}, CancellationToken)"/>
+        /// Note that for project case, this method returns diagnostics from all project documents as well. Use <see cref="GetProjectDiagnosticsForIdsAsync(Solution, ProjectId, ImmutableHashSet{string}, CancellationToken)"/>
         /// if you want to fetch only project diagnostics without source locations.
         /// </summary>
         Task<ImmutableArray<DiagnosticData>> GetDiagnosticsForIdsAsync(Solution solution, ProjectId projectId = null, DocumentId documentId = null, ImmutableHashSet<string> diagnosticIds = null, CancellationToken cancellationToken = default(CancellationToken));

src/Workspaces/Core/Portable/CodeFixes/FixAllOccurrences/BatchFixAllProvider.cs

@@ -217,10 +217,11 @@ public virtual string GetFixAllTitle(FixAllContext fixAllContext)
         {
             using (Logger.LogBlock(FunctionId.CodeFixes_FixAllOccurrencesComputation_Diagnostics, fixAllContext.CancellationToken))
             {
-                IEnumerable<Document> documentsToFix = null;
+                var allDiagnostics = ImmutableArray<Diagnostic>.Empty;
+                var projectsToFix = ImmutableArray<Project>.Empty;
+
                 var document = fixAllContext.Document;
                 var project = fixAllContext.Project;
-
                 var generatedCodeServices = project.Solution.Workspace.Services.GetService<IGeneratedCodeRecognitionService>();
 
                 switch (fixAllContext.Scope)
@@ -228,47 +229,101 @@ public virtual string GetFixAllTitle(FixAllContext fixAllContext)
                     case FixAllScope.Document:
                         if (document != null && !generatedCodeServices.IsGeneratedCode(document))
                         {
-                            var diagnostics = await fixAllContext.GetDocumentDiagnosticsAsync(document).ConfigureAwait(false);
-                            var kvp = SpecializedCollections.SingletonEnumerable(KeyValuePair.Create(document, diagnostics));
+                            var documentDiagnostics = await fixAllContext.GetDocumentDiagnosticsAsync(document).ConfigureAwait(false);
+                            var kvp = SpecializedCollections.SingletonEnumerable(KeyValuePair.Create(document, documentDiagnostics));
                             return ImmutableDictionary.CreateRange(kvp);
                         }
 
                         break;
 
                     case FixAllScope.Project:
-                        documentsToFix = project.Documents;
+                        projectsToFix = ImmutableArray.Create(project);
+                        allDiagnostics = await fixAllContext.GetAllDiagnosticsAsync(project).ConfigureAwait(false);
                         break;
 
                     case FixAllScope.Solution:
-                        documentsToFix = project.Solution.Projects
+                        projectsToFix = project.Solution.Projects
                             .Where(p => p.Language == project.Language)
-                            .SelectMany(p => p.Documents);
+                            .ToImmutableArray();
+
+                        var diagnostics = new ConcurrentBag<Diagnostic>();
+                        var tasks = new Task[projectsToFix.Length];
+                        for (int i = 0; i < projectsToFix.Length; i++)
+                        {
+                            var projectToFix = projectsToFix[i];
+                            tasks[i] = Task.Run(async () =>
+                            {
+                                var projectDiagnostics = await fixAllContext.GetAllDiagnosticsAsync(projectToFix).ConfigureAwait(false);
+                                foreach (var diagnostic in projectDiagnostics)
+                                {
+                                    diagnostics.Add(diagnostic);
+                                }
+                            }, fixAllContext.CancellationToken);
+                        }
+
+                        await Task.WhenAll(tasks).ConfigureAwait(false);
+                        allDiagnostics = allDiagnostics.AddRange(diagnostics);
                         break;
                 }
 
-                if (documentsToFix != null && documentsToFix.Any())
+                if (allDiagnostics.IsEmpty)
                 {
-                    var documentAndDiagnostics = new ConcurrentDictionary<Document, ImmutableArray<Diagnostic>>();
-                    var options = new ParallelOptions() { CancellationToken = fixAllContext.CancellationToken };
-                    Parallel.ForEach(documentsToFix, options, doc =>
-                    {
-                        fixAllContext.CancellationToken.ThrowIfCancellationRequested();
+                    return ImmutableDictionary<Document, ImmutableArray<Diagnostic>>.Empty;
+                }
 
-                        if (!generatedCodeServices.IsGeneratedCode(doc))
-                        {
-                            var documentDiagnostics = fixAllContext.GetDocumentDiagnosticsAsync(doc).WaitAndGetResult(fixAllContext.CancellationToken);
-                            if (documentDiagnostics.Any())
-                            {
-                                documentAndDiagnostics.TryAdd(doc, documentDiagnostics);
-                            }
-                        }
-                    });
+                return await GetDocumentDiagnosticsToFixAsync(allDiagnostics, projectsToFix, generatedCodeServices.IsGeneratedCode, fixAllContext.CancellationToken).ConfigureAwait(false);
+            }
+        }
 
-                    return documentAndDiagnostics.ToImmutableDictionary();
+        private async static Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(
+            ImmutableArray<Diagnostic> diagnostics,
+            ImmutableArray<Project> projects,
+            Func<Document, bool> isGeneratedCode, 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)))
+            {
+                var document = documentAndDiagnostics.Key;
+                if (!isGeneratedCode(document))
+                {
+                    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)
+            {
+                foreach (var document in project.Documents)
+                {
+                    var tree = await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
+                    builder.Add(tree, document);
                 }
+            }
 
-                return ImmutableDictionary<Document, ImmutableArray<Diagnostic>>.Empty;
+            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;
         }
 
         public virtual async Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(FixAllContext fixAllContext)
@@ -407,7 +462,7 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
                     });
                 }
 
-                Task.WaitAll(mergeTasks, cancellationToken);
+                await Task.WhenAll(mergeTasks).ConfigureAwait(false);
 
                 if (mergeFailed)
                 {