Messy code that doesn't work :(

src/Analyzers/MSTest.Analyzers.CodeFixes/AssertionArgsShouldAvoidConditionalAccessFixer.cs

@@ -0,0 +1,229 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System.Collections.Immutable;
+using System.Composition;
+using System.Threading;
+
+using Analyzer.Utilities;
+
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.CodeFixes;
+using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+using Microsoft.CodeAnalysis.Editing;
+
+using MSTest.Analyzers.Helpers;
+
+namespace MSTest.Analyzers;
+
+[ExportCodeFixProvider(LanguageNames.CSharp, Name = nameof(AssertionArgsShouldAvoidConditionalAccessFixer))]
+[Shared]
+public sealed class AssertionArgsShouldAvoidConditionalAccessFixer : CodeFixProvider
+{
+    /// <summary>
+    /// The scenario that is complicating this code fix is if we have multiple diagnostics that are doing conditional access
+    /// on the same expression. In that case, we need to ensure that we don't add multiple Assert.IsNotNull calls.
+    /// The first idea was to iterate through the existing statements, and if we found Assert.IsNotNull with
+    /// the relevant expression, we don't add it again. However, this approach works for iterative codefix application
+    /// only, and doesn't work with the BatchFixAllProvider. The BatchFixAllProvider works by applying individual fixes
+    /// completely in isolation, then merging the text changes.
+    /// This means, every invocation of the code action will not see that Assert.IsNotNull was added by another.
+    /// So, we provide our own FixAllProvider.
+    /// This FixAllProvider will reuse the same DocumentEditor across all the code actions.
+    /// </summary>
+    private sealed class CustomFixAll : DocumentBasedFixAllProvider
+    {
+        protected override async Task<Document?> FixAllAsync(FixAllContext fixAllContext, Document document, ImmutableArray<Diagnostic> diagnostics)
+        {
+            SyntaxNode root = await document.GetRequiredSyntaxRootAsync(fixAllContext.CancellationToken).ConfigureAwait(false);
+            DocumentEditor editor = await DocumentEditor.CreateAsync(document, fixAllContext.CancellationToken);
+            Document currentDocument = document;
+            foreach (Diagnostic diagnostic in diagnostics)
+            {
+                SyntaxNode assertInvocation = root.FindNode(diagnostic.Location.SourceSpan, getInnermostNodeForTie: true);
+                editor.TrackNode(assertInvocation);
+            }
+
+            foreach (Diagnostic diagnostic in diagnostics)
+            {
+                SyntaxNode conditionalAccess = root.FindNode(diagnostic.AdditionalLocations[0].SourceSpan, getInnermostNodeForTie: true);
+                SyntaxNode assertInvocation = root.FindNode(diagnostic.Location.SourceSpan, getInnermostNodeForTie: true);
+                if (conditionalAccess is not ConditionalAccessExpressionSyntax conditionalAccessExpressionSyntax ||
+                    assertInvocation is not InvocationExpressionSyntax invocationExpressionSyntax)
+                {
+                    continue;
+                }
+
+                new SingleFixCodeAction(currentDocument, conditionalAccessExpressionSyntax, invocationExpressionSyntax).ApplyFix(editor);
+                currentDocument = editor.GetChangedDocument();
+            }
+
+            return editor.GetChangedDocument();
+        }
+    }
+
+    public sealed override ImmutableArray<string> FixableDiagnosticIds { get; }
+        = ImmutableArray.Create(DiagnosticIds.AssertionArgsShouldAvoidConditionalAccessRuleId);
+
+    public override FixAllProvider GetFixAllProvider()
+        // See https://github.com/dotnet/roslyn/blob/main/docs/analyzers/FixAllProvider.md for more information on Fix All Providers
+        => new CustomFixAll();
+
+    public override async Task RegisterCodeFixesAsync(CodeFixContext context)
+    {
+        SyntaxNode root = await context.Document.GetRequiredSyntaxRootAsync(context.CancellationToken).ConfigureAwait(false);
+        Diagnostic diagnostic = context.Diagnostics[0];
+
+        SyntaxNode conditionalAccess = root.FindNode(diagnostic.AdditionalLocations[0].SourceSpan, getInnermostNodeForTie: true);
+        SyntaxNode assertInvocation = root.FindNode(diagnostic.Location.SourceSpan, getInnermostNodeForTie: true);
+        if (conditionalAccess is not ConditionalAccessExpressionSyntax conditionalAccessExpressionSyntax ||
+            assertInvocation is not InvocationExpressionSyntax invocationExpressionSyntax)
+        {
+            return;
+        }
+
+        context.RegisterCodeFix(
+            new SingleFixCodeAction(context.Document, conditionalAccessExpressionSyntax, invocationExpressionSyntax),
+            diagnostic);
+    }
+
+    private sealed class SingleFixCodeAction : CodeAction
+    {
+        private readonly Document _document;
+        private readonly ConditionalAccessExpressionSyntax _conditionalAccessExpressionSyntax;
+        private readonly InvocationExpressionSyntax _invocationExpressionSyntax;
+
+        public SingleFixCodeAction(Document document, ConditionalAccessExpressionSyntax conditionalAccessExpressionSyntax, InvocationExpressionSyntax invocationExpressionSyntax)
+        {
+            _document = document;
+            _conditionalAccessExpressionSyntax = conditionalAccessExpressionSyntax;
+            _invocationExpressionSyntax = invocationExpressionSyntax;
+        }
+
+        public override string Title { get; } = "TODO" /*CodeFixResources.AssertionArgsShouldAvoidConditionalAccessFix*/;
+
+        public override string? EquivalenceKey => nameof(AssertionArgsShouldAvoidConditionalAccessFixer);
+
+        protected override async Task<Document> GetChangedDocumentAsync(CancellationToken cancellationToken)
+        {
+            DocumentEditor editor = await DocumentEditor.CreateAsync(_document, cancellationToken).ConfigureAwait(false);
+            ApplyFix(editor);
+            return editor.GetChangedDocument();
+        }
+
+        internal void ApplyFix(DocumentEditor editor)
+        {
+            ExpressionSyntax expressionCheckedForNull = _conditionalAccessExpressionSyntax.Expression;
+            bool isNullAssertAlreadyPresent = IsNullAssertAlreadyPresent(expressionCheckedForNull, editor.GetChangedRoot().GetCurrentNode(_invocationExpressionSyntax) ?? _invocationExpressionSyntax);
+
+            // Easier than correctly reconstructing the syntax node manually, but not ideal.
+            ExpressionSyntax parsedExpression = SyntaxFactory.ParseExpression($"{expressionCheckedForNull.ToFullString()}{_conditionalAccessExpressionSyntax.WhenNotNull}");
+
+            editor.ReplaceNode(_conditionalAccessExpressionSyntax, parsedExpression);
+
+            if (!isNullAssertAlreadyPresent)
+            {
+                ExpressionStatementSyntax assertIsNotNull = SyntaxFactory.ExpressionStatement(
+                    SyntaxFactory.InvocationExpression(
+                        SyntaxFactory.MemberAccessExpression(
+                            SyntaxKind.SimpleMemberAccessExpression,
+                            SyntaxFactory.IdentifierName("Assert"),
+                            SyntaxFactory.IdentifierName("IsNotNull")))
+                    .WithArgumentList(
+                        SyntaxFactory.ArgumentList(
+                            SyntaxFactory.SingletonSeparatedList(
+                                SyntaxFactory.Argument(expressionCheckedForNull)))));
+                if (_invocationExpressionSyntax.Parent is ExpressionStatementSyntax expressionStatement)
+                {
+                    editor.InsertBefore(expressionStatement, assertIsNotNull);
+                }
+                else if (_invocationExpressionSyntax.Parent is ArrowExpressionClauseSyntax arrowExpressionClauseSyntax)
+                {
+                    // The following types are where ArrowExpressionClause can appear.
+                    // BaseMethodDeclarationSyntax: ConstructorDeclarationSyntax, ConversionOperatorDeclarationSyntax, DestructorDeclarationSyntax, MethodDeclarationSyntax, OperatorDeclarationSyntax
+                    // AccessorDeclarationSyntax, IndexerDeclarationSyntax, PropertyDeclarationSyntax, LocalFunctionStatementSyntax
+                    //
+                    // PropertyDeclarationSyntax and IndexerDeclarationSyntax don't make sense so we won't handle it.
+                    if (arrowExpressionClauseSyntax.Parent is BaseMethodDeclarationSyntax parentBaseMethod)
+                    {
+                        editor.ReplaceNode(
+                            parentBaseMethod,
+                            (node, _) =>
+                            {
+                                var parentBaseMethod = (BaseMethodDeclarationSyntax)node;
+                                return parentBaseMethod
+                                    .WithExpressionBody(null)
+                                    .WithBody(SyntaxFactory.Block(
+                                        assertIsNotNull,
+                                        SyntaxFactory.ExpressionStatement(parentBaseMethod.ExpressionBody!.Expression)));
+                            });
+                    }
+                    else if (arrowExpressionClauseSyntax.Parent is AccessorDeclarationSyntax parentAccessor)
+                    {
+                        editor.ReplaceNode(
+                            parentAccessor,
+                            (node, _) =>
+                            {
+                                var parentAccessor = (AccessorDeclarationSyntax)node;
+                                return parentAccessor
+                                    .WithExpressionBody(null)
+                                    .WithBody(SyntaxFactory.Block(
+                                        assertIsNotNull,
+                                        SyntaxFactory.ExpressionStatement(parentAccessor.ExpressionBody!.Expression)));
+                            });
+                    }
+                    else if (arrowExpressionClauseSyntax.Parent is LocalFunctionStatementSyntax parentLocalFunction)
+                    {
+                        editor.ReplaceNode(
+                            parentLocalFunction,
+                            (node, _) =>
+                            {
+                                var parentLocalFunction = (LocalFunctionStatementSyntax)node;
+                                return parentLocalFunction
+                                    .WithExpressionBody(null)
+                                    .WithBody(SyntaxFactory.Block(
+                                        assertIsNotNull,
+                                        SyntaxFactory.ExpressionStatement(parentLocalFunction.ExpressionBody!.Expression)));
+                            });
+                    }
+                }
+            }
+        }
+    }
+
+    private static bool IsNullAssertAlreadyPresent(SyntaxNode expressionCheckedForNull, InvocationExpressionSyntax invocationExpressionSyntax)
+    {
+        if (invocationExpressionSyntax.Parent?.Parent is not BlockSyntax blockSyntax)
+        {
+            return false;
+        }
+
+        foreach (StatementSyntax statement in blockSyntax.Statements)
+        {
+            if (statement is ExpressionStatementSyntax expressionStatement)
+            {
+                // We expect Assert.IsNull to be present before the invocation expression in question.
+                if (expressionStatement.Expression == invocationExpressionSyntax)
+                {
+                    return false;
+                }
+
+                if (expressionStatement.Expression is InvocationExpressionSyntax invocation)
+                {
+                    SimpleNameSyntax? methodName =
+                        invocation.Expression as IdentifierNameSyntax ?? (invocation.Expression as MemberAccessExpressionSyntax)?.Name;
+                    if ((methodName?.Identifier.Value as string) == "IsNotNull" &&
+                        invocation.ArgumentList.Arguments.Count > 0 &&
+                        invocation.ArgumentList.Arguments[0].Expression.IsEquivalentTo(expressionCheckedForNull))
+                    {
+                        return true;
+                    }
+                }
+            }
+        }
+
+        return false;
+    }
+}

