using Fantasy.SourceGenerator.Common;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using System.Collections.Generic;
using System.Linq;
using System.Text;
#pragma warning disable CS8602 // Dereference of a possibly null reference.
namespace Fantasy.SourceGenerator.Generators
{
///
/// Entity System 注册代码生成器
/// 自动生成 EntityComponent 所需的 System 注册代码,替代运行时反射
///
[Generator]
public partial class EntitySystemGenerator : IIncrementalGenerator
{
public void Initialize(IncrementalGeneratorInitializationContext context)
{
// 查找所有实现了 IEntitySystem 相关接口的类
var systemTypes = context.SyntaxProvider
.CreateSyntaxProvider(
predicate: static (node, _) => IsSystemClass(node),
transform: static (ctx, _) => GetSystemTypeInfo(ctx))
.Where(static info => info != null)
.Collect();
// 组合编译信息和找到的类型
var compilationAndTypes = context.CompilationProvider.Combine(systemTypes);
// 注册源代码输出
context.RegisterSourceOutput(compilationAndTypes, static (spc, source) =>
{
// 检查1: 是否定义了 FANTASY_NET 或 FANTASY_UNITY 预编译符号
if (!CompilationHelper.HasFantasyDefine(source.Left))
{
return;
}
// 检查2: 是否引用了 Fantasy 框架的核心类型
if (source.Left.GetTypeByMetadataName("Fantasy.Assembly.IEntitySystemRegistrar") == null)
{
return;
}
GenerateRegistrationCode(spc, source.Left, source.Right!);
});
}
private static EntitySystemTypeInfo? GetSystemTypeInfo(GeneratorSyntaxContext context)
{
var classDecl = (ClassDeclarationSyntax)context.Node;
if (context.SemanticModel.GetDeclaredSymbol(classDecl) is not INamedTypeSymbol symbol || !symbol.IsInstantiable())
{
return null;
}
var baseType = symbol.BaseType;
if (!baseType.IsGenericType)
{
return null;
}
return baseType.Name switch
{
"AwakeSystem" => EntitySystemTypeInfo.Create(EntitySystemType.AwakeSystem, symbol),
"UpdateSystem" => EntitySystemTypeInfo.Create(EntitySystemType.UpdateSystem, symbol),
"DestroySystem" => EntitySystemTypeInfo.Create(EntitySystemType.DestroySystem, symbol),
"DeserializeSystem" => EntitySystemTypeInfo.Create(EntitySystemType.DeserializeSystem, symbol),
"LateUpdateSystem" => EntitySystemTypeInfo.Create(EntitySystemType.LateUpdateSystem, symbol),
_ => null
};
}
private static void GenerateRegistrationCode(
SourceProductionContext context,
Compilation compilation,
IEnumerable systemTypes)
{
var entitySystemTypeInfos = systemTypes.ToList();
// 获取当前程序集名称(仅用于注释)
var assemblyName = compilation.AssemblyName ?? "Unknown";
// 生成代码文件
var builder = new SourceCodeBuilder();
// 添加文件头
builder.AppendLine(GeneratorConstants.AutoGeneratedHeader);
// 添加 using
builder.AddUsings(
"System",
"System.Collections.Generic",
"Fantasy.Assembly",
"Fantasy.Entitas",
"Fantasy.Entitas.Interface"
);
builder.AppendLine();
// 开始命名空间(固定使用 Fantasy.Generated)
builder.BeginNamespace("Fantasy.Generated");
// 开始类定义(实现 IEntitySystemRegistrar 接口)
builder.AddXmlComment($"Auto-generated Entity System registration class for {assemblyName}");
builder.BeginClass("EntitySystemRegistrar", "internal sealed", "IEntitySystemRegistrar");
// 生成字段用于存储 System 实例
GenerateFields(builder, entitySystemTypeInfos);
// 生成注册方法
GenerateRegisterMethod(builder, entitySystemTypeInfos);
// 生成反注册方法
GenerateUnRegisterMethod(builder, entitySystemTypeInfos);
// 结束类和命名空间
builder.EndClass();
builder.EndNamespace();
// 输出源代码
context.AddSource("EntitySystemRegistrar.g.cs", builder.ToString());
}
private static void GenerateFields(SourceCodeBuilder builder, List entitySystemTypeInfos)
{
builder.AddComment("Store registered entity system for UnRegister");
if (!entitySystemTypeInfos.Any())
{
return;
}
foreach (var eventSystemTypeInfo in entitySystemTypeInfos)
{
var fieldName = $"_{eventSystemTypeInfo.TypeName.ToCamelCase()}";
builder.AppendLine($"private {eventSystemTypeInfo.GlobalTypeFullName} {fieldName} = new {eventSystemTypeInfo.GlobalTypeFullName}();");
builder.AppendLine($"private const long _typeHashCode{fieldName} = {eventSystemTypeInfo.EntityTypeHashCode};");
}
builder.AppendLine();
}
private static void GenerateRegisterMethod(SourceCodeBuilder builder, List entitySystemTypeInfos)
{
builder.AddXmlComment("Register all Entity Systems to the dictionaries");
builder.AppendLine("#if FANTASY_NET", false);
builder.BeginMethod(
"public void RegisterSystems(" +
"Dictionary> awakeSystems, " +
"Dictionary> updateSystems, " +
"Dictionary> destroySystems, " +
"Dictionary> deserializeSystems)");
builder.AppendLine("#endif", false);
builder.Unindent();
builder.AppendLine("#if FANTASY_UNITY", false);
builder.BeginMethod(
"public void RegisterSystems(" +
"Dictionary> awakeSystems, " +
"Dictionary> updateSystems, " +
"Dictionary> destroySystems, " +
"Dictionary> deserializeSystems, " +
"Dictionary> lateUpdateSystems)");
builder.AppendLine("#endif", false);
if (entitySystemTypeInfos.Any())
{
foreach (var systemTypeInfo in entitySystemTypeInfos)
{
var fieldName = $"_{systemTypeInfo.TypeName.ToCamelCase()}";
switch (systemTypeInfo.EntitySystemType)
{
case EntitySystemType.AwakeSystem:
{
builder.AppendLine($"awakeSystems.Add(_typeHashCode{fieldName}, {fieldName}.Invoke);");
continue;
}
case EntitySystemType.UpdateSystem:
{
builder.AppendLine($"updateSystems.Add(_typeHashCode{fieldName}, {fieldName}.Invoke);");
continue;
}
case EntitySystemType.DestroySystem:
{
builder.AppendLine($"destroySystems.Add(_typeHashCode{fieldName}, {fieldName}.Invoke);");
continue;
}
case EntitySystemType.DeserializeSystem:
{
builder.AppendLine($"deserializeSystems.Add(_typeHashCode{fieldName}, {fieldName}.Invoke);");
continue;
}
case EntitySystemType.LateUpdateSystem:
{
builder.AppendLine("#if FANTASY_UNITY", false);
builder.AppendLine($"awakeSystems.Add(_typeHashCode{fieldName}, {fieldName}.Invoke);");
builder.AppendLine("#endif", false);
continue;
}
}
}
}
builder.EndMethod();
builder.AppendLine();
}
private static void GenerateUnRegisterMethod(SourceCodeBuilder builder, List entitySystemTypeInfos)
{
builder.AddXmlComment("Unregister all Entity Systems from the dictionaries");
builder.AppendLine("#if FANTASY_NET", false);
builder.BeginMethod(
"public void UnRegisterSystems(" +
"Dictionary> awakeSystems, " +
"Dictionary> updateSystems, " +
"Dictionary> destroySystems, " +
"Dictionary> deserializeSystems)");
builder.AppendLine("#endif", false);
builder.Unindent();
builder.AppendLine("#if FANTASY_UNITY", false);
builder.BeginMethod(
"public void UnRegisterSystems(" +
"Dictionary> awakeSystems, " +
"Dictionary> updateSystems, " +
"Dictionary> destroySystems, " +
"Dictionary> deserializeSystems, " +
"Dictionary> lateUpdateSystems)");
builder.AppendLine("#endif", false);
if (entitySystemTypeInfos.Any())
{
foreach (var systemTypeInfo in entitySystemTypeInfos)
{
var fieldName = $"_{systemTypeInfo.TypeName.ToCamelCase()}";
switch (systemTypeInfo.EntitySystemType)
{
case EntitySystemType.AwakeSystem:
{
builder.AppendLine($"awakeSystems.Remove(_typeHashCode{fieldName});");
continue;
}
case EntitySystemType.UpdateSystem:
{
builder.AppendLine($"updateSystems.Remove(_typeHashCode{fieldName});");
continue;
}
case EntitySystemType.DestroySystem:
{
builder.AppendLine($"destroySystems.Remove(_typeHashCode{fieldName});");
continue;
}
case EntitySystemType.DeserializeSystem:
{
builder.AppendLine($"deserializeSystems.Remove(_typeHashCode{fieldName});");
continue;
}
case EntitySystemType.LateUpdateSystem:
{
builder.AppendLine("#if FANTASY_UNITY", false);
builder.AppendLine($"lateUpdateSystem.Remove(_typeHashCode{fieldName});");
builder.AppendLine("#endif", false);
continue;
}
}
}
}
builder.EndMethod();
builder.AppendLine();
}
private static bool IsSystemClass(SyntaxNode node)
{
if (node is not ClassDeclarationSyntax classDecl)
{
return false;
}
// 必须有基类型列表(继承抽象类)
if (classDecl.BaseList == null || classDecl.BaseList.Types.Count == 0)
{
return false;
}
// 快速检查是否包含可能的 EntitySystem 基类名称
var baseListText = classDecl.BaseList.ToString();
return baseListText.Contains("AwakeSystem") ||
baseListText.Contains("UpdateSystem") ||
baseListText.Contains("DestroySystem") ||
baseListText.Contains("DeserializeSystem") ||
baseListText.Contains("LateUpdateSystem");
}
private enum EntitySystemType
{
None,
AwakeSystem,
UpdateSystem,
DestroySystem,
DeserializeSystem,
LateUpdateSystem
}
private sealed class EntitySystemTypeInfo
{
public readonly EntitySystemType EntitySystemType;
public readonly string GlobalTypeFullName;
public readonly string TypeFullName;
public readonly string TypeName;
public readonly long EntityTypeHashCode; // 预计算的实体类型哈希值
public readonly string EntityTypeFullName; // 实体类型的完整名称(用于注释)
private EntitySystemTypeInfo(
EntitySystemType entitySystemType,
string globalTypeFullName,
string typeFullName,
string typeName,
long entityTypeHashCode,
string entityTypeFullName)
{
EntitySystemType = entitySystemType;
GlobalTypeFullName = globalTypeFullName;
TypeFullName = typeFullName;
TypeName = typeName;
EntityTypeHashCode = entityTypeHashCode;
EntityTypeFullName = entityTypeFullName;
}
public static EntitySystemTypeInfo Create(EntitySystemType entitySystemType, INamedTypeSymbol symbol)
{
// 获取泛型参数 T (例如:AwakeSystem 中的 T)
var entityType = symbol.BaseType?.TypeArguments.FirstOrDefault();
var entityTypeFullName = entityType?.GetFullName(false) ?? "Unknown";
// 使用与运行时相同的算法计算哈希值
var entityTypeHashCode = HashCodeHelper.ComputeHash64(entityTypeFullName);
return new EntitySystemTypeInfo(
entitySystemType,
symbol.GetFullName(),
symbol.GetFullName(false),
symbol.Name,
entityTypeHashCode,
entityTypeFullName);
}
}
}
}