Generate code for classes with an attribute

I have the following setup :

public class CustomAttribute : Attribute
{
   [...]
   public CustomAttribute(Type type)
   {
    [...]
   }
}

[Custom(typeof(Class2))]
public class Class1
{
    public void M1(Class2) {}
    public void M2(Class2) {}
}


public partial class Class2
{
[...]
}

What I am trying to achieve using the new Code Generation mechanism added in .NET 5 is at compile-time, find every class in the project referencing my generator being annotated by the Custom attribute, and then, create a partial class for the type in its constructor containing methods having the same name and parameters (It won't be the same parameters, it's just to simplify a bit).

Before, I was planning to use TTs to generate the partial file but creating one per type was announcing itself to be both tedious and hard to maintain.

Thing is...

I'm a little lost.

What I did manage to do:

Create a generator, making sure it is called at generation and the code it generates is usable (~ a hello world version)
Find my attribute symbol in the compilation context (not sure I'll need it, but I found it)
Found a way to identify the classes being annotated by my attribute by relying on the syntax trees present in the compilation context.

Now, though I don't know how to proceed further, the syntax tree has at the same level the identifier nodes for my attribute and the class being used as a parameter, meaning if I ever use another attribute, I fear they will all get at the same level (might use the order getting the position of the identifier for my attribute and then getting the next one).

But then even if we omit that... How can I list all methods and their parameters for a given class that I have the name of? Reflection is obviously out of the picture since the assembly is not loaded.

I only found Rosly examples, based on using the solution or Analyzers who don't really have the same type of objects available, and thus the proposed solutions are not applicable. And I'm not sure starting another Roslyn analysis on single files is really the way it is supposed to be done.

Please keep in mind that this is my first foray with the Semantic/Syntaxic API.

Preparatory work: Introduction to C# source generators

This will guide you towards settings up a code generator project. As far as I understood, tooling is on the way to automate this part.

TL;DR there will be the full ExecuteMethod at the end of this answer

Filtering out the syntaxic trees containing no classes decorated with an attribute

This is our first step, we only want to work with the classes that are decorated by an attribute, we'll then make sure it's the one that interests us. This also has the secondary benefit of filtering out any source files that do not contain classes (think AssemblyInfo.cs)

In the Execute method of our new Generator, we will be able to use Linq to filter out the trees:

var treesWithlassWithAttributes = context.Compilation.SyntaxTrees.Where(st => st.GetRoot().DescendantNodes().OfType<ClassDeclarationSyntax>()
                    .Any(p => p.DescendantNodes().OfType<AttributeSyntax>().Any()));

We will then be able to loop on our syntaxic trees (from what I could see, one syntaxic tree corresponds roughly to one file)

Filter-out classes not being annotated with an attribute

The next step is to make sure that in our current syntaxic tree, we only work on classes being decorated by an attribute (for the cases where several classes are declared in one file).

var declaredClass = tree
                    .GetRoot()
                    .DescendantNodes()
                    .OfType<ClassDeclarationSyntax>()
                    .Where(cd => cd.DescendantNodes().OfType<AttributeSyntax>().Any())

This is all pretty similar to the previous step, tree being one of the items we got in our treesWithClassWithAttributes collection.

Once again, we'll loop on that collection.

Filter-out classes that are not annotated with our specific attribute

Now, that we are working on single classes, we can dive in and check if any of the attributes are the one we are looking for. This will also be the first time we will need the semantic API, since the attribute identifier is not it's class name (PropertyAttribute, will be used as [Property] for example), and the semantic API, allows us to find the original class name without us having to guess.

We will first need to initialize our semantic model (this should placed in our top level loop):

var semanticModel = context.Compilation.GetSemanticModel(tree);

Once initialized, we get to our searching :

var nodes = declaredClass
                    .DescendantNodes()
                    .OfType<AttributeSyntax>()
                    .FirstOrDefault(a => a.DescendantTokens().Any(dt => dt.IsKind(SyntaxKind.IdentifierToken) && semanticModel.GetTypeInfo(dt.Parent).Type.Name == attributeSymbol.Name))
                    ?.DescendantTokens()
                    ?.Where(dt => dt.IsKind(SyntaxKind.IdentifierToken))
                    ?.ToList();

Note: attributeSymbol is a variable holding the Type of the attribute I am searching for

What we are doing here, is for each syntaxic node related to our class declaration, we only look at the ones describing an attribute declaration.

We then take the first one (my attribute can only be placed once on a class) that has an IdentifierToken for which the parent node Is of the type of my attribute (the semantic API does NOT return a Type hence the name Comparison).

For the next steps, I will need the IdentifiersToken, so we'll use the Elvis operator to get them if we found our attribute, we'll get a null result otherwise which will allow us to get to the next iteration of our loop.

Get the class type used as my Attribute parameter

This is where it gets really specific to my use case, but it's part of the question, so I'll cover it anyway.

What we got at the end of the last step was a list of Identifier Tokens, which mean, we will have only two for my attribute : The first one Identifying the attribute itself, and the second one identifying the class I want to get the name of.

We will be using the semantic API again, this allows me to avoid looking in all the syntax trees to find the class we identified :

var relatedClass = semanticModel.GetTypeInfo(nodes.Last().Parent);

This gives us an object similar to the ones we were manipulating until now.

This is a good point to start generating our new class file (so a new stringbuilder, with all the test needed to have a partial class in the same namespace the other one was, it will always be the same in my case, so I went and wrote it directly)

To get the name of the type in relatedClass => relatedClass.Type.Name

List all methods used in a class

So now, to list all the methods in the annotated class. Remember we are looping on classes here, coming from our syntactic tree.

To obtain a list of all the methods declared in this class we will ask to list the member of type method

IEnumerable<MethodDeclarationSyntax> classMethod = declaredClass.Members.Where(m => m.IsKind(SyntaxKind.MethodDeclaration)).OfType<MethodDeclarationSyntax>()

I will strongly recommend either casting to MethodDeclarationSyntax or assigning to a variable with the explicit type, because it is stored as a base type that does not expose all of the properties we'll need.

Once we've got our methods, we will once again loop on them. Here are the few properties I needed for my use case :

methodDeclaration.Modifiers //public, static, etc...
methodDeclaration.Identifier // Quite obvious => the name
methodDeclaration.ParameterList //The list of the parameters, including type, name, default values

The rest was just a matter of constructing a string representing my target partial class which is now a pretty simple matter.

The final solution

Remember that's what I came up with as my first try, I will most probably submit it on the CodeReview StackExchange to see what could be improved.

And the RelatedModelaAttribute is basically the CustomAttribute class from my question.

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Text;
using SpeedifyCliWrapper.SourceGenerators.Annotations;
using System.Linq;
using System.Text;

namespace SpeedifyCliWrapper.SourceGenerators
{
    [Generator]
    class ModuleModelGenerator : ISourceGenerator
    {
        public void Execute(GeneratorExecutionContext context)
        {
            var attributeSymbol = context.Compilation.GetTypeByMetadataName(typeof(RelatedModelAttribute).FullName);

            var classWithAttributes = context.Compilation.SyntaxTrees.Where(st => st.GetRoot().DescendantNodes().OfType<ClassDeclarationSyntax>()
                    .Any(p => p.DescendantNodes().OfType<AttributeSyntax>().Any()));

            foreach (SyntaxTree tree in classWithAttributes)
            {
                var semanticModel = context.Compilation.GetSemanticModel(tree);
                
                foreach(var declaredClass in tree
                    .GetRoot()
                    .DescendantNodes()
                    .OfType<ClassDeclarationSyntax>()
                    .Where(cd => cd.DescendantNodes().OfType<AttributeSyntax>().Any()))
                {
                    var nodes = declaredClass
                    .DescendantNodes()
                    .OfType<AttributeSyntax>()
                    .FirstOrDefault(a => a.DescendantTokens().Any(dt => dt.IsKind(SyntaxKind.IdentifierToken) && semanticModel.GetTypeInfo(dt.Parent).Type.Name == attributeSymbol.Name))
                    ?.DescendantTokens()
                    ?.Where(dt => dt.IsKind(SyntaxKind.IdentifierToken))
                    ?.ToList();

                    if(nodes == null)
                    {
                        continue;
                    }

                    var relatedClass = semanticModel.GetTypeInfo(nodes.Last().Parent);

                    var generatedClass = this.GenerateClass(relatedClass);

                    foreach(MethodDeclarationSyntax classMethod in declaredClass.Members.Where(m => m.IsKind(SyntaxKind.MethodDeclaration)).OfType<MethodDeclarationSyntax>())
                    {
                        this.GenerateMethod(declaredClass.Identifier, relatedClass, classMethod, ref generatedClass);
                    }

                    this.CloseClass(generatedClass);

                    context.AddSource($"{declaredClass.Identifier}_{relatedClass.Type.Name}", SourceText.From(generatedClass.ToString(), Encoding.UTF8));
                }
            }
        }

        public void Initialize(GeneratorInitializationContext context)
        {
            // Nothing to do here
        }

        private void GenerateMethod(SyntaxToken moduleName, TypeInfo relatedClass, MethodDeclarationSyntax methodDeclaration, ref StringBuilder builder)
        {
            var signature = $"{methodDeclaration.Modifiers} {relatedClass.Type.Name} {methodDeclaration.Identifier}(";

            var parameters = methodDeclaration.ParameterList.Parameters.Skip(1);

            signature += string.Join(", ", parameters.Select(p => p.ToString())) + ")";

            var methodCall = $"return this._wrapper.{moduleName}.{methodDeclaration.Identifier}(this, {string.Join(", ", parameters.Select(p => p.Identifier.ToString()))});";

            builder.AppendLine(@"
        " + signature + @"
        {
            " + methodCall + @"
        }");
        }

        private StringBuilder GenerateClass(TypeInfo relatedClass)
        {
            var sb = new StringBuilder();

            sb.Append(@"
using System;
using System.Collections.Generic;
using SpeedifyCliWrapper.Common;

namespace SpeedifyCliWrapper.ReturnTypes
{
    public partial class " + relatedClass.Type.Name);

            sb.Append(@"
    {");

            return sb;
        }
        private void CloseClass(StringBuilder generatedClass)
        {
            generatedClass.Append(
@"    }
}");
        }
    }
}

Inside the Execute method of your generator, add this:

var classesWithAttribute = context.Compilation.SyntaxTrees
                .SelectMany(st => st.GetRoot()
                        .DescendantNodes()
                        .Where(n => n is ClassDeclarationSyntax)
                        .Select(n => n as ClassDeclarationSyntax)
                        .Where(r => r.AttributeLists
                            .SelectMany(al => al.Attributes)
                            .Any(a => a.Name.GetText().ToString() == "Foo")));

This basically takes all the nodes of all the trees, filters out nodes that are not class declarations, and, for each class declaration, looks if any of its attributes matches our custom attribute, "Foo" here.

Note: if your attribute is named FooAttribute, then you look for Foo, not FooAttribute.