gaemstone/src/gaemstone.SourceGen/RelevantSymbolReceiver.cs

using System;
using System.Collections.Generic;
using System.Linq;
using gaemstone.SourceGen.Structure;
using gaemstone.SourceGen.Utility;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace gaemstone.SourceGen;

public class RelevantSymbolReceiver
	: ISyntaxContextReceiver
{
	private static readonly HashSet<string> RelevantAttributeNames = new(){
		// Base entity attributes
		"Module",     // Can also be [Singleton]
		"Entity",
		"Relation",   // Can also be [Tag] or [Component]
		"Tag",
		"Component",
		"Singleton",  // Implies [Component]
		"System",
		"Observer",

		// Entity properties that specify additional info / behavior
		"Public",
		"Private",
		"Path",
		"Add",
		"BuiltIn",    // Valid on [Module]
		"Expression", // Valid on [System] and [Observer]

		// Term properties (on [System] and [Observer] parameters)
		"Source",
		"Pair",
		"Has",
		"Not",
		"Or",
	};

	public Dictionary<ISymbol, BaseInfo> Symbols { get; } = new(SymbolEqualityComparer.Default);

	public void OnVisitSyntaxNode(GeneratorSyntaxContext context)
	{
		var model = context.SemanticModel;
		if (context.Node is not AttributeSyntax node) return;
		if (node.Parent?.Parent is not SyntaxNode parent) return;
		if (model.GetDeclaredSymbol(parent) is not ISymbol symbol) return;
		if (model.GetTypeInfo(node).Type is not INamedTypeSymbol type) return;

		// Go through the attribute's type hierarchy to see if it matches any
		// of the attributes we care about. This is to make sure attributes
		// based on [Add<...>] are picked up correctly, including custom ones.
		for (var baseType = type; baseType != null; baseType = baseType.BaseType) {
			if ((ToRelevantAttributeName(baseType) is string name)
				&& RelevantAttributeNames.Contains(name) // Check if we found a relevant attribute.
				&& !Symbols.ContainsKey(symbol))         // Check if this is already a known symbol.
			{
				Symbols.Add(symbol, symbol switch {
					INamedTypeSymbol typeSymbol =>
						typeSymbol.GetAttributes().Any(attr => attr.AttributeClass!
								.GetFullName() == "gaemstone.ECS.ModuleAttribute")
							? new ModuleEntityInfo(typeSymbol)
							: new TypeEntityInfo(typeSymbol),
					IMethodSymbol methodSymbol => new MethodEntityInfo(methodSymbol),
					IParameterSymbol paramSymbol => new ParameterInfo(paramSymbol),
					_ => throw new InvalidOperationException(
						$"Unhandled symbol type {symbol.GetType()}"),
				});
				break;
			}
		}
	}

	public static string? ToRelevantAttributeName(INamedTypeSymbol symbol)
	{
		if (symbol.GetNamespace() != "gaemstone.ECS") return null;
		var name = symbol.MetadataName.Split('`')[0];
		return name.EndsWith("Attribute") ? name[..^"Attribute".Length] : null;
	}
}