src/Analyzers/MSTest.Analyzers/AssertionArgsShouldAvoidConditionalAccessAnalyzer.cs

@@ -101,15 +101,15 @@ private static void AnalyzeOperation(OperationAnalysisContext context, INamedTyp
         (_, int argumentCountToCheck) = supportedMethodNames.FirstOrDefault(x => x.MethodName == invocationOperation.TargetMethod.Name);
         if (argumentCountToCheck == 0
             || !SymbolEqualityComparer.Default.Equals(assertSymbol, invocationOperation.TargetMethod.ContainingType)
-            || !HasAnyConditionalAccessOperationChild(invocationOperation, argumentCountToCheck))
+            || !HasAnyConditionalAccessOperationChild(invocationOperation, argumentCountToCheck, out Location? additionalLocation))
         {
             return;
         }
 
-        context.ReportDiagnostic(invocationOperation.CreateDiagnostic(Rule));
+        context.ReportDiagnostic(invocationOperation.CreateDiagnostic(Rule, additionalLocations: ImmutableArray.Create(additionalLocation), properties: null));
     }
 
-    private static bool HasAnyConditionalAccessOperationChild(IInvocationOperation invocationOperation, int argumentCountToCheck)
+    private static bool HasAnyConditionalAccessOperationChild(IInvocationOperation invocationOperation, int argumentCountToCheck, [NotNullWhen(true)] out Location? additionalLocation)
     {
         foreach (IArgumentOperation argument in invocationOperation.Arguments)
         {
@@ -127,19 +127,27 @@ private static bool HasAnyConditionalAccessOperationChild(IInvocationOperation i
             //      a.b?.c
             if (value.Kind == OperationKind.ConditionalAccess)
             {
+                additionalLocation = value.Syntax.GetLocation();
                 return true;
             }
 
             // Check for binary operations with conditional access => s?.Length > 1.
             if (value is IBinaryOperation binaryOperation)
             {
-                if (binaryOperation.LeftOperand.Kind == OperationKind.ConditionalAccess || binaryOperation.RightOperand.Kind == OperationKind.ConditionalAccess)
+                if (binaryOperation.LeftOperand.Kind == OperationKind.ConditionalAccess)
                 {
+                    additionalLocation = binaryOperation.LeftOperand.Syntax.GetLocation();
+                    return true;
+                }
+                else if (binaryOperation.RightOperand.Kind == OperationKind.ConditionalAccess)
+                {
+                    additionalLocation = binaryOperation.RightOperand.Syntax.GetLocation();
                     return true;
                 }
             }
         }
 
+        additionalLocation = null;
         return false;
     }
 }

