Provide a specialized fix-all provider as the stock BatchFixAllProvider can't...

Provide a specialized fix-all provider as the stock BatchFixAllProvider can't handle these close edits that this feature causes.

src/EditorFeatures/CSharpTest/CSharpEditorServicesTest.csproj

@@ -309,6 +309,7 @@
     <Compile Include="Organizing\OrganizeModifiersTests.cs" />
     <Compile Include="Organizing\OrganizeTypeDeclarationTests.cs" />
     <Compile Include="Organizing\OrganizeUsingsTests.cs" />
+    <Compile Include="SimplifyNullCheck\SimplifyNullCheckTests_FixAllTests.cs" />
     <Compile Include="SimplifyNullCheck\SimplifyNullCheckTests.cs" />
     <Compile Include="Structure\AbstractCSharpSyntaxNodeStructureTests.cs" />
     <Compile Include="Structure\AbstractCSharpSyntaxTriviaStructureTests.cs" />

src/EditorFeatures/CSharpTest/SimplifyNullCheck/SimplifyNullCheckTests.cs

@@ -65,6 +65,32 @@ void M(string s)
 }",
 @"using System;
 
+class C
+{
+    void M(string s)
+    {
+        _s = s ?? throw new ArgumentNullException(nameof(s));
+    }
+}");
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsSimplifyNullCheck)]
+        public async Task TestOnAssign()
+        {
+            await TestAsync(
+@"
+using System;
+
+class C
+{
+    void M(string s)
+    {
+        if (s == null) throw new ArgumentNullException(nameof(s));
+        _s = [|s|];
+    }
+}",
+@"using System;
+
 class C
 {
     void M(string s)

src/EditorFeatures/CSharpTest/SimplifyNullCheck/SimplifyNullCheckTests_FixAllTests.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.Threading.Tasks;
+using Microsoft.CodeAnalysis.Editor.CSharp.UnitTests.Diagnostics;
+using Microsoft.CodeAnalysis.SimplifyNullCheck;
+using Roslyn.Test.Utilities;
+using Xunit;
+
+namespace Microsoft.CodeAnalysis.Editor.CSharp.UnitTests.SimplifyNullCheck
+{
+    public partial class SimplifyNullCheckTests : AbstractCSharpDiagnosticProviderBasedUserDiagnosticTest
+    {
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsSimplifyNullCheck)]
+        public async Task FixAllInDocument1()
+        {
+            await TestAsync(
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        {|FixAllInDocument:if|} (s == null) { throw new ArgumentNullException(nameof(s)); }
+        if (t == null) { throw new ArgumentNullException(nameof(t)); }
+        _s = s;
+        _t = t;
+    }
+}",
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        _s = s ?? throw new ArgumentNullException(nameof(s));
+        _t = t ?? throw new ArgumentNullException(nameof(t));
+    }
+}", fixAllActionEquivalenceKey: SimplifyNullCheckCodeFixProvider.IfStatementEquivalenceKey);
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsSimplifyNullCheck)]
+        public async Task FixAllInDocument2()
+        {
+            await TestAsync(
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        if (s == null) { throw new ArgumentNullException(nameof(s)); }
+        {|FixAllInDocument:if|} (t == null) { throw new ArgumentNullException(nameof(t)); }
+        _s = s;
+        _t = t;
+    }
+}",
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        _s = s ?? throw new ArgumentNullException(nameof(s));
+        _t = t ?? throw new ArgumentNullException(nameof(t));
+    }
+}", fixAllActionEquivalenceKey: SimplifyNullCheckCodeFixProvider.IfStatementEquivalenceKey);
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsSimplifyNullCheck)]
+        public async Task FixAllInDocument3()
+        {
+            await TestAsync(
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        if (s == null) { throw new ArgumentNullException(nameof(s)); }
+        if (t == null) { throw new ArgumentNullException(nameof(t)); }
+        _s = {|FixAllInDocument:s|};
+        _t = t;
+    }
+}",
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        _s = s ?? throw new ArgumentNullException(nameof(s));
+        _t = t ?? throw new ArgumentNullException(nameof(t));
+    }
+}", fixAllActionEquivalenceKey: SimplifyNullCheckCodeFixProvider.IfStatementEquivalenceKey);
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsSimplifyNullCheck)]
+        public async Task FixAllInDocument4()
+        {
+            await TestAsync(
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        if (s == null) { throw new ArgumentNullException(nameof(s)); }
+        if (t == null) { throw new ArgumentNullException(nameof(t)); }
+        _s = s;
+        _t = {|FixAllInDocument:t|};
+    }
+}",
+@"
+using System;
+
+class C
+{
+    void M(string s, string t)
+    {
+        _s = s ?? throw new ArgumentNullException(nameof(s));
+        _t = t ?? throw new ArgumentNullException(nameof(t));
+    }
+}", fixAllActionEquivalenceKey: SimplifyNullCheckCodeFixProvider.IfStatementEquivalenceKey);
+        }
+    }
+}
\ No newline at end of file

