Implementing Code Action with Roslyn

Original author: Brian Rasmussen
  • Transfer
  • Tutorial
The Roslyn Services API makes it easy to implement extensions that find and fix code problems right in Visual Studio. The Roslyn Services API is available as part of the Roslyn CTP .

In this post, we are implementing an extension for Visual Studio that detects calls to the Count () method of Enumerable, after which the result is checked for equality greater than zero, for example, someSequence.Count ()> 0. The problem with the code is that Count () must go through the entire sequence before returning the result. A more appropriate approach in this case is to call the Enumerable.Any () method.

To fix this, we implement CodeIssueProvider, which detects the problem, and CodeAction, which replaces the condition with the call to Enumerable.Any (), as required. Those. our CodeAction will change something like someSequence.Count ()> 0 to someSequence.Any ().

There are a couple of additional conditions that we would also like to fulfill: first of all, the expression can be inverted and written as 0 <someSequence.Count (). The next case is a record of type> = 1 instead of> 0, which is logically the same as before. We need the extension to work in both cases.

Obviously, we would not want to change not all calls with the signature Count (), but only if they relate to the extension method from IEnumerable defined in Enumerable.

Beginning of work

Roslyn CTP comes with a set of templates to help you get started with your API. To get started, we will create a new project of type Code Issue from Roslyn templates. Let's name the project as ReplaceCountWithAny.



The template generates a simple working provider that highlights words containing the letter “a”. To see an example in action, we will build and run a project created by a template. This starts a new instance of Visual Studio with the extension enabled. From the just launched Visual Studio we will create a console application and see how the keywords namespace, class, etc. underlined by our extension.



Although the example is not as useful as an extension for Visual Studio, it prepares everything you need to start implementing your own extension. We only have to replace the contents of the generated GetIssue method. I note that there are three overloads for GetIssues. We will work with overloading, where one of the parameters is of type CommonSyntaxNode. The remaining two overloads can be left in our case.

The generated CodeIssueProvider class implements the ICodeIssueProvider interface and is decorated with the ExportSyntaxNodeCodeIssueProvide attribute. This allows Visual Studio to import this type as an extension containing the contract provided by the ICodeIssueProvide interface.

We implement GetIssues

Our GetIssues method will be called for each syntax construct, so the first thing we should do is weed out the nodes we are not interested in. Since we need constructs of the type someSequence.Count ()> 0, we only need nodes of the BinaryExpressionSyntax type. We can tell Visual Studio to use our provider only for specific nodes by providing a list of types through the ExportSyntaxNodeCodeIssueProvide attribute. So, update the attribute as shown below:
[ExportSyntaxNodeCodeIssueProvider("ReplaceCountWithAny", 
LanguageNames.CSharp, typeof(BinaryExpressionSyntax))]
class CodeIssueProvider : ICodeIssueProvider ...

This allows you to safely cast the CommonSyntaxNode node to the BinaryExpressionSyntax type in the GetIssues method.

To highlight the cases that we want to handle, you need to check part of the expression for the call to Enumerable.Count (), and another for the comparison itself. We will isolate the verification data into auxiliary methods, so our GetIssues implementation will look like this:
public IEnumerable GetIssues(IDocument document, 
    CommonSyntaxNode node, CancellationToken cancellationToken)
{
    var binaryExpression = (BinaryExpressionSyntax)node;
    var left = binaryExpression.Left;
    var right = binaryExpression.Right;
    var kind = binaryExpression.Kind;
    if (IsCallToEnumerableCount(document, left, cancellationToken) && 
        IsRelevantRightSideComparison(document, right, kind, cancellationToken) ||
        IsCallToEnumerableCount(document, right, cancellationToken) && 
        IsRelevantLeftSideComparison(document, left, kind, cancellationToken))
    {
        yield return new CodeIssue(CodeIssue.Severity.Info, binaryExpression.Span,
            string.Format("Change {0} to use Any() instead of Count() to avoid " +
                          "possible enumeration of entire sequence.", 
                          binaryExpression));
    }
}


