gaemstone/src/gaemstone.SourceGen/Generators/ModuleGenerator.cs

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

namespace gaemstone.SourceGen.Generators;

[Generator]
public class ModuleGenerator
	: ISourceGenerator
{
	private static readonly DiagnosticDescriptor ModuleMayNotBeNested = new(
		"gaem0001", "Module may not be nested",
		"Type {0} marked with [Module] may not be a nested type",
		nameof(ModuleGenerator), DiagnosticSeverity.Error, true);

	private static readonly DiagnosticDescriptor ModuleMustBePartial = new(
		"gaem0002", "Module must be partial",
		"Type {0} marked with [Module] must be a partial type",
		nameof(ModuleGenerator), DiagnosticSeverity.Error, true);

	private static readonly DiagnosticDescriptor ModuleBuiltInMustHavePath = new(
		"gaem0003", "Built-in module must have [Path]",
		"Type {0} marked with [Module] is a built-in module (static), and therefore must have [Path] set",
		nameof(ModuleGenerator), DiagnosticSeverity.Error, true);


	public void Initialize(GeneratorInitializationContext context)
		=> context.RegisterForSyntaxNotifications(() => new SyntaxReceiver());

	private class SyntaxReceiver
		: ISyntaxContextReceiver
	{
		public HashSet<INamedTypeSymbol> Symbols { get; }
			= new(SymbolEqualityComparer.Default);

		public void OnVisitSyntaxNode(GeneratorSyntaxContext context)
		{
			if (context.Node is not AttributeSyntax attrNode) return;
			var model = context.SemanticModel;

			var attrType = model.GetTypeInfo(attrNode).Type!;
			if (attrType.GetFullName(true) != "gaemstone.ECS.ModuleAttribute") return;

			var memberNode   = attrNode.Parent?.Parent!;
			var memberSymbol = model.GetDeclaredSymbol(memberNode) as INamedTypeSymbol;
			Symbols.Add(memberSymbol!);
		}
	}

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

		foreach (var symbol in receiver.Symbols) {
			var isNested = (symbol.ContainingType != null);
			if (isNested)
				context.ReportDiagnostic(Diagnostic.Create(ModuleMayNotBeNested,
					symbol.Locations.FirstOrDefault(), symbol.GetFullName()));

			var isPartial = symbol.DeclaringSyntaxReferences
				.Any(r => (r.GetSyntax() as ClassDeclarationSyntax)?.Modifiers
					.Any(t => t.IsKind(SyntaxKind.PartialKeyword)) ?? false);
			if (!isPartial)
				context.ReportDiagnostic(Diagnostic.Create(ModuleMustBePartial,
					symbol.Locations.FirstOrDefault(), symbol.GetFullName()));

			if (symbol.IsStatic && (symbol.GetAttribute("gaemstone.ECS.PathAttribute")?
			                              .ConstructorArguments.FirstOrDefault().Value == null))
				context.ReportDiagnostic(Diagnostic.Create(ModuleBuiltInMustHavePath,
					symbol.Locations.FirstOrDefault(), symbol.GetFullName()));

			// Static classes can't implement interfaces.
			if (symbol.IsStatic) continue;

			var modulePath = GetModulePath(symbol).ToStringLiteral();
			var dependencies = new List<string>();
			foreach (var attr in symbol.GetAttributes())
				for (var type = attr.AttributeClass; type != null; type = type.BaseType)
					if ((type.GetFullName(true) == "gaemstone.ECS.AddAttribute`2")
							&& type.TypeArguments[0].GetFullName() == "gaemstone.Flecs.Core.DependsOn")
						dependencies.Add(GetModulePath(type.TypeArguments[1]));

			var sb = new StringBuilder();
			sb.AppendLine($$"""
				// <auto-generated/>
				using System.Collections.Generic;
				using System.Linq;
				using gaemstone.ECS;

				namespace {{ symbol.GetNamespace() }};

				public partial class {{ symbol.Name }}
					: IModule
				{
					public static string ModulePath { get; }
						= {{ modulePath }};

				""");

			if (dependencies.Count > 0) {
				sb.AppendLine($$"""
						public static IEnumerable<string> Dependencies { get; } = new[] {
					""");
				foreach (var dependency in dependencies)
					sb.AppendLine($$"""		{{ dependency.ToStringLiteral() }},""");
				sb.AppendLine($$"""
						};
					""");
			} else sb.AppendLine($$"""
					public static IEnumerable<string> Dependencies { get; } = Enumerable.Empty<string>();
				""");


			sb.AppendLine($$"""
				}

				""");

			context.AddSource($"{symbol.GetFullName()}_Module.g.cs", sb.ToString());
		}
	}

	private static string GetModulePath(ISymbol module)
	{
		var path = module.GetAttribute("gaemstone.ECS.PathAttribute")?
							.ConstructorArguments.FirstOrDefault().Value as string;
		var isAbsolute = (path?.FirstOrDefault() == '/');
		if (isAbsolute) return path!;

		var fullPath = module.GetFullName().Replace('.', '/');
		if (path != null) {
			var index = fullPath.LastIndexOf('/');
			fullPath = $"{fullPath.Substring(0, index)}/{path}";
		}
		return $"/{fullPath}";
	}
}