Merge pull request #21503 from zaytsev-victor/Fixed21501

Fixed 'Add null check' refactoring

src/EditorFeatures/CSharpTest/InitializeParameter/AddParameterCheckTests.cs

@@ -886,6 +886,158 @@ public C(string s)
             throw new ArgumentNullException(nameof(s));
         }
     }
+}");
+        }
+
+        [WorkItem(21501, "https://github.com/dotnet/roslyn/issues/21501")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsInitializeParameter)]
+        public async Task TestInArrowExpression1()
+        {
+            await TestInRegularAndScript1Async(
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public int Foo(int[] array[||]) =>
+        array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count();
+}",
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public int Foo(int[] array)
+    {
+        if (array == null)
+        {
+            throw new ArgumentNullException(nameof(array));
+        }
+
+        return array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count();
+    }
+}");
+        }
+
+        [WorkItem(21501, "https://github.com/dotnet/roslyn/issues/21501")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsInitializeParameter)]
+        public async Task TestInArrowExpression2()
+        {
+            await TestInRegularAndScript1Async(
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public int Foo(int[] array[||]) /* Bar */ => /* Bar */
+        array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count(); /* Bar */
+}",
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public int Foo(int[] array) /* Bar */
+    {
+        if (array == null)
+        {
+            throw new ArgumentNullException(nameof(array));
+        }
+        /* Bar */
+        return array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count(); /* Bar */
+    }
+}");
+        }
+
+        [WorkItem(21501, "https://github.com/dotnet/roslyn/issues/21501")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsInitializeParameter)]
+        public async Task TestMissingInArrowExpression1()
+        {
+            await TestMissingInRegularAndScriptAsync(
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public void Foo(string bar[||]) =>
+#if DEBUG
+        Console.WriteLine(""debug"" + bar);
+#else
+        Console.WriteLine(""release"" + bar);
+#endif
+}");
+        }
+
+        [WorkItem(21501, "https://github.com/dotnet/roslyn/issues/21501")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsInitializeParameter)]
+        public async Task TestMissingInArrowExpression2()
+        {
+            await TestMissingInRegularAndScriptAsync(
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public int Foo(int[] array[||]) =>
+#if DEBUG
+        array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count();
+#else
+        array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count();
+#endif
+}");
+        }
+
+        [WorkItem(21501, "https://github.com/dotnet/roslyn/issues/21501")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsInitializeParameter)]
+        public async Task TestInArrowExpression3()
+        {
+            await TestInRegularAndScript1Async(
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public void Foo(int[] array[||]) =>
+        array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count();
+}",
+@"
+using System;
+using System.Linq;
+
+class C
+{
+    public void Foo(int[] array)
+    {
+        if (array == null)
+        {
+            throw new ArgumentNullException(nameof(array));
+        }
+
+        array.Where(x => x > 3)
+            .OrderBy(x => x)
+            .Count();
+    }
 }");
         }
     }

src/Features/CSharp/Portable/InitializeParameter/CSharpAddParameterCheckCodeRefactoringProvider.cs

@@ -2,6 +2,7 @@
 
 using System.Composition;
 using Microsoft.CodeAnalysis.CodeRefactorings;
+using Microsoft.CodeAnalysis.CSharp.Extensions;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Editing;
 using Microsoft.CodeAnalysis.InitializeParameter;
@@ -30,5 +31,18 @@ protected override void InsertStatement(SyntaxEditor editor, BaseMethodDeclarati
 
         protected override bool IsImplicitConversion(Compilation compilation, ITypeSymbol source, ITypeSymbol destination)
             => InitializeParameterHelpers.IsImplicitConversion(compilation, source, destination);
+
+        protected override bool CanOffer(SyntaxNode body)
+        {
+            if (body is ArrowExpressionClauseSyntax arrowExpressionClauseSyntax)
+            {
+                return arrowExpressionClauseSyntax.TryConvertToStatement(
+                    semicolonToken: SyntaxFactory.Token(SyntaxKind.SemicolonToken), 
+                    createReturnStatementForExpression: false, 
+                    statement: out var _);              
+            }
+
+            return true;
+        }
     }
 }

src/Features/CSharp/Portable/InitializeParameter/InitializeParameterHelpers.cs

