UnityCodeGenerators/ExampleGenerator/Unity/Ui/UIBackingClassGenerator.cs

using System;
using System.Collections.Generic;
using System.Linq;

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.Text;

using System.Text;

using Microsoft.CodeAnalysis.CSharp.Syntax;


namespace ExampleGenerator.Unity.Ui
{
    public static class Helpers
    {
        public const string UiElementAttribute = "UiElementAttribute";

        internal static bool IsDerivedFrom( INamedTypeSymbol baseType , string targetType )
        {
            while ( baseType != null )
            {
                if ( baseType.Name == targetType )
                    return true;

                baseType = baseType.BaseType;
            }

            return false;
        }
    }

    [Generator]
    public class UiBackingClassGenerator : ISourceGenerator
    {

        private static readonly string UiElementAttributeText = $@"// <auto-generated/>
using System;

[AttributeUsage(AttributeTargets.Property | AttributeTargets.Field, Inherited = true, AllowMultiple = false)]
internal class {Helpers.UiElementAttribute} : Attribute
{{
    public {Helpers.UiElementAttribute}(string name=null) {{ }}
}}
";

    #region Implementation of ISourceGenerator

        public void Initialize( GeneratorInitializationContext context )
        {
            context.RegisterForPostInitialization( i =>
            {
                i.AddSource( $"{Helpers.UiElementAttribute}_g.cs"
                           , SourceText.From( UiElementAttributeText , Encoding.UTF8 ) );
            } );

            context.RegisterForSyntaxNotifications( () => new SyntaxReceiver() );
        }

        public void Execute( GeneratorExecutionContext context )
        {
            if ( !(context.SyntaxContextReceiver is SyntaxReceiver receiver) )
                return;

            INamedTypeSymbol uiElementAttributeSymbol = context.Compilation.GetTypeByMetadataName( Helpers.UiElementAttribute );
            foreach ( IGrouping<INamedTypeSymbol , ISymbol> group in receiver.Fields
                                                                             .GroupBy<ISymbol , INamedTypeSymbol>( f => f.ContainingType
                                                                                                                 , SymbolEqualityComparer.Default ) )
            {
                var classSource = ProcessClass( group.Key , group , uiElementAttributeSymbol );
                if ( classSource == null )
                    continue;

                context.AddSource( $"{group.Key.Name}_ui_g.cs" , SourceText.From( classSource , Encoding.UTF8 ) );
            }
        }

        private string ProcessClass( INamedTypeSymbol     classSymbol
                                   , IEnumerable<ISymbol> fields
                                   , INamedTypeSymbol     uiElementAttributeSymbol )
        {
            var fieldsList = fields.ToList();
            if ( !fieldsList.Any() )
                return null;

            List<ISymbol> elementFields = fieldsList.Where( f => GetUiElementAttributeData( f , uiElementAttributeSymbol ) != null ).ToList();

            var source = new StringBuilder( $@"// <auto-generated/>

using UnityEngine.UIElements;
namespace {classSymbol.ContainingNamespace}
{{
public partial class {classSymbol.Name} 
{{" );

            source.Append( $@"
    protected override void QueryElements()
    {{
" );

            foreach ( ISymbol fieldSymbol in elementFields )
            {
                source.AppendLine( $"         {fieldSymbol.Name} = this.Q<{GetQualifyingTypeNameFromSymbol( fieldSymbol )}>(\"{GetUiElementAttributeData( fieldSymbol , uiElementAttributeSymbol )?.Name}\");" );
            }

            source.Append( $@"    }}
}}
}}
" );

            return source.ToString();
        }

        private static string GetQualifyingTypeName( ITypeSymbol type ) { return type.ToDisplayString( SymbolDisplayFormat.FullyQualifiedFormat ); }
        private string GetQualifyingTypeNameFromSymbol( ISymbol  symbol ) => GetQualifyingTypeName( GetTypeFromSymbol( symbol ) );

        private static UiElementAttributeData? GetUiElementAttributeData( ISymbol fieldSymbol , INamedTypeSymbol uiElementAttributeSymbol )
        {
            var attr = GetSingleAttributeData( fieldSymbol , uiElementAttributeSymbol );
            if ( attr == null )
                return null;

            var args = attr.ConstructorArguments.ToList();
            if ( args.Count > 1 )
            {
                throw new NotImplementedException( $"Attribute did not have enough parameters: expected 1 got {args.Count} {attr}: args: {args}" );
            }

            string name = null;
            if ( args.Count == 1 )
            {
                name = args[0].Value as string;
            }

            if ( name is null )
            {
                name = fieldSymbol.Name;
            }

            return new UiElementAttributeData()
            {
                Name = name ,
            };
        }

        private static AttributeData GetSingleAttributeData( ISymbol fieldSymbol , INamedTypeSymbol attributeSymbol )
        {
            var attr = fieldSymbol.GetAttributes()
                                  .SingleOrDefault( ad =>
                                                        ad?.AttributeClass?.Equals( attributeSymbol , SymbolEqualityComparer.Default ) ?? false );

            return attr;
        }


        private static ITypeSymbol GetTypeFromSymbol( ISymbol symbol )
        {
            switch ( symbol )
            {
                case IFieldSymbol fieldSymbol:
                    return fieldSymbol.Type;

                case IPropertySymbol propertySymbol:
                    return propertySymbol.Type;

                default:
                    throw new InvalidCastException( $"symbol was not property or field: {symbol}" );
            }
        }

        struct UiElementAttributeData
        {
            public string Name;
        }

    #endregion

    }

    public class SyntaxReceiver : ISyntaxContextReceiver
    {

        public List<ISymbol> Fields { get; } = new List<ISymbol>();

    #region Implementation of ISyntaxContextReceiver

        public void OnVisitSyntaxNode( GeneratorSyntaxContext context )
        {
            switch ( context.Node )
            {
                case FieldDeclarationSyntax fieldDeclarationSyntax when fieldDeclarationSyntax.AttributeLists.Count > 0:
                    {
                        foreach ( VariableDeclaratorSyntax variable in fieldDeclarationSyntax.Declaration.Variables )
                        {
                            if ( context.SemanticModel.GetDeclaredSymbol( variable ) is IFieldSymbol symbol
                              && Helpers.IsDerivedFrom( symbol.ContainingType.BaseType , "AtVisualElement" )
                              && symbol.GetAttributes()
                                       .Any( ad => ad.AttributeClass?.ToDisplayString() == Helpers.UiElementAttribute ) )
                            {
                                Fields.Add( symbol );
                            }
                        }
                        break;
                    }

                case PropertyDeclarationSyntax propertyDeclarationSyntax when propertyDeclarationSyntax.AttributeLists.Count > 0:
                    {
                        if ( context.SemanticModel.GetDeclaredSymbol( propertyDeclarationSyntax ) is IPropertySymbol symbol
                          && Helpers.IsDerivedFrom( symbol.ContainingType.BaseType , "AtVisualElement" )
                          && symbol.GetAttributes()
                                   .Any( ad => ad.AttributeClass?.ToDisplayString() == Helpers.UiElementAttribute ) )
                        {
                            Fields.Add( symbol );
                        }
                        break;
                    }
            }
        }

    #endregion

    }
}