Replacing a method node using Roslyn

Question

While exploring Roslyn I put together a small app that should include a trace statement as the first statement in every method found in a Visual Studio Solution. My code is buggy and is only updating the first method.

The line that is not working as expected is flagged with a “TODO” comment. Please, advise.

I also welcome style recommendations that would create a more streamlined/readable solution.

Thanks in advance.

...

    private void TraceBtn_Click(object sender, RoutedEventArgs e) {
        var myWorkSpace = new MyWorkspace("...Visual Studio 2012\Projects\Tests.sln"); 
        myWorkSpace.InjectTrace();
        myWorkSpace.ApplyChanges();
    }

...

using System;
using System.Linq;
using Roslyn.Compilers;
using Roslyn.Compilers.CSharp;
using Roslyn.Services;

namespace InjectTrace
{
    public class MyWorkspace
    {
    private string solutionFile;
    public string SolutionFile {
        get { return solutionFile; }
        set { 
            if (string.IsNullOrEmpty(value)) throw new Exception("Invalid Solution File");
            solutionFile = value;
        }
    }

    private IWorkspace loadedWorkSpace;
    public IWorkspace LoadedWorkSpace { get { return loadedWorkSpace; } }

    public ISolution CurrentSolution { get; private set; }
    public IProject CurrentProject { get; private set; }
    public IDocument CurrentDocument { get; private set; }
    public ISolution NewSolution { get; private set; }


    public MyWorkspace(string solutionFile) {
        this.SolutionFile = solutionFile;
        this.loadedWorkSpace = Workspace.LoadSolution(SolutionFile);
    }

    public void InjectTrace()
    {

        int projectCtr = 0;
        int documentsCtr = 0;
        int transformedMembers = 0;
        int transformedClasses = 0;
        this.CurrentSolution = this.LoadedWorkSpace.CurrentSolution;
        this.NewSolution = this.CurrentSolution;

        //For Each Project...
        foreach (var projectId in LoadedWorkSpace.CurrentSolution.ProjectIds)
        {
            CurrentProject = NewSolution.GetProject(projectId);

            //..for each Document in the Project..
            foreach (var docId in CurrentProject.DocumentIds)
            {
                CurrentDocument = NewSolution.GetDocument(docId);
                var docRoot = CurrentDocument.GetSyntaxRoot();
                var newDocRoot = docRoot;
                var classes = docRoot.DescendantNodes().OfType<ClassDeclarationSyntax>();
                IDocument newDocument = null;

                //..for each Class in the Document..
                foreach (var @class in classes) {
                    var methods = @class.Members.OfType<MethodDeclarationSyntax>();

                    //..for each Member in the Class..
                    foreach (var currMethod in methods) {
                        //..insert a Trace Statement
                        var newMethod = InsertTrace(currMethod);
                        transformedMembers++;
                        //TODO: PROBLEM IS HERE
                        newDocRoot = newDocRoot.ReplaceNode(currMethod, newMethod);                             
                    }
                    if (transformedMembers != 0) {
                        newDocument = CurrentDocument.UpdateSyntaxRoot(newDocRoot);
                        transformedMembers = 0;
                        transformedClasses++;
                    }
                }

                if (transformedClasses != 0) {
                    NewSolution = NewSolution.UpdateDocument(newDocument);
                    transformedClasses = 0;
                }

                documentsCtr++;

            }
            projectCtr++;
            if (projectCtr > 2) return;
        }
    }

    public MethodDeclarationSyntax InsertTrace(MethodDeclarationSyntax currMethod) {
        var traceText =
        @"System.Diagnostics.Trace.WriteLine(""Tracing: '" + currMethod.Ancestors().OfType<NamespaceDeclarationSyntax>().Single().Name + "." + currMethod.Identifier.ValueText + "'\");";
        var traceStatement = Syntax.ParseStatement(traceText);
        var bodyStatementsWithTrace = currMethod.Body.Statements.Insert(0, traceStatement);
        var newBody = currMethod.Body.Update(Syntax.Token(SyntaxKind.OpenBraceToken), bodyStatementsWithTrace,
                                            Syntax.Token(SyntaxKind.CloseBraceToken));
        var newMethod = currMethod.ReplaceNode(currMethod.Body, newBody);
        return newMethod;

    }

    public void ApplyChanges() {
        LoadedWorkSpace.ApplyChanges(CurrentSolution, NewSolution);
    }


}

}

Answer 1

The root problem of you code is that newDocRoot = newDocRoot.ReplaceNode(currMethod, newMethod); somehow rebuilds newDocRoot internal representation of code so next currMethod elements won't be find in it and next ReplaceNode calls will do nothing. It is a situation similar to modifying a collection within its foreach loop.

The solution is to gather all necessary changes and apply them at once with ReplaceNodes method. And this in fact naturally leads to simplification of code, because we do not need to trace all those counters. We simply store all needed transformation and apply them for whole document at once.

Working code after changes:

public void InjectTrace()
{
    this.CurrentSolution = this.LoadedWorkSpace.CurrentSolution;
    this.NewSolution = this.CurrentSolution;

    //For Each Project...
    foreach (var projectId in LoadedWorkSpace.CurrentSolution.ProjectIds)
    {
        CurrentProject = NewSolution.GetProject(projectId);
        //..for each Document in the Project..
        foreach (var docId in CurrentProject.DocumentIds)
        {
            var dict = new Dictionary<CommonSyntaxNode, CommonSyntaxNode>();
            CurrentDocument = NewSolution.GetDocument(docId);
            var docRoot = CurrentDocument.GetSyntaxRoot();
            var classes = docRoot.DescendantNodes().OfType<ClassDeclarationSyntax>();

            //..for each Class in the Document..
            foreach (var @class in classes)
            {
                var methods = @class.Members.OfType<MethodDeclarationSyntax>();

                //..for each Member in the Class..
                foreach (var currMethod in methods)
                {
                    //..insert a Trace Statement
                    dict.Add(currMethod, InsertTrace(currMethod));
                }
            }

            if (dict.Any())
            {
                var newDocRoot = docRoot.ReplaceNodes(dict.Keys, (n1, n2) => dict[n1]);
                var newDocument = CurrentDocument.UpdateSyntaxRoot(newDocRoot);
                NewSolution = NewSolution.UpdateDocument(newDocument);
            }
        }
    }
}