@@ -28,28 +28,23 @@ public static SyntaxNode TryGetLastStatement(IBlockOperation blockStatementOpt)
             SyntaxNode statementToAddAfterOpt,
             StatementSyntax statement)
         {
-            var generator = editor.Generator;
-
             if (methodDeclaration.ExpressionBody != null)
             {
                 // If this is a => method, then we'll have to convert the method to have a block
                 // body.  Add the new statement as the first/last statement of the new block 
                 // depending if we were asked to go after something or not.
+                if (!methodDeclaration.ExpressionBody.TryConvertToStatement(methodDeclaration.SemicolonToken, CreateReturnStatement(methodDeclaration), out var declaration))
+                {
+                    return;
+                }
+
                 if (statementToAddAfterOpt == null)
                 {
-                    editor.SetStatements(
-                        methodDeclaration,
-                        ImmutableArray.Create(
-                            statement,
-                            generator.ExpressionStatement(methodDeclaration.ExpressionBody.Expression)));
+                    editor.SetStatements(methodDeclaration, ImmutableArray.Create(statement, declaration));
                 }
                 else
                 {
-                    editor.SetStatements(
-                        methodDeclaration,
-                        ImmutableArray.Create(
-                            generator.ExpressionStatement(methodDeclaration.ExpressionBody.Expression),
-                            statement));
+                    editor.SetStatements(methodDeclaration, ImmutableArray.Create(declaration, statement));
                 }
             }
             else if (methodDeclaration.Body != null)
@@ -84,5 +79,22 @@ public static SyntaxNode TryGetLastStatement(IBlockOperation blockStatementOpt)
                                      .WithBody(SyntaxFactory.Block(statement)));
             }
         }
+
+        private static bool CreateReturnStatement(BaseMethodDeclarationSyntax declarationSyntax)
+        {
+            switch(declarationSyntax)
+            {
+                case MethodDeclarationSyntax method:
+                    return !method.ReturnType.IsVoid();
+                case ConversionOperatorDeclarationSyntax _:
+                case OperatorDeclarationSyntax _:
+                    return true;
+                case DestructorDeclarationSyntax _:
+                case ConstructorDeclarationSyntax _:
+                    return false;
+            }
+
+            return false;
+        }
     }
 }

src/Features/Core/Portable/InitializeParameter/AbstractAddParameterCheckCodeRefactoringProvider.cs

@@ -55,6 +55,11 @@ internal abstract partial class AbstractAddParameterCheckCodeRefactoringProvider
             // people do strange things in their constructors.
             if (blockStatementOpt != null)
             {
+                if (!CanOffer(blockStatementOpt.Syntax))
+                { 
+                    return ImmutableArray<CodeAction>.Empty;
+                }
+
                 foreach (var statement in blockStatementOpt.Operations)
                 {
                     if (IsIfNullCheck(statement, parameter))
@@ -93,6 +98,8 @@ internal abstract partial class AbstractAddParameterCheckCodeRefactoringProvider
             return result.ToImmutableAndFree();
         }
 
+        protected abstract bool CanOffer(SyntaxNode body);
+
         private bool ContainsNullCoalesceCheck(
             ISyntaxFactsService syntaxFacts, SemanticModel semanticModel,
             IOperation statement, IParameterSymbol parameter,

src/Features/VisualBasic/Portable/InitializeParameter/VisualBasicAddParameterCheckCodeRefactoringProvider.vb

@@ -33,5 +33,9 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.InitializeParameter
         Protected Overrides Sub InsertStatement(editor As SyntaxEditor, methodDeclaration As MethodBlockBaseSyntax, statementToAddAfterOpt As SyntaxNode, statement As StatementSyntax)
             InitializeParameterHelpers.InsertStatement(editor, methodDeclaration, statementToAddAfterOpt, statement)
         End Sub
+
+        Protected Overrides Function CanOffer(body As SyntaxNode) As Boolean
+            Return True
+        End Function
     End Class
 End Namespace

src/Workspaces/CSharp/Portable/Extensions/ArrowExpressionClauseSyntaxExtensions.cs

@@ -13,23 +13,38 @@ internal static class ArrowExpressionClauseSyntaxExtensions
             SyntaxToken semicolonToken,
             bool createReturnStatementForExpression,
             out BlockSyntax block)
+        {
+            if (!arrowExpression.TryConvertToStatement(semicolonToken, createReturnStatementForExpression, out var statement))
+            {
+                block = null;
+                return false;
+            }
+
+            block = SyntaxFactory.Block(statement);
+            return true;
+        }
+
+        public static bool TryConvertToStatement(
+            this ArrowExpressionClauseSyntax arrowExpression,
+            SyntaxToken semicolonToken,
+            bool createReturnStatementForExpression,
+            out StatementSyntax statement)
         {
             // It's tricky to convert an arrow expression with directives over to a block.
             // We'd need to find and remove the directives *after* the arrow expression and
             // move them accordingly.  So, for now, we just disallow this.
             if (arrowExpression.Expression.GetLeadingTrivia().Any(t => t.IsDirective))
             {
-                block = null;
+                statement = null;
                 return false;
             }
 
-            var statement = ConvertToStatement(arrowExpression.Expression, semicolonToken, createReturnStatementForExpression);
+            statement = ConvertToStatement(arrowExpression.Expression, semicolonToken, createReturnStatementForExpression);
             if (arrowExpression.ArrowToken.TrailingTrivia.Any(t => t.IsSingleOrMultiLineComment()))
             {
                 statement = statement.WithPrependedLeadingTrivia(arrowExpression.ArrowToken.TrailingTrivia);
             }
 
-            block = SyntaxFactory.Block(statement);
             return true;
         }