src/Features/Core/Portable/Features.csproj

@@ -94,6 +94,7 @@
     <Compile Include="AddMissingReference\CodeAction.cs" />
     <Compile Include="SimplifyNullCheck\AbstractSimplifyNullCheckDiagnosticAnalyzer.cs" />
     <Compile Include="SimplifyNullCheck\SimplifyNullCheckCodeFixProvider.cs" />
+    <Compile Include="SimplifyNullCheck\SimplifyNullCheckCodeFixProvider.FixAllProvider.cs" />
     <Compile Include="Structure\BlockTypes.cs" />
     <Compile Include="CodeLens\CodeLensReferencesServiceFactory.cs" />
     <Compile Include="Structure\Syntax\BlockStructureExtensions.cs" />

src/Features/Core/Portable/SimplifyNullCheck/AbstractSimplifyNullCheckDiagnosticAnalyzer.cs

@@ -4,6 +4,7 @@
 using System.Collections.Immutable;
 using System.Reflection;
 using System.Threading;
+using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Semantics;
 
@@ -130,9 +131,23 @@ private void AnalyzeOperation(OperationAnalysisContext context)
                 assignmentExpression.Value.Syntax.GetLocation());
 
             context.ReportDiagnostic(
-                Diagnostic.Create(s_descriptor, ifOperation.Syntax.GetLocation(), additionalLocations: allLocations));
+                Diagnostic.Create(s_descriptor, ifOperation.Syntax.GetLocation(), additionalLocations: allLocations,
+                    properties: GetProperties(ifStatement: true)));
             context.ReportDiagnostic(
-                Diagnostic.Create(s_descriptor, expressionStatement.Syntax.GetLocation(), additionalLocations: allLocations));
+                Diagnostic.Create(s_descriptor, expressionStatement.Syntax.GetLocation(), additionalLocations: allLocations,
+                    properties: GetProperties(ifStatement: false)));
+        }
+
+        private static ImmutableDictionary<string, string> GetProperties(
+            bool ifStatement)
+        {
+            var equivalenceKey = ifStatement
+                ? SimplifyNullCheckCodeFixProvider.IfStatementEquivalenceKey
+                : SimplifyNullCheckCodeFixProvider.ExpressionStatementEquivalenceKey;
+
+            return ImmutableDictionary<string, string>.Empty.Add(
+                nameof(CodeAction.EquivalenceKey),
+                equivalenceKey);
         }
 
         private bool TryFindAssignmentExpression(

src/Features/Core/Portable/SimplifyNullCheck/SimplifyNullCheckCodeFixProvider.FixAllProvider.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.Immutable;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.CodeFixes;
+
+namespace Microsoft.CodeAnalysis.SimplifyNullCheck
+{
+    internal partial class SimplifyNullCheckCodeFixProvider
+    {
+        /// <summary>
+        /// A specialized <see cref="FixAllProvider"/> for this CodeFixProvider.
+        /// SimplifyNullCheck needs a specialized <see cref="FixAllProvider"/> because
+        /// the normal <see cref="BatchFixAllProvider"/> is insufficient for our needs.
+        /// Specifically, when doing a bulk fix-all, it's very common for many edits to
+        /// be near each other. The simplest example of this is just the following code:
+        /// 
+        /// <code>
+        ///     if (s == null) throw ...
+        ///     if (t == null) throw ...
+        ///     _s = s;
+        ///     _t = t;
+        /// </code>
+        /// 
+        /// If we use the normal batch-fixer then the underlying merge algorithm gets
+        /// throw off by hte sequence of edits.  specifically, the removal of 
+        /// "if (s == null) throw ..." actually gets seen by it as the removal of 
+        /// "(s == null) throw ... \r\n if".  That text change then intersects the
+        /// removal of "if (t == null) throw ...".  because of the intersection, one
+        /// of the edits is ignored.
+        /// 
+        /// This FixAllProvider avoids this entirely by not doing any textual merging.
+        /// Instead, we just take all the fixes to apply in the document, as we use
+        /// the core <see cref="SimplifyNullCheckCodeFixProvider.FixAllAsync"/> to do
+        /// all the editing at once on the SyntaxTree.  Because we're doing real tree
+        /// edits with actual nodes, there is no issue with anything getting messed up.
+        /// </summary>
+        private class SimplifyNullCheckFixAllProvider : BatchFixAllProvider
+        {
+            private readonly SimplifyNullCheckCodeFixProvider _provider;
+
+            public SimplifyNullCheckFixAllProvider(SimplifyNullCheckCodeFixProvider provider)
+            {
+                _provider = provider;
+            }
+
+            internal override async Task<CodeAction> GetFixAsync(
+                ImmutableDictionary<Document, ImmutableArray<Diagnostic>> documentsAndDiagnosticsToFixMap,
+                FixAllState fixAllState,
+                CancellationToken cancellationToken)
+            {
+                // Process all documents in parallel.
+                var updatedDocumentTasks = documentsAndDiagnosticsToFixMap.Select(
+                    kvp => FixAsync(fixAllState, kvp.Key, kvp.Value, cancellationToken));
+
+                await Task.WhenAll(updatedDocumentTasks).ConfigureAwait(false);
+
+                var currentSolution = fixAllState.Solution;
+                foreach (var task in updatedDocumentTasks)
+                {
+                    var updatedDocument = await task.ConfigureAwait(false);
+                    currentSolution = currentSolution.WithDocumentSyntaxRoot(
+                        updatedDocument.Id,
+                        await updatedDocument.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false));
+                }
+
+                var title = GetFixAllTitle(fixAllState);
+                return new CodeAction.SolutionChangeAction(title, _ => Task.FromResult(currentSolution));
+            }
+
+            private Task<Document> FixAsync(
+                FixAllState fixAllState, Document document, ImmutableArray<Diagnostic> diagnostics, CancellationToken cancellationToken)
+            {
+                var filteredDiagnostics = diagnostics.WhereAsArray(
+                    d => d.Properties[nameof(CodeAction.EquivalenceKey)] == fixAllState.CodeActionEquivalenceKey);
+
+                // Defer to the actual SimplifyNullCheckCodeFixProvider to process htis
+                // document.  It can process all the diagnostics and apply them properly.
+                return _provider.FixAllAsync(document, filteredDiagnostics, cancellationToken);
+            }
+        }
+    }
+}
\ No newline at end of file

