Make constructors of FixAllContext public and add a new nested abstract public...

Make constructors of FixAllContext public and add a new nested abstract public type "FixAllContext.DiagnosticProvider" for fix all diagnostic queries.

This unblocks a few product/test FixAll scenarios, see #2709 for details.

src/EditorFeatures/Test/Diagnostics/AbstractDiagnosticProviderBasedUserDiagnosticTest.cs

@@ -87,8 +87,9 @@ internal override IEnumerable<Diagnostic> GetDiagnostics(TestWorkspace workspace
                             return Task.FromResult(diags);
                         };
 
-                    var fixAllContext = new FixAllContext(document, fixer, scope, fixAllActionId,
-                        SpecializedCollections.SingletonEnumerable(diagnostic.Id), getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync, CancellationToken.None);
+                    var diagnosticIds = ImmutableHashSet.Create(diagnostic.Id);
+                    var fixAllDiagnosticProvider = new FixAllCodeActionContext.FixAllDiagnosticProvider(diagnosticIds, getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync);
+                    var fixAllContext = new FixAllContext(document, fixer, scope, fixAllActionId, diagnosticIds, fixAllDiagnosticProvider, CancellationToken.None);
                     var fixAllFix = fixAllProvider.GetFixAsync(fixAllContext).WaitAndGetResult(CancellationToken.None);
                     if (fixAllFix != null)
                     {

src/Features/Core/CodeFixes/CodeFixService.cs

@@ -277,7 +277,7 @@ public async Task<IEnumerable<CodeFixCollection>> GetFixesAsync(Document documen
 
                     if (fixAllProviderInfo != null)
                     {
-                        fixAllContext = new FixAllCodeActionContext(document, fixAllProviderInfo, codeFixProvider, diagnostics, this.GetDocumentDiagnosticsAsync, this.GetProjectDiagnosticsAsync, cancellationToken);
+                        fixAllContext = FixAllCodeActionContext.Create(document, fixAllProviderInfo, codeFixProvider, diagnostics, this.GetDocumentDiagnosticsAsync, this.GetProjectDiagnosticsAsync, cancellationToken);
                     }
                 }
 

src/Features/Core/CodeFixes/FixAllOccurrences/FixAllCodeActionContext.DiagnosticProvider.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.Generic;
+using System.Collections.Immutable;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.CodeFixes
+{
+    /// <summary>
+    /// FixAll context with some additional information specifically for <see cref="FixAllCodeAction"/>.
+    /// </summary>
+    internal partial class FixAllCodeActionContext : FixAllContext
+    {
+        internal class FixAllDiagnosticProvider : DiagnosticProvider
+        {
+            private readonly ImmutableHashSet<string> _diagnosticIds;
+
+            /// <summary>
+            /// Delegate to fetch diagnostics for any given document within the given fix all scope.
+            /// This delegate is invoked by <see cref="GetDocumentDiagnosticsAsync(Document, CancellationToken)"/> with the given <see cref="_diagnosticIds"/> as arguments.
+            /// </summary>
+            private readonly Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> _getDocumentDiagnosticsAsync;
+
+            /// <summary>
+            /// Delegate to fetch diagnostics for any given project within the given fix all scope.
+            /// This delegate is invoked by <see cref="GetProjectDiagnosticsAsync(Project, CancellationToken)"/> and <see cref="GetAllDiagnosticsAsync(Project, CancellationToken)"/>
+            /// with the given <see cref="_diagnosticIds"/> as arguments.
+            /// The boolean argument to the delegate indicates whether or not to return location-based diagnostics, i.e.
+            /// (a) False => Return only diagnostics with <see cref="Location.None"/>.
+            /// (b) True => Return all project diagnostics, regardless of whether or not they have a location.
+            /// </summary>
+            private readonly Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> _getProjectDiagnosticsAsync;
+
+            public FixAllDiagnosticProvider(
+                ImmutableHashSet<string> diagnosticIds,
+                Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getDocumentDiagnosticsAsync,
+                Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getProjectDiagnosticsAsync)
+            {
+                _diagnosticIds = diagnosticIds;
+                _getDocumentDiagnosticsAsync = getDocumentDiagnosticsAsync;
+                _getProjectDiagnosticsAsync = getProjectDiagnosticsAsync;
+            }
+
+            public override Task<IEnumerable<Diagnostic>> GetDocumentDiagnosticsAsync(Document document, CancellationToken cancellationToken)
+            {
+                return _getDocumentDiagnosticsAsync(document, _diagnosticIds, cancellationToken);
+            }
+
+            public override Task<IEnumerable<Diagnostic>> GetAllDiagnosticsAsync(Project project, CancellationToken cancellationToken)
+            {
+                return _getProjectDiagnosticsAsync(project, true, _diagnosticIds, cancellationToken);
+            }
+
+            public override Task<IEnumerable<Diagnostic>> GetProjectDiagnosticsAsync(Project project, CancellationToken cancellationToken)
+            {
+                return _getProjectDiagnosticsAsync(project, false, _diagnosticIds, cancellationToken);
+            }
+        }
+    }
+}

src/Features/Core/CodeFixes/FixAllOccurrences/FixAllCodeActionContext.cs

@@ -13,14 +13,13 @@ namespace Microsoft.CodeAnalysis.CodeFixes
     /// <summary>
     /// FixAll context with some additional information specifically for <see cref="FixAllCodeAction"/>.
     /// </summary>
-    internal class FixAllCodeActionContext : FixAllContext
+    internal partial class FixAllCodeActionContext : FixAllContext
     {
         private readonly FixAllProviderInfo _fixAllProviderInfo;
         private readonly IEnumerable<Diagnostic> _originalFixDiagnostics;
-        private readonly Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> _getDocumentDiagnosticsAsync;
-        private readonly Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> _getProjectDiagnosticsAsync;
+        private readonly FixAllDiagnosticProvider _diagnosticProvider;
 
-        internal FixAllCodeActionContext(
+        internal static FixAllCodeActionContext Create(
             Document document,
             FixAllProviderInfo fixAllProviderInfo,
             CodeFixProvider originalFixProvider,
@@ -28,17 +27,13 @@ internal class FixAllCodeActionContext : FixAllContext
             Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getDocumentDiagnosticsAsync,
             Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getProjectDiagnosticsAsync,
             CancellationToken cancellationToken)
-            : base(document, originalFixProvider, FixAllScope.Document,
-                  null, GetFixAllDiagnosticIds(fixAllProviderInfo, originalFixDiagnostics),
-                  getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync, cancellationToken)
         {
-            _fixAllProviderInfo = fixAllProviderInfo;
-            _originalFixDiagnostics = originalFixDiagnostics;
-            _getDocumentDiagnosticsAsync = getDocumentDiagnosticsAsync;
-            _getProjectDiagnosticsAsync = getProjectDiagnosticsAsync;
+            var diagnosticIds = GetFixAllDiagnosticIds(fixAllProviderInfo, originalFixDiagnostics).ToImmutableHashSet();
+            var diagnosticProvider = new FixAllDiagnosticProvider(diagnosticIds, getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync);
+            return new FixAllCodeActionContext(document, fixAllProviderInfo, originalFixProvider, originalFixDiagnostics, diagnosticIds, diagnosticProvider, cancellationToken);
         }
 
-        internal FixAllCodeActionContext(
+        internal static FixAllCodeActionContext Create(
             Project project,
             FixAllProviderInfo fixAllProviderInfo,
             CodeFixProvider originalFixProvider,
@@ -46,14 +41,40 @@ internal class FixAllCodeActionContext : FixAllContext
             Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getDocumentDiagnosticsAsync,
             Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getProjectDiagnosticsAsync,
             CancellationToken cancellationToken)
-            : base(project, originalFixProvider, FixAllScope.Project,
-                  null, GetFixAllDiagnosticIds(fixAllProviderInfo, originalFixDiagnostics),
-                  getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync, cancellationToken)
+        {
+            var diagnosticIds = GetFixAllDiagnosticIds(fixAllProviderInfo, originalFixDiagnostics).ToImmutableHashSet();
+            var diagnosticProvider = new FixAllDiagnosticProvider(diagnosticIds, getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync);
+            return new FixAllCodeActionContext(project, fixAllProviderInfo, originalFixProvider, originalFixDiagnostics, diagnosticIds, diagnosticProvider, cancellationToken);
+        }
+
+        private FixAllCodeActionContext(
+            Document document,
+            FixAllProviderInfo fixAllProviderInfo,
+            CodeFixProvider originalFixProvider,
+            IEnumerable<Diagnostic> originalFixDiagnostics,
+            ImmutableHashSet<string> diagnosticIds,
+            FixAllDiagnosticProvider diagnosticProvider,
+            CancellationToken cancellationToken)
+            : base(document, originalFixProvider, FixAllScope.Document, null, diagnosticIds, diagnosticProvider, cancellationToken)
+        {
+            _fixAllProviderInfo = fixAllProviderInfo;
+            _originalFixDiagnostics = originalFixDiagnostics;
+            _diagnosticProvider = diagnosticProvider;
+        }
+
+        private FixAllCodeActionContext(
+            Project project,
+            FixAllProviderInfo fixAllProviderInfo,
+            CodeFixProvider originalFixProvider,
+            IEnumerable<Diagnostic> originalFixDiagnostics,
+            ImmutableHashSet<string> diagnosticIds,
+            FixAllDiagnosticProvider diagnosticProvider,
+            CancellationToken cancellationToken)
+            : base(project, originalFixProvider, FixAllScope.Project, null, diagnosticIds, diagnosticProvider, cancellationToken)
         {
             _fixAllProviderInfo = fixAllProviderInfo;
             _originalFixDiagnostics = originalFixDiagnostics;
-            _getDocumentDiagnosticsAsync = getDocumentDiagnosticsAsync;
-            _getProjectDiagnosticsAsync = getProjectDiagnosticsAsync;
+            _diagnosticProvider = diagnosticProvider;
         }
 
         private static IEnumerable<string> GetFixAllDiagnosticIds(FixAllProviderInfo fixAllProviderInfo, IEnumerable<Diagnostic> originalFixDiagnostics)
@@ -91,11 +112,11 @@ internal FixAllContext GetContextForScopeAndActionId(FixAllScope scope, string c
             if (this.Document != null)
             {
                 return new FixAllContext(this.Document, this.CodeFixProvider, scope, codeActionEquivalenceKey,
-                    this.DiagnosticIds, _getDocumentDiagnosticsAsync, _getProjectDiagnosticsAsync, this.CancellationToken);
+                    this.DiagnosticIds, _diagnosticProvider, this.CancellationToken);
             }
 
             return new FixAllContext(this.Project, this.CodeFixProvider, scope, codeActionEquivalenceKey,
-                    this.DiagnosticIds, _getDocumentDiagnosticsAsync, _getProjectDiagnosticsAsync, this.CancellationToken);
+                    this.DiagnosticIds, _diagnosticProvider, this.CancellationToken);
         }
     }
 }

