Semantic snippets: allow cw snippet in void-returning lambdas (#73706)

* Allow `Console.WriteLine` snippet in void-returning lambdas * Strengthen the constraint Co-authored-by: Cyrus Najmabadi <[email protected]> --------- Co-authored-by: Cyrus Najmabadi <[email protected]>
dotnet · May 29, 2024 · 392c26b · 392c26b
1 parent 463678f
commit 392c26b
Show file tree

Hide file tree

Showing 3 changed files with 145 additions and 33 deletions.
diff --git a/...st/Completion/CompletionProviders/Snippets/CSharpConsoleSnippetCompletionProviderTests.cs b/...st/Completion/CompletionProviders/Snippets/CSharpConsoleSnippetCompletionProviderTests.cs
@@ -491,4 +491,94 @@ public void Method()
             """;
         await VerifyCustomCommitProviderAsync(markupBeforeCommit, ItemToCommit, expectedCodeAfterCommit);
     }
+
+    [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/72266")]
+    public async Task InsertConsoleSnippetInVoidReturningLambdaTest1()
+    {
+        var markupBeforeCommit = """
+            using System;
+
+            M(() => $$);
+
+            void M(Action a)
+            {
+            }
+            """;
+
+        var expectedCodeAfterCommit = """
+            using System;
+
+            M(() => Console.WriteLine($$));
+
+            void M(Action a)
+            {
+            }
+            """;
+
+        await VerifyCustomCommitProviderAsync(markupBeforeCommit, ItemToCommit, expectedCodeAfterCommit);
+    }
+
+    [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/72266")]
+    public async Task InsertConsoleSnippetInVoidReturningLambdaTest2()
+    {
+        var markupBeforeCommit = """
+            using System;
+
+            Action action = () => $$
+            """;
+
+        var expectedCodeAfterCommit = """
+            using System;
+
+            Action action = () => Console.WriteLine($$)
+            """;
+
+        await VerifyCustomCommitProviderAsync(markupBeforeCommit, ItemToCommit, expectedCodeAfterCommit);
+    }
+
+    [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/72266")]
+    public async Task InsertConsoleSnippetInVoidReturningLambdaTest_TypeInference()
+    {
+        var markupBeforeCommit = """
+            using System;
+
+            var action = () => $$
+            """;
+
+        var expectedCodeAfterCommit = """
+            using System;
+
+            var action = () => Console.WriteLine($$)
+            """;
+
+        await VerifyCustomCommitProviderAsync(markupBeforeCommit, ItemToCommit, expectedCodeAfterCommit);
+    }
+
+    [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/72266")]
+    public async Task NoConsoleSnippetInNonVoidReturningLambdaTest1()
+    {
+        var markupBeforeCommit = """
+            using System;
+
+            M(() => $$);
+
+            void M(Func<int> f)
+            {
+            }
+            """;
+
+        await VerifyItemIsAbsentAsync(markupBeforeCommit, ItemToCommit);
+    }
+
+    [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/72266")]
+    public async Task NoConsoleSnippetInNonVoidReturningLambdaTest2()
+    {
+        var markupBeforeCommit = """
+            using System;
+
+            Func<int> f = () => $$
+            """;
+
+        await VerifyItemIsAbsentAsync(markupBeforeCommit, ItemToCommit);
+    }
 }
diff --git a/src/Features/CSharp/Portable/Snippets/CSharpConsoleSnippetProvider.cs b/src/Features/CSharp/Portable/Snippets/CSharpConsoleSnippetProvider.cs
@@ -4,6 +4,7 @@
 
 using System;
 using System.Composition;
+using System.Threading;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Snippets;
@@ -17,10 +18,33 @@ namespace Microsoft.CodeAnalysis.CSharp.Snippets;
 internal sealed class CSharpConsoleSnippetProvider() : AbstractConsoleSnippetProvider<
     ExpressionStatementSyntax,
     ExpressionSyntax,
-    ArgumentListSyntax>
+    ArgumentListSyntax,
+    LambdaExpressionSyntax>
 {
-    protected override ExpressionSyntax GetExpression(ExpressionStatementSyntax expressionStatement)
-        => expressionStatement.Expression;
+    protected override bool IsValidSnippetLocation(in SnippetContext context, CancellationToken cancellationToken)
+    {
+        var syntaxContext = context.SyntaxContext;
+
+        var consoleSymbol = GetConsoleSymbolFromMetaDataName(syntaxContext.SemanticModel.Compilation);
+        if (consoleSymbol is null)
+            return false;
+
+        // Console.WriteLine snippet is legal after an arrow token of a void-returning lambda, e.g.
+        // Action a = () => Console.WriteLine("Action called");
+        if (syntaxContext.TargetToken is { RawKind: (int)SyntaxKind.EqualsGreaterThanToken, Parent: LambdaExpressionSyntax lambda })
+        {
+            var semanticModel = syntaxContext.SemanticModel;
+            var lambdaSymbol = semanticModel.GetSymbolInfo(lambda, cancellationToken).Symbol;
+
+            // Given that we are in a partially written lambda state compiler might not always infer return type correctly.
+            // In such cases an error type is returned. We allow them to provide snippet in locations
+            // where lambda return type isn't yet known, but it might be a void type after fully completing the lambda
+            if (lambdaSymbol is IMethodSymbol { ReturnType: { SpecialType: SpecialType.System_Void } or { TypeKind: TypeKind.Error } })
+                return true;
+        }
+
+        return syntaxContext.IsStatementContext || syntaxContext.IsGlobalStatementContext;
+    }
 
     protected override ArgumentListSyntax GetArgumentList(ExpressionSyntax expression)
         => ((InvocationExpressionSyntax)expression).ArgumentList;

diff --git a/src/Features/Core/Portable/Snippets/SnippetProviders/AbstractConsoleSnippetProvider.cs b/src/Features/Core/Portable/Snippets/SnippetProviders/AbstractConsoleSnippetProvider.cs
@@ -10,6 +10,7 @@
 using Microsoft.CodeAnalysis.Formatting;
 using Microsoft.CodeAnalysis.LanguageService;
 using Microsoft.CodeAnalysis.Shared.Extensions;
+using Microsoft.CodeAnalysis.Shared.Extensions.ContextQuery;
 using Microsoft.CodeAnalysis.Simplification;
 using Microsoft.CodeAnalysis.Snippets.SnippetProviders;
 using Microsoft.CodeAnalysis.Text;
@@ -20,52 +21,49 @@ namespace Microsoft.CodeAnalysis.Snippets;
 internal abstract class AbstractConsoleSnippetProvider<
     TExpressionStatementSyntax,
     TExpressionSyntax,
-    TArgumentListSyntax> : AbstractStatementSnippetProvider<TExpressionStatementSyntax>
+    TArgumentListSyntax,
+    TLambdaExpressionSyntax> : AbstractSingleChangeSnippetProvider<TExpressionSyntax>
     where TExpressionStatementSyntax : SyntaxNode
     where TExpressionSyntax : SyntaxNode
     where TArgumentListSyntax : SyntaxNode
+    where TLambdaExpressionSyntax : TExpressionSyntax
 {
     public sealed override string Identifier => CommonSnippetIdentifiers.ConsoleWriteLine;
 
     public sealed override string Description => FeaturesResources.console_writeline;
 
     public sealed override ImmutableArray<string> AdditionalFilterTexts { get; } = ["WriteLine"];
 
-    protected abstract TExpressionSyntax GetExpression(TExpressionStatementSyntax expressionStatement);
     protected abstract TArgumentListSyntax GetArgumentList(TExpressionSyntax expression);
     protected abstract SyntaxToken GetOpenParenToken(TArgumentListSyntax argumentList);
 
-    protected sealed override bool IsValidSnippetLocation(in SnippetContext context, CancellationToken cancellationToken)
+    protected sealed override async Task<TextChange> GenerateSnippetTextChangeAsync(Document document, int position, CancellationToken cancellationToken)
     {
-        var consoleSymbol = GetConsoleSymbolFromMetaDataName(context.SyntaxContext.SemanticModel.Compilation);
-        if (consoleSymbol is null)
-            return false;
+        var generator = SyntaxGenerator.GetGenerator(document);
 
-        return base.IsValidSnippetLocation(in context, cancellationToken);
-    }
+        var resultingNode = generator.InvocationExpression(generator.MemberAccessExpression(generator.IdentifierName(nameof(Console)), nameof(Console.WriteLine)));
 
-    protected sealed override Task<TextChange> GenerateSnippetTextChangeAsync(Document document, int position, CancellationToken cancellationToken)
-    {
-        var generator = SyntaxGenerator.GetGenerator(document);
+        var semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
+        var syntaxContext = document.GetRequiredLanguageService<ISyntaxContextService>().CreateContext(document, semanticModel, position, cancellationToken);
 
-        var invocation = generator.InvocationExpression(generator.MemberAccessExpression(generator.IdentifierName(nameof(Console)), nameof(Console.WriteLine)));
-        var expressionStatement = generator.ExpressionStatement(invocation);
+        // In case we are after an arrow token in lambda, Console.WriteLine acts like an expression,
+        // so it doesn't need to be wrapped into a statement
+        if (syntaxContext.TargetToken.Parent is not TLambdaExpressionSyntax)
+        {
+            resultingNode = generator.ExpressionStatement(resultingNode);
+        }
 
-        var change = new TextChange(TextSpan.FromBounds(position, position), expressionStatement.ToFullString());
-        return Task.FromResult(change);
+        var change = new TextChange(TextSpan.FromBounds(position, position), resultingNode.ToFullString());
+        return change;
     }
 
     /// <summary>
     /// Tries to get the location after the open parentheses in the argument list.
     /// If it can't, then we default to the end of the snippet's span.
     /// </summary>
-    protected sealed override int GetTargetCaretPosition(TExpressionStatementSyntax caretTarget, SourceText sourceText)
+    protected sealed override int GetTargetCaretPosition(TExpressionSyntax caretTarget, SourceText sourceText)
     {
-        var invocationExpression = GetExpression(caretTarget);
-        if (invocationExpression is null)
-            return caretTarget.Span.End;
-
-        var argumentListNode = GetArgumentList(invocationExpression);
+        var argumentListNode = GetArgumentList(caretTarget);
         if (argumentListNode is null)
             return caretTarget.Span.End;
 
@@ -86,25 +84,25 @@ protected sealed override async Task<SyntaxNode> AnnotateNodesToReformatAsync(
         return root.ReplaceNode(snippetExpressionNode, reformatSnippetNode);
     }
 
-    protected sealed override ImmutableArray<SnippetPlaceholder> GetPlaceHolderLocationsList(TExpressionStatementSyntax node, ISyntaxFacts syntaxFacts, CancellationToken cancellationToken)
+    protected sealed override ImmutableArray<SnippetPlaceholder> GetPlaceHolderLocationsList(TExpressionSyntax node, ISyntaxFacts syntaxFacts, CancellationToken cancellationToken)
         => [];
 
-    private static INamedTypeSymbol? GetConsoleSymbolFromMetaDataName(Compilation compilation)
+    protected static INamedTypeSymbol? GetConsoleSymbolFromMetaDataName(Compilation compilation)
         => compilation.GetBestTypeByMetadataName(typeof(Console).FullName!);
 
-    protected sealed override TExpressionStatementSyntax? FindAddedSnippetSyntaxNode(SyntaxNode root, int position)
+    protected sealed override TExpressionSyntax? FindAddedSnippetSyntaxNode(SyntaxNode root, int position)
     {
         var closestNode = root.FindNode(TextSpan.FromBounds(position, position));
-        var nearestExpressionStatement = closestNode.FirstAncestorOrSelf<TExpressionStatementSyntax>();
-        if (nearestExpressionStatement is null)
+        var nearestExpression = closestNode.FirstAncestorOrSelf<TExpressionSyntax>(static exp => exp.Parent is TExpressionStatementSyntax or TLambdaExpressionSyntax);
+        if (nearestExpression is null)
             return null;
 
-        // Checking to see if that expression statement that we found is
+        // Checking to see if that expression that we found is
         // starting at the same position as the position we inserted
-        // the Console WriteLine expression statement.
-        if (nearestExpressionStatement.SpanStart != position)
+        // the Console WriteLine expression.
+        if (nearestExpression.SpanStart != position)
             return null;
 
-        return nearestExpressionStatement;
+        return nearestExpression;
     }
 }