Stop making FixMultipleContext a subclass of FixAllContex.t

src/EditorFeatures/Core/Implementation/Suggestions/FixMultipleOccurrencesService.cs

@@ -75,7 +75,7 @@ public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
         }
 
         private FixMultipleSuggestedAction GetSuggestedAction(
-            FixMultipleContext fixMultipleContext,
+            FixAllContext fixAllContext,
             Diagnostic triggerDiagnostic,
             Workspace workspace,
             FixAllProvider fixAllProvider,
@@ -84,7 +84,7 @@ public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
             bool showPreviewChangesDialog,
             CancellationToken cancellationToken)
         {
-            var fixMultipleCodeAction = new FixMultipleCodeAction(fixMultipleContext, triggerDiagnostic, fixAllProvider, title, waitDialogMessage, showPreviewChangesDialog);
+            var fixMultipleCodeAction = new FixMultipleCodeAction(fixAllContext, triggerDiagnostic, fixAllProvider, title, waitDialogMessage, showPreviewChangesDialog);
             return new FixMultipleSuggestedAction(_listener, workspace, _editHandler, _waitIndicator, fixMultipleCodeAction, fixAllProvider);
         }
     }

src/Features/Core/Portable/CodeFixes/FixAllOccurrences/FixMultipleCodeAction.cs

@@ -12,13 +12,13 @@ internal partial class FixMultipleCodeAction : FixAllCodeAction
         private readonly string _computingFixWaitDialogMessage;
 
         internal FixMultipleCodeAction(
-            FixMultipleContext fixMultipleContext,
+            FixAllContext fixAllContext,
             Diagnostic triggerDiagnostic,
             FixAllProvider fixAllProvider,
             string title,
             string computingFixWaitDialogMessage,
             bool showPreviewChangesDialog)