src/Features/Core/Features.csproj

@@ -118,6 +118,7 @@
     <Compile Include="CodeFixes\Async\AbstractChangeToAsyncCodeFixProvider.cs" />
     <Compile Include="CodeFixes\Async\AbstractAsyncCodeFix.cs" />
     <Compile Include="CodeFixes\CodeFixCollection.cs" />
+    <Compile Include="CodeFixes\FixAllOccurrences\FixAllCodeActionContext.DiagnosticProvider.cs" />
     <Compile Include="CodeFixes\FixAllOccurrences\FixAllProviderInfo.cs" />
     <Compile Include="CodeFixes\FixAllOccurrences\FixAllCodeActionContext.cs" />
     <Compile Include="CodeFixes\CodeFixService.cs" />

src/Workspaces/Core/Portable/CodeFixes/FixAllOccurrences/FixAllContext.DiagnosticProvider.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.Collections.Generic;
+using System.Threading.Tasks;
+using System.Threading;
+
+namespace Microsoft.CodeAnalysis.CodeFixes
+{
+    /// <summary>
+    /// Context for "Fix all occurrences" code fixes provided by an <see cref="FixAllProvider"/>.
+    /// </summary>
+    public partial class FixAllContext
+    {
+        /// <summary>
+        /// Diagnostic provider to fetch document/project diagnostics to fix in a <see cref="FixAllContext"/>.
+        /// </summary>
+        public abstract class DiagnosticProvider
+        {
+            /// <summary>
+            /// Gets all the diagnostics to fix in the given document in a <see cref="FixAllContext"/>.
+            /// </summary>
+            public abstract Task<IEnumerable<Diagnostic>> GetDocumentDiagnosticsAsync(Document document, CancellationToken cancellationToken);
+
+            /// <summary>
+            /// Gets all the project-level diagnostics to fix, i.e. diagnostics with no source location, in the given project in a <see cref="FixAllContext"/>.
+            /// </summary>
+            public abstract Task<IEnumerable<Diagnostic>> GetProjectDiagnosticsAsync(Project project, CancellationToken cancellationToken);
+
+            /// <summary>
+            /// Gets all the diagnostics to fix in the given project in a <see cref="FixAllContext"/>.
+            /// This includes both document-level diagnostics for all documents in the given project and project-level diagnostics, i.e. diagnostics with no source location, in the given project. 
+            /// </summary>
+            public abstract Task<IEnumerable<Diagnostic>> GetAllDiagnosticsAsync(Project project, CancellationToken cancellationToken);
+        }
+    }
+}

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

