using System.Collections.Generic; using System.Linq; using System.Text; using Microsoft.CodeAnalysis; using Microsoft.CodeAnalysis.CSharp.Syntax; using Microsoft.CodeAnalysis.Text; namespace GetComponentGenerator { [Generator] public class GetComponentGenerator : ISourceGenerator { private const string AttributeText = @" using System; [AttributeUsage(AttributeTargets.Field, Inherited = true, AllowMultiple = false)] internal class GetComponentAttribute : Attribute { public enum TargetType { This = 0, Parent = 1, Child = 2, } public GetComponentAttribute(TargetType targetType = TargetType.This) { } } "; public void Initialize( GeneratorInitializationContext context ) { context.RegisterForPostInitialization( i => i.AddSource( "GetComponentAttribute_g.cs" , AttributeText ) ); context.RegisterForSyntaxNotifications( () => new SyntaxReceiver() ); } public void Execute( GeneratorExecutionContext context ) { if ( !(context.SyntaxContextReceiver is SyntaxReceiver receiver) ) return; INamedTypeSymbol attributeSymbol = context.Compilation.GetTypeByMetadataName( "GetComponentAttribute" ); foreach ( IGrouping group in receiver.Fields .GroupBy( f => f.ContainingType , SymbolEqualityComparer.Default ) ) { var classSource = ProcessClass( group.Key , group , attributeSymbol ); context.AddSource( $"{group.Key.Name}_Components_g.cs" , SourceText.From( classSource , Encoding.UTF8 ) ); } } private string ProcessClass( INamedTypeSymbol classSymbol , IEnumerable fields , ISymbol attributeSymbol ) { var source = new StringBuilder( $@" public partial class {classSymbol.Name} {{ private void t() {{ " ); foreach ( IFieldSymbol fieldSymbol in fields ) { ProcessField( source , fieldSymbol , attributeSymbol ); } source.Append( "}\n\n}" ); return source.ToString(); } private void ProcessField( StringBuilder source , IFieldSymbol fieldSymbol , ISymbol attributeSymbol ) { var fieldName = fieldSymbol.Name; ITypeSymbol fieldType = fieldSymbol.Type; AttributeData attributeData = fieldSymbol.GetAttributes() .Single( ad => ad.AttributeClass?.Equals( attributeSymbol , SymbolEqualityComparer.Default ) ?? false ); var methodType = ProcessAttribute( attributeData ); source.AppendLine( $@"{fieldName} = {methodType}<{fieldType}>();" ); } private string ProcessAttribute( AttributeData attributeData ) { var stringBuilder = new StringBuilder( "GetComponent" ); var args = attributeData.ConstructorArguments; if ( args.Length > 0 && int.TryParse( args[0].Value?.ToString() , out var enumValue ) ) { switch ( enumValue ) { case 1: stringBuilder.Append( "InParent" ); break; case 2: stringBuilder.Append( "InChildren" ); break; } } return stringBuilder.ToString(); } } internal class SyntaxReceiver : ISyntaxContextReceiver { public List Fields { get; } = new List(); public void OnVisitSyntaxNode( GeneratorSyntaxContext context ) { if ( !(context.Node is FieldDeclarationSyntax fieldDeclarationSyntax) || fieldDeclarationSyntax.AttributeLists.Count <= 0 ) return; foreach ( VariableDeclaratorSyntax variable in fieldDeclarationSyntax.Declaration.Variables ) { if ( context.SemanticModel.GetDeclaredSymbol( variable ) is IFieldSymbol fieldSymbol && IsDerivedFrom( fieldSymbol.ContainingType.BaseType , "MonoBehaviour" ) && IsDerivedFrom( fieldSymbol.Type.BaseType , "Component" ) && fieldSymbol.GetAttributes().Any( ad => ad.AttributeClass?.ToDisplayString() == "GetComponentAttribute" ) ) { Fields.Add( fieldSymbol ); } } } private static bool IsDerivedFrom( INamedTypeSymbol baseType , string targetType ) { while ( baseType != null ) { if ( baseType.Name == targetType ) return true; baseType = baseType.BaseType; } return false; } } }