test/UnitTests/MSTest.Analyzers.UnitTests/AssertionArgsShouldAvoidConditionalAccessAnalyzerTests.cs

@@ -3,7 +3,7 @@
 
 using VerifyCS = MSTest.Analyzers.Test.CSharpCodeFixVerifier<
     MSTest.Analyzers.AssertionArgsShouldAvoidConditionalAccessAnalyzer,
-    Microsoft.CodeAnalysis.Testing.EmptyCodeFixProvider>;
+    MSTest.Analyzers.AssertionArgsShouldAvoidConditionalAccessFixer>;
 
 namespace MSTest.Analyzers.UnitTests;
 
@@ -93,11 +93,110 @@ public void Compliant()
                     Assert.AreNotEqual(s.Length, s.Length);
                     Assert.AreNotEqual(a.S.Length, 32);
                     Assert.AreNotEqual(b.S.Length, 32);
+                }
             }
+            """;
+
+        string fixedCode = """
+            #nullable enable
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+            using System.Collections.Generic;
+
+            public class A
+            {
+                public string? S { get; set; }
+                public List<string>? T { get; set; }
+            }
+
+            [TestClass]
+            public class MyTestClass
+            {
+                [TestMethod]
+                public void NonCompliant()
+                {
+                    A? a = new A();
+                    A b = new A();
+                    string? s = "";
+                    Assert.IsNotNull(s);
+                    Assert.AreEqual(s.Length, 32);
+                    Assert.AreEqual(((s.Length)), 32);
+                    Assert.AreEqual(s.Length, s.Length);
+                    Assert.IsNotNull(a);
+                    Assert.IsNotNull(a.S);
+                    Assert.AreEqual(a.S.Length, 32);
+                    Assert.IsNotNull(b.S);
+                    Assert.AreEqual(b.S.Length, 32);
+                    Assert.IsNotNull(a.T);
+                    Assert.IsNotNull(a.T[3]);
+                    Assert.AreEqual(a.T[3].Length, 32);
+                    Assert.IsNotNull(b.T);
+                    Assert.IsNotNull(b.T[3]);
+                    Assert.AreEqual(b.T[3].Length, 32);
+
+                    Assert.AreNotEqual(s.Length, 32);
+                    Assert.AreNotEqual(((s.Length)), 32);
+                    Assert.AreNotEqual(s.Length, s.Length);
+                    Assert.AreNotEqual(a.S.Length, 32);
+                    Assert.AreNotEqual(b.S.Length, 32);
+                    Assert.AreNotEqual(a.T[3].Length, 32);
+                    Assert.AreNotEqual(b.T[3].Length, 32);
+
+                    Assert.AreSame(s.Length, 32);
+                    Assert.AreSame(((s.Length)), 32);
+                    Assert.AreSame(s.Length, s.Length);
+                    Assert.AreSame(a.S.Length, 32);
+                    Assert.AreSame(b.S.Length, 32);
+                    Assert.AreSame(a.T[3].Length, 32);
+                    Assert.AreSame(b.T[3].Length, 32);
+
+                    Assert.AreNotSame(s.Length, 32);
+                    Assert.AreNotSame(((s.Length)), 32);
+                    Assert.AreNotSame(s.Length, s.Length);
+                    Assert.AreNotSame(a.S.Length, 32);
+                    Assert.AreNotSame(b.S.Length, 32);
+                    Assert.AreNotSame(a.T[3].Length, 32);
+                    Assert.AreNotSame(b.T[3].Length, 32);
+                }
+
+                [TestMethod]
+                public void Compliant()
+                {
+                    string? s = "";
+                    A? a = new A();
+                    A b = new A();
+
+                    Assert.IsNotNull(s);
+                    Assert.IsNotNull(a);
+                    Assert.IsNotNull(a.S);
+                    Assert.IsNotNull(a.T);
+                    Assert.IsNotNull(b);
+                    Assert.IsNotNull(b.S);
+                    Assert.IsNotNull(b.T);
+                    Assert.AreEqual(s.Length, 32);
+                    Assert.AreEqual(((s.Length)), 32);
+                    Assert.AreEqual(s.Length, s.Length);
+                    Assert.AreEqual(a.S.Length, 32);
+                    Assert.AreEqual(b.S.Length, 32);
+                    Assert.AreEqual(a.T[3].Length, 32);
+                    Assert.AreEqual(b.T[3].Length, 32);
+                    Assert.AreNotEqual(s.Length, 32);
+                    Assert.AreNotEqual(((s.Length)), 32);
+                    Assert.AreNotEqual(s.Length, s.Length);
+                    Assert.AreNotEqual(a.S.Length, 32);
+                    Assert.AreNotEqual(b.S.Length, 32);
+                }
             }
             """;
 
-        await VerifyCS.VerifyAnalyzerAsync(code);
+        await new VerifyCS.Test
+        {
+            TestCode = code,
+            FixedCode = fixedCode,
+            NumberOfFixAllIterations = 48,
+            NumberOfFixAllInDocumentIterations = 48,
+            NumberOfFixAllInProjectIterations = 48,
+            NumberOfIncrementalIterations = 48,
+        }.RunAsync();
     }
 
     [TestMethod]