@@ -14,10 +14,9 @@ namespace Microsoft.CodeAnalysis.CodeFixes
     /// <summary>
     /// Context for "Fix all occurrences" code fixes provided by an <see cref="FixAllProvider"/>.
     /// </summary>
-    public class FixAllContext
+    public partial class FixAllContext
     {
-        private readonly Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> _getDocumentDiagnosticsAsync;
-        private readonly Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> _getProjectDiagnosticsAsync;
+        private readonly DiagnosticProvider _diagnosticProvider;
 
         /// <summary>
         /// Solution to fix all occurrences.
@@ -31,6 +30,7 @@ public class FixAllContext
 
         /// <summary>
         /// Document within which fix all occurrences was triggered.
+        /// Can be null if the context was created using <see cref="FixAllContext.FixAllContext(Project, CodeFixProvider, FixAllScope, string, IEnumerable{string}, DiagnosticProvider, CancellationToken)"/>.
         /// </summary>
         public Document Document { get; }
 
@@ -40,7 +40,7 @@ public class FixAllContext
         public CodeFixProvider CodeFixProvider { get; }
 
         /// <summary>
-        /// FixAllScope to fix all occurrences.
+        /// <see cref="FixAllScope"/> to fix all occurrences.
         /// </summary>
         public FixAllScope Scope { get; }
 
@@ -61,30 +61,62 @@ public class FixAllContext
         /// </summary>
         public CancellationToken CancellationToken { get; }
 
-        internal FixAllContext(
+        /// <summary>
+        /// Creates a new <see cref="FixAllContext"/>.
+        /// Use this overload when applying fix all to a diagnostic with a source location.
+        /// </summary>
+        /// <param name="document">Document within which fix all occurrences was triggered.</param>
+        /// <param name="codeFixProvider">Underlying <see cref="CodeFixes.CodeFixProvider"/> which triggered this fix all.</param>
+        /// <param name="scope"><see cref="FixAllScope"/> to fix all occurrences.</param>
+        /// <param name="codeActionEquivalenceKey">The <see cref="CodeAction.EquivalenceKey"/> value expected of a <see cref="CodeAction"/> participating in this fix all.</param>
+        /// <param name="diagnosticIds">Diagnostic Ids to fix.</param>
+        /// <param name="fixAllDiagnosticProvider">
+        /// <see cref="DiagnosticProvider"/> to fetch document/project diagnostics to fix in a <see cref="FixAllContext"/>.
+        /// </param>
+        /// <param name="cancellationToken">Cancellation token for fix all computation.</param>
+        public FixAllContext(
             Document document,
             CodeFixProvider codeFixProvider,
             FixAllScope scope,
-            string codeActionId,
+            string codeActionEquivalenceKey,
             IEnumerable<string> diagnosticIds,
-            Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getDocumentDiagnosticsAsync,
-            Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getProjectDiagnosticsAsync,
+            DiagnosticProvider fixAllDiagnosticProvider,
             CancellationToken cancellationToken)
-            : this(document, document.Project, codeFixProvider, scope, codeActionId, diagnosticIds, getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync, cancellationToken)
+            : this(document, document.Project, codeFixProvider, scope, codeActionEquivalenceKey, diagnosticIds, fixAllDiagnosticProvider, cancellationToken)
         {
+            if (document == null)
+            {
+                throw new ArgumentNullException(nameof(document));
+            }
         }
 
-        internal FixAllContext(
+        /// <summary>
+        /// Creates a new <see cref="FixAllContext"/>.
+        /// Use this overload when applying fix all to a diagnostic with no source location, i.e. <see cref="Location.None"/>.
+        /// </summary>
+        /// <param name="project">Project within which fix all occurrences was triggered.</param>
+        /// <param name="codeFixProvider">Underlying <see cref="CodeFixes.CodeFixProvider"/> which triggered this fix all.</param>
+        /// <param name="scope"><see cref="FixAllScope"/> to fix all occurrences.</param>
+        /// <param name="codeActionEquivalenceKey">The <see cref="CodeAction.EquivalenceKey"/> value expected of a <see cref="CodeAction"/> participating in this fix all.</param>
+        /// <param name="diagnosticIds">Diagnostic Ids to fix.</param>
+        /// <param name="fixAllDiagnosticProvider">
+        /// <see cref="DiagnosticProvider"/> to fetch document/project diagnostics to fix in a <see cref="FixAllContext"/>.
+        /// </param>
+        /// <param name="cancellationToken">Cancellation token for fix all computation.</param>
+        public FixAllContext(
             Project project,
             CodeFixProvider codeFixProvider,
             FixAllScope scope,
-            string codeActionId,
+            string codeActionEquivalenceKey,
             IEnumerable<string> diagnosticIds,
-            Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getDocumentDiagnosticsAsync,
-            Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getProjectDiagnosticsAsync,
+            DiagnosticProvider fixAllDiagnosticProvider,
             CancellationToken cancellationToken)
-            : this(null, project, codeFixProvider, scope, codeActionId, diagnosticIds, getDocumentDiagnosticsAsync, getProjectDiagnosticsAsync, cancellationToken)
+            : this(null, project, codeFixProvider, scope, codeActionEquivalenceKey, diagnosticIds, fixAllDiagnosticProvider, cancellationToken)
+        {
+            if (project == null)
             {
+                throw new ArgumentNullException(nameof(project));
+            }
         }
 
         private FixAllContext(
@@ -94,18 +126,38 @@ public class FixAllContext
             FixAllScope scope,
             string codeActionEquivalenceKey,
             IEnumerable<string> diagnosticIds,
-            Func<Document, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getDocumentDiagnosticsAsync,
-            Func<Project, bool, ImmutableHashSet<string>, CancellationToken, Task<IEnumerable<Diagnostic>>> getProjectDiagnosticsAsync,
+            DiagnosticProvider fixAllDiagnosticProvider,
             CancellationToken cancellationToken)
         {
+            Contract.ThrowIfNull(project);
+
+            if (codeFixProvider == null)
+            {
+                throw new ArgumentNullException(nameof(codeFixProvider));
+            }
+
+            if (diagnosticIds == null)
+            {
+                throw new ArgumentNullException(nameof(diagnosticIds));
+            }
+
+            if (diagnosticIds.Any(d => d == null))
+            {
+                throw new ArgumentException(WorkspacesResources.DiagnosticCannotBeNull, nameof(diagnosticIds));
+            }
+
+            if (fixAllDiagnosticProvider == null)
+            {
+                throw new ArgumentNullException(nameof(fixAllDiagnosticProvider));
+            }
+
             this.Document = document;
             this.Project = project;
             this.CodeFixProvider = codeFixProvider;
             this.Scope = scope;
             this.CodeActionEquivalenceKey = codeActionEquivalenceKey;
             this.DiagnosticIds = ImmutableHashSet.CreateRange(diagnosticIds);
-            _getDocumentDiagnosticsAsync = getDocumentDiagnosticsAsync;
-            _getProjectDiagnosticsAsync = getProjectDiagnosticsAsync;
+            _diagnosticProvider = fixAllDiagnosticProvider;
             this.CancellationToken = cancellationToken;
         }
 
@@ -114,16 +166,32 @@ public class FixAllContext
         /// </summary>
         public async Task<ImmutableArray<Diagnostic>> GetDocumentDiagnosticsAsync(Document document)
         {
-            Contract.ThrowIfNull(document);
+            if (document == null)
+            {
+                throw new ArgumentNullException(nameof(document));
+            }
 
             if (this.Project.Language != document.Project.Language)
             {
                 return ImmutableArray<Diagnostic>.Empty;
             }
 
-            var diagnostics = await _getDocumentDiagnosticsAsync(document, this.DiagnosticIds, this.CancellationToken).ConfigureAwait(false);
-            Contract.ThrowIfFalse(diagnostics.All(d => this.DiagnosticIds.Contains(d.Id)));
-            return diagnostics.ToImmutableArray();
+            var getDiagnosticsTask = _diagnosticProvider.GetDocumentDiagnosticsAsync(document, this.CancellationToken);
+            return await GetFilteredDiagnosticsAsync(getDiagnosticsTask, this.DiagnosticIds).ConfigureAwait(false);
+        }
+
+        private static async Task<ImmutableArray<Diagnostic>> GetFilteredDiagnosticsAsync(Task<IEnumerable<Diagnostic>> getDiagnosticsTask, ImmutableHashSet<string> diagnosticIds)
+        {
+            if (getDiagnosticsTask != null)
+            {
+                var diagnostics = await getDiagnosticsTask.ConfigureAwait(false);
+                if (diagnostics != null)
+                {
+                    return diagnostics.Where(d => d != null && diagnosticIds.Contains(d.Id)).ToImmutableArray();
+                }
+            }
+
+            return ImmutableArray<Diagnostic>.Empty;
         }
 
         /// <summary>
@@ -131,6 +199,11 @@ public async Task<ImmutableArray<Diagnostic>> GetDocumentDiagnosticsAsync(Docume
         /// </summary>
         public async Task<ImmutableArray<Diagnostic>> GetProjectDiagnosticsAsync(Project project)
         {
+            if (project == null)
+            {
+                throw new ArgumentNullException(nameof(project));
+            }
+
             return await GetProjectDiagnosticsAsync(project, includeAllDocumentDiagnostics: false).ConfigureAwait(false);
         }
 
@@ -140,6 +213,11 @@ public async Task<ImmutableArray<Diagnostic>> GetProjectDiagnosticsAsync(Project
         /// </summary>
         public async Task<ImmutableArray<Diagnostic>> GetAllDiagnosticsAsync(Project project)
         {
+            if (project == null)
+            {
+                throw new ArgumentNullException(nameof(project));
+            }
+
             return await GetProjectDiagnosticsAsync(project, includeAllDocumentDiagnostics: true).ConfigureAwait(false);
         }
 
@@ -157,9 +235,10 @@ private async Task<ImmutableArray<Diagnostic>> GetProjectDiagnosticsAsync(Projec
                 return ImmutableArray<Diagnostic>.Empty;
             }
 
-            var diagnostics = await _getProjectDiagnosticsAsync(project, includeAllDocumentDiagnostics, this.DiagnosticIds, this.CancellationToken).ConfigureAwait(false);
-            Contract.ThrowIfFalse(diagnostics.All(d => this.DiagnosticIds.Contains(d.Id)));
-            return diagnostics.ToImmutableArray();
+            var getDiagnosticsTask = includeAllDocumentDiagnostics ?
+                _diagnosticProvider.GetAllDiagnosticsAsync(project, CancellationToken) :
+                _diagnosticProvider.GetProjectDiagnosticsAsync(project, CancellationToken);
+            return await GetFilteredDiagnosticsAsync(getDiagnosticsTask, this.DiagnosticIds).ConfigureAwait(false);
         }
 
         /// <summary>
@@ -179,8 +258,7 @@ public FixAllContext WithCancellationToken(CancellationToken cancellationToken)
                 this.Scope,
                 this.CodeActionEquivalenceKey,
                 this.DiagnosticIds,
-                _getDocumentDiagnosticsAsync,
-                _getProjectDiagnosticsAsync,
+                _diagnosticProvider,
                 cancellationToken);
         }
     }

src/Workspaces/Core/Portable/PublicAPI.txt

@@ -72,7 +72,10 @@ Microsoft.CodeAnalysis.CodeFixes.FixAllContext.CancellationToken.get -> System.T
 Microsoft.CodeAnalysis.CodeFixes.FixAllContext.CodeActionEquivalenceKey.get -> string
 Microsoft.CodeAnalysis.CodeFixes.FixAllContext.CodeFixProvider.get -> Microsoft.CodeAnalysis.CodeFixes.CodeFixProvider
 Microsoft.CodeAnalysis.CodeFixes.FixAllContext.DiagnosticIds.get -> System.Collections.Immutable.ImmutableHashSet<string>
+Microsoft.CodeAnalysis.CodeFixes.FixAllContext.DiagnosticProvider
 Microsoft.CodeAnalysis.CodeFixes.FixAllContext.Document.get -> Microsoft.CodeAnalysis.Document
+Microsoft.CodeAnalysis.CodeFixes.FixAllContext.FixAllContext(Microsoft.CodeAnalysis.Document document, Microsoft.CodeAnalysis.CodeFixes.CodeFixProvider codeFixProvider, Microsoft.CodeAnalysis.CodeFixes.FixAllScope scope, string codeActionEquivalenceKey, System.Collections.Generic.IEnumerable<string> diagnosticIds, Microsoft.CodeAnalysis.CodeFixes.FixAllContext.DiagnosticProvider fixAllDiagnosticProvider, System.Threading.CancellationToken cancellationToken) -> void
+Microsoft.CodeAnalysis.CodeFixes.FixAllContext.FixAllContext(Microsoft.CodeAnalysis.Project project, Microsoft.CodeAnalysis.CodeFixes.CodeFixProvider codeFixProvider, Microsoft.CodeAnalysis.CodeFixes.FixAllScope scope, string codeActionEquivalenceKey, System.Collections.Generic.IEnumerable<string> diagnosticIds, Microsoft.CodeAnalysis.CodeFixes.FixAllContext.DiagnosticProvider fixAllDiagnosticProvider, System.Threading.CancellationToken cancellationToken) -> void
 Microsoft.CodeAnalysis.CodeFixes.FixAllContext.GetAllDiagnosticsAsync(Microsoft.CodeAnalysis.Project project) -> System.Threading.Tasks.Task<System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.Diagnostic>>
 Microsoft.CodeAnalysis.CodeFixes.FixAllContext.GetDocumentDiagnosticsAsync(Microsoft.CodeAnalysis.Document document) -> System.Threading.Tasks.Task<System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.Diagnostic>>
 Microsoft.CodeAnalysis.CodeFixes.FixAllContext.GetProjectDiagnosticsAsync(Microsoft.CodeAnalysis.Project project) -> System.Threading.Tasks.Task<System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.Diagnostic>>
@@ -834,6 +837,9 @@ abstract Microsoft.CodeAnalysis.CodeActions.CodeActionWithOptions.GetOptions(Sys
 abstract Microsoft.CodeAnalysis.CodeActions.PreviewOperation.GetPreviewAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<object>
 abstract Microsoft.CodeAnalysis.CodeFixes.CodeFixProvider.FixableDiagnosticIds.get -> System.Collections.Immutable.ImmutableArray<string>
 abstract Microsoft.CodeAnalysis.CodeFixes.CodeFixProvider.RegisterCodeFixesAsync(Microsoft.CodeAnalysis.CodeFixes.CodeFixContext context) -> System.Threading.Tasks.Task
+abstract Microsoft.CodeAnalysis.CodeFixes.FixAllContext.DiagnosticProvider.GetAllDiagnosticsAsync(Microsoft.CodeAnalysis.Project project, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<System.Collections.Generic.IEnumerable<Microsoft.CodeAnalysis.Diagnostic>>
+abstract Microsoft.CodeAnalysis.CodeFixes.FixAllContext.DiagnosticProvider.GetDocumentDiagnosticsAsync(Microsoft.CodeAnalysis.Document document, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<System.Collections.Generic.IEnumerable<Microsoft.CodeAnalysis.Diagnostic>>
+abstract Microsoft.CodeAnalysis.CodeFixes.FixAllContext.DiagnosticProvider.GetProjectDiagnosticsAsync(Microsoft.CodeAnalysis.Project project, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<System.Collections.Generic.IEnumerable<Microsoft.CodeAnalysis.Diagnostic>>
 abstract Microsoft.CodeAnalysis.CodeFixes.FixAllProvider.GetFixAsync(Microsoft.CodeAnalysis.CodeFixes.FixAllContext fixAllContext) -> System.Threading.Tasks.Task<Microsoft.CodeAnalysis.CodeActions.CodeAction>
 abstract Microsoft.CodeAnalysis.CodeRefactorings.CodeRefactoringProvider.ComputeRefactoringsAsync(Microsoft.CodeAnalysis.CodeRefactorings.CodeRefactoringContext context) -> System.Threading.Tasks.Task
 abstract Microsoft.CodeAnalysis.Differencing.TreeComparer<TNode>.GetChildren(TNode node) -> System.Collections.Generic.IEnumerable<TNode>

src/Workspaces/Core/Portable/Workspaces.csproj

@@ -303,6 +303,7 @@
     <Compile Include="CodeFixes\CodeFixContext.cs" />
     <Compile Include="CodeFixes\FixAllOccurrences\BatchSimplificationFixAllProvider.cs" />
     <Compile Include="CodeFixes\FixAllOccurrences\FixAllLogger.cs" />
+    <Compile Include="CodeFixes\FixAllOccurrences\FixAllContext.DiagnosticProvider.cs" />
     <Compile Include="CodeFixes\FixAllOccurrences\WellKnownFixAllProviders.cs" />
     <Compile Include="CodeFixes\FixAllOccurrences\BatchFixAllProvider.cs" />
     <Compile Include="CodeFixes\FixAllOccurrences\FixAllContext.cs" />