-            : base(fixMultipleContext, fixAllProvider, showPreviewChangesDialog)
+            : base(fixAllContext, fixAllProvider, showPreviewChangesDialog)
         {
             _triggerDiagnostic = triggerDiagnostic;
             _title = title;

src/Features/Core/Portable/CodeFixes/Suppression/AbstractSuppressionCodeFixProvider.PragmaWarningBatchFixAllProvider.cs

@@ -38,7 +38,7 @@ public override async Task AddDocumentFixesAsync(Document document, ImmutableArr
                     var pragmaSuppression = pragmaSuppressions.SingleOrDefault();
                     if (pragmaSuppression != null)
                     {
-                        if (fixAllContext is FixMultipleContext)
+                        if (fixAllContext.IsFixMultiple)
                         {
                             pragmaSuppression = pragmaSuppression.CloneForFixMultipleContext();
                         }

src/Features/Core/Portable/CodeFixes/Suppression/AbstractSuppressionCodeFixProvider.RemoveSuppressionCodeAction.BatchFixer.cs

@@ -47,7 +47,7 @@ public override async Task AddDocumentFixesAsync(Document document, ImmutableArr
                             var codeAction = removeSuppressionFix.Action as RemoveSuppressionCodeAction;
                             if (codeAction != null)
                             {
-                                if (fixAllContext is FixMultipleContext)
+                                if (fixAllContext.IsFixMultiple)
                                 {
                                     codeAction = codeAction.CloneForFixMultipleContext();
                                 }
@@ -84,7 +84,7 @@ public async override Task AddProjectFixesAsync(Project project, ImmutableArray<
                         var removeSuppressionCodeAction = removeSuppressionFixes.SingleOrDefault()?.Action as RemoveSuppressionCodeAction;
                         if (removeSuppressionCodeAction != null)
                         {
-                            if (fixAllContext is FixMultipleContext)
+                            if (fixAllContext.IsFixMultiple)
                             {
                                 removeSuppressionCodeAction = removeSuppressionCodeAction.CloneForFixMultipleContext();
                             }

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

@@ -18,6 +18,8 @@ public partial class FixAllContext
         /// </summary>
         public abstract class DiagnosticProvider
         {
+            internal virtual bool IsFixMultiple => false;
+
             /// <summary>
             /// Gets all the diagnostics to fix in the given document in a <see cref="FixAllContext"/>.
             /// </summary>
@@ -33,16 +35,30 @@ public abstract class DiagnosticProvider
             /// 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);
+
+            internal virtual Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(
+                BatchFixAllProvider batchFixer, FixAllContext context)
+            {
+                return batchFixer.GetDocumentDiagnosticsToFixAsync(context);
+            }
+
+            internal virtual Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(
+                BatchFixAllProvider batchFixer, FixAllContext context)
+            {
+                return batchFixer.GetProjectDiagnosticsToFixAsync(context);
+            }
         }
 
-        internal virtual Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(BatchFixAllProvider batchFixer)
+        internal virtual Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(
+            BatchFixAllProvider batchFixer)
         {
-            return batchFixer.GetDocumentDiagnosticsToFixAsync(this);
+            return _diagnosticProvider.GetDocumentDiagnosticsToFixAsync(batchFixer, this);
         }
 
-        internal virtual Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(BatchFixAllProvider batchFixer)
+        internal virtual Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(
+            BatchFixAllProvider batchFixer)
         {
-            return batchFixer.GetProjectDiagnosticsToFixAsync(this);
+            return _diagnosticProvider.GetProjectDiagnosticsToFixAsync(batchFixer, this);
         }
     }
 }

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

@@ -7,7 +7,8 @@
 namespace Microsoft.CodeAnalysis.CodeFixes
 {
     public partial class FixAllContext
-    {        // Internal for testing purposes.
+    {
+        // Internal for testing purposes.
         internal class FixAllDiagnosticProvider : FixAllContext.DiagnosticProvider
         {
             private readonly ImmutableHashSet<string> _diagnosticIds;

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

@@ -161,6 +161,8 @@ public partial class FixAllContext
             this.CancellationToken = cancellationToken;
         }
 
+        internal bool IsFixMultiple => _diagnosticProvider.IsFixMultiple;
+
         internal DiagnosticProvider GetDiagnosticProvider() => _diagnosticProvider;
 
         /// <summary>
@@ -249,13 +251,6 @@ private async Task<ImmutableArray<Diagnostic>> GetProjectDiagnosticsAsync(Projec
         public FixAllContext WithCancellationToken(CancellationToken cancellationToken)
         {
             // TODO: We should change this API to be a virtual method, as the class is not sealed.
-            //       For now, special case FixMultipleContext.
-            var fixMultipleContext = this as FixMultipleContext;
-            if (fixMultipleContext != null)
-            {
-                return fixMultipleContext.WithCancellationToken(cancellationToken);
-            }
-
             if (this.CancellationToken == cancellationToken)
             {
                 return this;

src/Workspaces/Core/Portable/CodeFixes/FixAllOccurrences/FixMultipleContext.DiagnosticProvider.cs

@@ -12,13 +12,15 @@ namespace Microsoft.CodeAnalysis.CodeFixes
     /// <summary>
     /// Context for "Fix multiple occurrences" code fixes provided by an <see cref="FixAllProvider"/>.
     /// </summary>
-    internal partial class FixMultipleContext
+    internal static partial class FixMultipleContext
     {
         /// <summary>
         /// Diagnostic provider to fetch document/project diagnostics to fix in a <see cref="FixMultipleContext"/>.
         /// </summary>
-        public sealed class FixMultipleDiagnosticProvider : DiagnosticProvider
+        public sealed class FixMultipleDiagnosticProvider : FixAllContext.DiagnosticProvider
         {
+            internal override bool IsFixMultiple => true;
+
             private readonly ImmutableDictionary<Document, ImmutableArray<Diagnostic>> _documentDiagnosticsMap;
             private readonly ImmutableDictionary<Project, ImmutableArray<Diagnostic>> _projectDiagnosticsMap;
 
@@ -34,6 +36,16 @@ public FixMultipleDiagnosticProvider(ImmutableDictionary<Project, ImmutableArray
                 _documentDiagnosticsMap = ImmutableDictionary<Document, ImmutableArray<Diagnostic>>.Empty;
             }
 
+            internal override Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(BatchFixAllProvider batchFixer, FixAllContext context)
+            {
+                return Task.FromResult(_documentDiagnosticsMap);
+            }
+
+            internal override Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(BatchFixAllProvider batchFixer, FixAllContext context)
+            {
+                return Task.FromResult(_projectDiagnosticsMap);
+            }
+
             public ImmutableDictionary<Document, ImmutableArray<Diagnostic>> DocumentDiagnosticsToFix => _documentDiagnosticsMap;
             public ImmutableDictionary<Project, ImmutableArray<Diagnostic>> ProjectDiagnosticsToFix => _projectDiagnosticsMap;
 

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

@@ -14,10 +14,8 @@ namespace Microsoft.CodeAnalysis.CodeFixes
     /// <summary>
     /// Context for "Fix multiple occurrences" code fixes provided by an <see cref="FixAllProvider"/>.
     /// </summary>
-    internal partial class FixMultipleContext : FixAllContext
+    internal static partial class FixMultipleContext
     {
-        private readonly FixMultipleDiagnosticProvider _diagnosticProvider;
-
         /// <summary>
         /// Creates a new <see cref="FixMultipleContext"/>.
         /// Use this overload when applying fix multiple diagnostics with a source location.
@@ -26,7 +24,7 @@ internal partial class FixMultipleContext : FixAllContext
         /// <param name="codeFixProvider">Underlying <see cref="CodeFixes.CodeFixProvider"/> which triggered this fix all.</param>
         /// <param name="codeActionEquivalenceKey">The <see cref="CodeAction.EquivalenceKey"/> value expected of a <see cref="CodeAction"/> participating in this fix all.</param>
         /// <param name="cancellationToken">Cancellation token for fix all computation.</param>
-        public static FixMultipleContext Create(
+        public static FixAllContext Create(
             ImmutableDictionary<Document, ImmutableArray<Diagnostic>> diagnosticsToFix,
             CodeFixProvider codeFixProvider,
             string codeActionEquivalenceKey,
@@ -35,7 +33,8 @@ internal partial class FixMultipleContext : FixAllContext
             var triggerDocument = diagnosticsToFix.First().Key;
             var diagnosticIds = GetDiagnosticsIds(diagnosticsToFix.Values);
             var diagnosticProvider = new FixMultipleDiagnosticProvider(diagnosticsToFix);
-            return new FixMultipleContext(triggerDocument, codeFixProvider, codeActionEquivalenceKey, diagnosticIds, diagnosticProvider, cancellationToken);
+            return new FixAllContext(
+                triggerDocument, codeFixProvider, FixAllScope.Custom, codeActionEquivalenceKey, diagnosticIds, diagnosticProvider, cancellationToken);
         }
 
         /// <summary>
@@ -46,7 +45,7 @@ internal partial class FixMultipleContext : FixAllContext
         /// <param name="codeFixProvider">Underlying <see cref="CodeFixes.CodeFixProvider"/> which triggered this fix all.</param>
         /// <param name="codeActionEquivalenceKey">The <see cref="CodeAction.EquivalenceKey"/> value expected of a <see cref="CodeAction"/> participating in this fix all.</param>
         /// <param name="cancellationToken">Cancellation token for fix all computation.</param>
-        public static FixMultipleContext Create(
+        public static FixAllContext Create(
             ImmutableDictionary<Project, ImmutableArray<Diagnostic>> diagnosticsToFix,
             CodeFixProvider codeFixProvider,
             string codeActionEquivalenceKey,
@@ -55,31 +54,7 @@ internal partial class FixMultipleContext : FixAllContext
             var triggerProject = diagnosticsToFix.First().Key;
             var diagnosticIds = GetDiagnosticsIds(diagnosticsToFix.Values);
             var diagnosticProvider = new FixMultipleDiagnosticProvider(diagnosticsToFix);
-            return new FixMultipleContext(triggerProject, codeFixProvider, codeActionEquivalenceKey, diagnosticIds, diagnosticProvider, cancellationToken);
-        }
-
-        private FixMultipleContext(
-            Document triggerDocument,
-            CodeFixProvider codeFixProvider,
-            string codeActionEquivalenceKey,
-            ImmutableHashSet<string> diagnosticIds,
-            FixMultipleDiagnosticProvider diagnosticProvider,
-            CancellationToken cancellationToken)
-            : base(triggerDocument, codeFixProvider, FixAllScope.Custom, codeActionEquivalenceKey, diagnosticIds, diagnosticProvider, cancellationToken)
-        {
-            _diagnosticProvider = diagnosticProvider;
-        }
-
-        private FixMultipleContext(
-            Project triggerProject,
-            CodeFixProvider codeFixProvider,
-            string codeActionEquivalenceKey,
-            ImmutableHashSet<string> diagnosticIds,
-            FixMultipleDiagnosticProvider diagnosticProvider,
-            CancellationToken cancellationToken)
-            : base(triggerProject, codeFixProvider, FixAllScope.Custom, codeActionEquivalenceKey, diagnosticIds, diagnosticProvider, cancellationToken)
-        {
-            _diagnosticProvider = diagnosticProvider;
+            return new FixAllContext(triggerProject, codeFixProvider, FixAllScope.Custom, codeActionEquivalenceKey, diagnosticIds, diagnosticProvider, cancellationToken);
         }
 
         private static ImmutableHashSet<string> GetDiagnosticsIds(IEnumerable<ImmutableArray<Diagnostic>> diagnosticsCollection)
@@ -95,27 +70,5 @@ private static ImmutableHashSet<string> GetDiagnosticsIds(IEnumerable<ImmutableA
 
             return uniqueIds.ToImmutable();
         }
-
-        public new FixAllContext WithCancellationToken(CancellationToken cancellationToken)
-        {
-            if (this.CancellationToken == cancellationToken)
-            {
-                return this;
-            }
-
-            return Document != null ?
-                new FixMultipleContext(Document, CodeFixProvider, CodeActionEquivalenceKey, DiagnosticIds, _diagnosticProvider, cancellationToken) :
-                new FixMultipleContext(Project, CodeFixProvider, CodeActionEquivalenceKey, DiagnosticIds, _diagnosticProvider, cancellationToken);
-        }
-
-        internal override Task<ImmutableDictionary<Document, ImmutableArray<Diagnostic>>> GetDocumentDiagnosticsToFixAsync(BatchFixAllProvider batchFixer)
-        {
-            return Task.FromResult(_diagnosticProvider.DocumentDiagnosticsToFix);
-        }
-
-        internal override Task<ImmutableDictionary<Project, ImmutableArray<Diagnostic>>> GetProjectDiagnosticsToFixAsync(BatchFixAllProvider batchFixer)
-        {
-            return Task.FromResult(_diagnosticProvider.ProjectDiagnosticsToFix);
-        }
     }
 }