The instance of the CodeIssue class that we return indicates the level of the problem, which may be Error, Warning, or Info, the description used to highlight the section of code, and text that describes the problem to the user.

Helper Methods

Now we turn our attention to the helper methods used in GetIssues. The IsCallToEnumerableCount method returns true if the part of the expression that we are considering is a call to the Count () method on some sequence. Let me remind you again: we start first by filtering unnecessary expressions.

First of all, the expression must be a method call. In this case, we get the necessary call from the expression property. So, if the design looks like someSequence.Count ()> 0, then we will have a part with Count (); but how to check whether it belongs to the Enumerable type?

To answer such questions, you need to request a sematic model. Fortunately, one of the parameters of the GetIssues method is IDocument, which is a document in the project and solution. We can get a semantic model through it, and already from it SymbolInfo itself, which we need.
With the help of SymbolInfo you can check whether our method call belongs to the desired [Enumerable.Count ()]. Since Count () is an extension method, working with it will be slightly different. Recall that C # allows extension methods to be called as part of a type. The semantic model provides this information through the ConstructedFrom property of the MethodSymbol class with reference to the original type. There is an opportunity to make this a little easier, so stay tuned for the API names.

All that remains for us to do is to indicate the type of extension method. If it matches Enumerable, then we found a call to Enumerable.Count ().

The implementation is as follows:
private bool IsCallToEnumerableCount(IDocument document, 
    ExpressionSyntax expression, CancellationToken cancellationToken)
{
    var invocation = expression as InvocationExpressionSyntax;
    if (invocation == null)
    {
        return false;
    }
    var call = invocation.Expression as MemberAccessExpressionSyntax;
    if (call == null)
    {
        return false;
    }
    var semanticModel = document.GetSemanticModel(cancellationToken);
    var methodSymbol = semanticModel.GetSemanticInfo(call, cancellationToken).Symbol 
        as MethodSymbol;
    if (methodSymbol == null || 
        methodSymbol.Name != "Count" || 
        methodSymbol.ConstructedFrom == null)
    {
        return false;
    }
    var enumerable = semanticModel.Compilation.GetTypeByMetadataName(
        typeof(Enumerable).FullName);
    if (enumerable == null || 
        !methodSymbol.ConstructedFrom.ContainingType.Equals(enumerable))
    {
        return false;
    }
    return true;
} 

Before moving forward, you must also check the expression for the correctness of comparison in another part of the binary expression; and this is the work for helper methods IsRelevantRightSideComparison and IsRelevantLeftSideComparison.

Below are their implementations:

private bool IsRelevantRightSideComparison(IDocument document, 
    ExpressionSyntax expression, SyntaxKind kind, 
    CancellationToken cancellationToken)
{
    var semanticInfo = document.GetSemanticModel(cancellationToken).
        GetSemanticInfo(expression);
    int? value;
    if (!semanticInfo.IsCompileTimeConstant || 
        (value = semanticInfo.ConstantValue as int?) == null)
    {
        return false;
    }
    if (kind == SyntaxKind.GreaterThanExpression && value == 0 ||
        kind == SyntaxKind.GreaterThanOrEqualExpression && value == 1)
    {
        return true;
    }
    return false;
}
private bool IsRelevantLeftSideComparison(IDocument document, 
    ExpressionSyntax expression, SyntaxKind kind, 
    CancellationToken cancellationToken)
{
    var semanticInfo = document.GetSemanticModel(cancellationToken).
        GetSemanticInfo(expression);
    int? value;
    if (!semanticInfo.IsCompileTimeConstant ||
        (value = semanticInfo.ConstantValue as int?) == null)
    {
        return false;
    }
    if (kind == SyntaxKind.LessThanExpression && value == 0 ||
        kind == SyntaxKind.LessThanOrEqualExpression && value == 1)
    {
        return true;
    }
    return false;
}

