Add Namespace to class with Roslyn CodeAnalyzer

Question

Piggy backing off Add a parameter to a method with a Roslyn CodeFixProvider,

I'm working on creating a CodeFixProvider that ensures all async methods take a CancellationToken:

//Before Code Fix:
public async Task Example(){}

//After Code Fix
public async Task Example(CancellationToken token){}

I'm able to add the parameter to the method, but I have to do so using the Type.FullName. Instead, I'd like to add a using statement for System.Threading to the top of the class file so the method doesn't need to use the full namespace. In other words:

// What I have thus far:
public class AClass{
  public async Task Example(System.Threading.CancellationToken token){}
}

// What I want:
using System.Threading;

public class AClass{
  public async Task Example(CancellationToken token){}
}

How can I add a using statement to a Document?

I've tried a couple ways, but it seams when I'm replacing multiple nodes in a SyntaxTree references are being lost (as the tree is immutable and reconstructed on every change).

I was able to get it partially working with the code below, but this only works if the CompilationUnitSyntax.Using property is populated, which is not the case when using statements come after a namespace. And this also relies on there being at least one using statement already in the file.

Is there a better way to do this?

private async Task<Document> HaveMethodTakeACancellationTokenParameter(
        Document document, SyntaxNode syntaxNode, CancellationToken cancellationToken)
    {
        var syntaxTree = 
            (await document.GetSyntaxTreeAsync(cancellationToken))
                .GetRoot(cancellationToken);

        var method = syntaxNode as MethodDeclarationSyntax;

        #region Add Parameter
        var newParameter =
            SyntaxFactory.Parameter(
                SyntaxFactory.Identifier("cancellationToken")
            )
            .WithType(
                SyntaxFactory.ParseTypeName(
                    typeof(CancellationToken).FullName));

        var updatedMethod = method.AddParameterListParameters(newParameter);

        syntaxTree = syntaxTree.ReplaceNode(method, updatedMethod);

        #endregion

        #region Add Using Statements

        var compilation =
            syntaxTree as CompilationUnitSyntax;                    

        var systemThreadingUsingName =
            SyntaxFactory.QualifiedName(
                SyntaxFactory.IdentifierName("System"),
                SyntaxFactory.IdentifierName("Threading"));

        if (compilation.Usings.All(u => u.Name.GetText().ToString() != typeof(CancellationToken).Namespace))
        {
            syntaxTree = syntaxTree.InsertNodesAfter(compilation.Usings.Last(), new[]
            {
                SyntaxFactory.UsingDirective(
                        systemThreadingUsingName)
            });
        }

        #endregion

        return document.WithSyntaxRoot(syntaxTree);
    }

Answer 1

An option is to mark all methods with an annotation, add the using statement, find the methods with annotations, change all methods, and remove the annotations.

Trees are immutable as you said, but annotations are not lost during modifications. So you'll need something like the following:

var annotation = new SyntaxAnnotation();
var newRoot = root.ReplaceNode(
     method,
     method.WithAdditionalAnnotations(annotation));
newRoot = AddUsing(newRoot);
method = newRoot.GetAnnotatedNodes(annotation).First();
var newMethod = ChangeParameters(method);
newRoot = root.ReplaceNode(method, newMethod.WithoutAnnotations(annotation));

---

Full Implementation:

 private async Task<Document> HaveMethodTakeACancellationTokenParameter(
        Document document, SyntaxNode syntaxNode, CancellationToken cancellationToken)
    {
        var method = syntaxNode as MethodDeclarationSyntax;

        var cancellationTokenParameter =
           SyntaxFactory.Parameter(
               SyntaxFactory.Identifier("cancellationToken")
           )
           .WithType(
               SyntaxFactory.ParseTypeName(
                   typeof(CancellationToken).Name));

        var root = 
           (await document.GetSyntaxTreeAsync(cancellationToken))
               .GetRoot(cancellationToken);

        var annotation = new SyntaxAnnotation();
        var newRoot = root.ReplaceNode(
             method,
             method.WithAdditionalAnnotations(annotation));

        #region Add Using Statements

        var systemThreadingUsingStatement =
            SyntaxFactory.UsingDirective(
                SyntaxFactory.QualifiedName(
                    SyntaxFactory.IdentifierName("System"),
                    SyntaxFactory.IdentifierName("Threading")));

        var compilation =
            newRoot as CompilationUnitSyntax;

        if (null == compilation)
        {
            newRoot =
                newRoot.InsertNodesBefore(
                    newRoot.ChildNodes().First(),
                    new[] {systemThreadingUsingStatement});
        }
        else if (compilation.Usings.All(u => u.Name.GetText().ToString() != typeof(CancellationToken).Namespace))
        {
            newRoot = 
                newRoot.InsertNodesAfter(compilation.Usings.Last(), 
                new[]{ systemThreadingUsingStatement });
        }

        #endregion

        method = (MethodDeclarationSyntax)newRoot.GetAnnotatedNodes(annotation).First();

        var updatedMethod = method.AddParameterListParameters(cancellationTokenParameter);
        newRoot = newRoot.ReplaceNode(method, updatedMethod.WithoutAnnotations(annotation));

        return document.WithSyntaxRoot(newRoot);
    }