src/Features/Core/Portable/SimplifyNullCheck/SimplifyNullCheckCodeFixProvider.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.Composition;
 using System.Linq;
@@ -16,42 +17,60 @@ namespace Microsoft.CodeAnalysis.SimplifyNullCheck
 {
     [ExportCodeFixProvider(LanguageNames.CSharp,
         Name = PredefinedCodeFixProviderNames.SimplifyNullCheck), Shared]
-    internal class SimplifyNullCheckCodeFixProvider : CodeFixProvider
+    internal partial class SimplifyNullCheckCodeFixProvider : CodeFixProvider
     {
+        public static readonly string IfStatementEquivalenceKey = nameof(IfStatementEquivalenceKey);
+        public static readonly string ExpressionStatementEquivalenceKey = nameof(ExpressionStatementEquivalenceKey);
+
         public override ImmutableArray<string> FixableDiagnosticIds
             => ImmutableArray.Create(IDEDiagnosticIds.SimplifyNullCheckDiagnosticId);
 
+        public override FixAllProvider GetFixAllProvider()
+            => new SimplifyNullCheckFixAllProvider(this);
+
         public override Task RegisterCodeFixesAsync(CodeFixContext context)
         {
             var diagnostic = context.Diagnostics.First();
+            var equivalenceKey = diagnostic.Properties[nameof(CodeAction.EquivalenceKey)];
             context.RegisterCodeFix(
-                new MyCodeAction(c => FixAsync(context.Document, diagnostic, c)),
+                new MyCodeAction(c => FixAsync(context.Document, diagnostic, c), equivalenceKey),
                 diagnostic);
 
             return SpecializedTasks.EmptyTask;
         }
 
-        private async Task<Document> FixAsync(
+        private Task<Document> FixAsync(
+            Document document,
+            Diagnostic diagnostic,
+            CancellationToken cancellationToken)
+        {
+            return FixAllAsync(document, ImmutableArray.Create(diagnostic), cancellationToken);
+        }
+
+        private async Task<Document> FixAllAsync(
             Document document, 
-            Diagnostic diagnostic, 
+            ImmutableArray<Diagnostic> diagnostics, 
             CancellationToken cancellationToken)
         {
             var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
 
-            var ifStatement = root.FindNode(diagnostic.AdditionalLocations[0].SourceSpan);
-            var throwStatementExpression = root.FindNode(diagnostic.AdditionalLocations[1].SourceSpan);
-            var assignmentValue = root.FindNode(diagnostic.AdditionalLocations[2].SourceSpan);
-
             var editor = new SyntaxEditor(root, document.Project.Solution.Workspace);
             var generator = editor.Generator;
 
-            // First, remote the if-statement entirely.
-            editor.RemoveNode(ifStatement);
+            foreach (var diagnostic in diagnostics)
+            {
+                var ifStatement = root.FindNode(diagnostic.AdditionalLocations[0].SourceSpan);
+                var throwStatementExpression = root.FindNode(diagnostic.AdditionalLocations[1].SourceSpan);
+                var assignmentValue = root.FindNode(diagnostic.AdditionalLocations[2].SourceSpan);
 
-            // Now, update the assignemnt value to go from 'a' to 'a ?? throw ...'.
-            editor.ReplaceNode(assignmentValue,
-                generator.CoalesceExpression(assignmentValue,
-                generator.ThrowExpression(throwStatementExpression)));
+                // First, remote the if-statement entirely.
+                editor.RemoveNode(ifStatement);
+
+                // Now, update the assignemnt value to go from 'a' to 'a ?? throw ...'.
+                editor.ReplaceNode(assignmentValue,
+                    generator.CoalesceExpression(assignmentValue,
+                    generator.ThrowExpression(throwStatementExpression)));
+            }
 
             var newRoot = editor.GetChangedRoot();
             return document.WithSyntaxRoot(newRoot);
@@ -60,10 +79,11 @@ public override Task RegisterCodeFixesAsync(CodeFixContext context)
         private class MyCodeAction : CodeAction.DocumentChangeAction
         {
             public MyCodeAction(
-                Func<CancellationToken, Task<Document>> createChangedDocument) 
-                : base(FeaturesResources.Simplify_null_check, createChangedDocument, FeaturesResources.Simplify_null_check)
+                Func<CancellationToken, Task<Document>> createChangedDocument,
+                string equivalenceKey) 
+                : base(FeaturesResources.Simplify_null_check, createChangedDocument, equivalenceKey)
             {
             }
         }
     }
-}
+}
\ No newline at end of file