Yes, they are almost identical with the only difference being that both variants of comparison are checked, as well as the correctness of the value itself, so that we do not have to highlight something like Count ()> = 0.

Testing CodeIssueProvider

At the moment, our provider is able to detect problems of interest to us. We compile and run the project along with a new instance of Visual Studio along with the extension included. Add code and notice that calls to Enumerable.Count () are underlined correctly, while calls to other methods with a signature of Count () are not.



The next step is to provide action to solve the problem.

CodeAction

To implement the action, we need a class that implements the ICodeAction interface. ICodeAction is a simple interface that defines a description and an icon for an action, as well as the only GetEdit method that returns a revision that transforms the current syntax tree. So, let's start with the constructor of our CodeAction class.

public CodeAction(ICodeActionEditFactory editFactory, 
    IDocument document, BinaryExpressionSyntax binaryExpression)
{
    this.editFactory = editFactory;
    this.document = document;
    this.binaryExpression = binaryExpression;
}

For each problem found, a new instance of the CodeAction class will be created, so for convenience we will omit some parameters and change the constructor itself. This requires the implementation of ICodeActionEditFactory to create a transformation of the newly created syntax tree. Since the syntax trees in the Roslyn project are immutable, returning a new tree is the only way to make any changes. Fortunately, Roslyn tries to reuse the tree as much as possible, thus preventing the creation of redundant syntax nodes.

In addition, you need a document that gives our code access to the syntax tree, project and solution, as well as a link to the syntax node that we want to replace.

So, we got closer to the GetEdit method. It is here that we create a transformation that replaces the detected binary expression with a new one with a call to the Any () method. Creating a new node is assigned to a simple helper method GetNewNode. The implementation of both methods is given below:

public ICodeActionEdit GetEdit(CancellationToken cancellationToken)
{
    var syntaxTree = (SyntaxTree)document.GetSyntaxTree(cancellationToken);
    var newExpression = GetNewNode(binaryExpression).
        WithLeadingTrivia(binaryExpression.GetLeadingTrivia()).
        WithTrailingTrivia(binaryExpression.GetTrailingTrivia());
    var newRoot = syntaxTree.Root.ReplaceNode(binaryExpression, newExpression);
    return editFactory.CreateTreeTransformEdit(
        document.Project.Solution,
        syntaxTree,
        newRoot,
        cancellationToken: cancellationToken);
}
private ExpressionSyntax GetNewNode(BinaryExpressionSyntax node)
{
    var invocation = node.DescendentNodes().
        OfType().Single();
    var caller = invocation.DescendentNodes().
        OfType().Single();
    return invocation.Update(
        caller.Update(caller.Expression, 
        caller.OperatorToken, 
        Syntax.IdentifierName("Any")),
        invocation.ArgumentList);
}

The Roslyn syntax tree is exactly the same as the original code, so each node in the tree can contain extra spaces and comments. Those. we save the original node along with comments and code structure when changing the nodes themselves. To do this, we call the extension methods WithLeadingTrivia and WithTrailingTrivia.

I also note that the GetNewNode method saves the list of parameters of the Count () method, so if the extension method was called via a lambda expression to count specific elements in the sequence, it will still be replaced with Any ().

To summarize

To enable our action, you must update the GetIssues method in the CodeIssueProvider class to return a CodeAction instance for each CodeIssue. Each problematic section of code can have several actions, allowing the user to choose between them. In this case, we return a single action as shown below.

The updated part of the GetIssues method is as follows:

yield return new CodeIssue(CodeIssue.Severity.Info, binaryExpression.Span,
    string.Format("Change {0} to use Any() instead of Count() to avoid " +
                  "possible enumeration of entire sequence.", binaryExpression),
    new CodeAction(editFactory, document, binaryExpression));

We rebuild and run the project to launch a new instance of Visual Studio with the downloaded extension. Now we see that the problem area provides a drop-down list with options for fixing the code.



Thus, we have implemented an extension for Visual Studio, which will help improve our code.

Also popular now: