diff --git a/Arch.System.SourceGenerator.SnapshotTests/ConcurrencyTest.cs b/Arch.System.SourceGenerator.SnapshotTests/ConcurrencyTest.cs
new file mode 100644
index 0000000..ec4e1c1
--- /dev/null
+++ b/Arch.System.SourceGenerator.SnapshotTests/ConcurrencyTest.cs
@@ -0,0 +1,128 @@
+using System.Collections.Immutable;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp;
+using NUnit.Framework;
+
+namespace Arch.System.SourceGenerator.Tests;
+
+///
+/// Tests that is thread-safe when invoked concurrently,
+/// as happens during parallel builds with multiple projects referencing the same generator.
+///
+[TestFixture]
+internal class ConcurrencyTest
+{
+ private const string SystemSourceA = """
+ using Arch.Core;
+ using Arch.System;
+
+ namespace TestA;
+
+ public partial class SystemA : BaseSystem
+ {
+ public SystemA(World world) : base(world) { }
+
+ [Query]
+ public void MethodA(ref int component) { component++; }
+
+ [Query]
+ public void MethodB(ref int component) { component += 2; }
+ }
+ """;
+
+ private const string SystemSourceB = """
+ using Arch.Core;
+ using Arch.System;
+
+ namespace TestB;
+
+ public partial class SystemB : BaseSystem
+ {
+ public SystemB(World world) : base(world) { }
+
+ [Query]
+ public void MethodC(ref int component) { component += 10; }
+
+ [Query]
+ public void MethodD(ref int component) { component += 20; }
+ }
+ """;
+
+ private static CSharpCompilation CreateCompilation(string source, string assemblyName)
+ {
+ var syntaxTree = CSharpSyntaxTree.ParseText(source, CSharpParseOptions.Default.WithLanguageVersion(LanguageVersion.Latest));
+
+ var baseAssemblyPath = Path.GetDirectoryName(typeof(object).Assembly.Location)!;
+ var systemAssemblies = Directory.GetFiles(baseAssemblyPath)
+ .Where(x => !x.EndsWith("Native.dll"))
+ .Where(x =>
+ {
+ var filename = Path.GetFileName(x);
+ return filename.StartsWith("System") || filename is "mscorlib.dll" or "netstandard.dll";
+ });
+
+ var references = new List
+ {
+ typeof(Arch.Core.World).Assembly.Location,
+ typeof(QueryAttribute).Assembly.Location,
+ typeof(CommunityToolkit.HighPerformance.ArrayExtensions).Assembly.Location
+ };
+ references.AddRange(systemAssemblies);
+
+ return CSharpCompilation.Create(assemblyName, new[] { syntaxTree },
+ references.Select(r => MetadataReference.CreateFromFile(r)),
+ new CSharpCompilationOptions(OutputKind.DynamicallyLinkedLibrary));
+ }
+
+ private static (ImmutableArray diagnostics, ImmutableArray errors) RunGenerator(CSharpCompilation compilation)
+ {
+ GeneratorDriver driver = CSharpGeneratorDriver.Create(new QueryGenerator());
+ driver = driver.RunGeneratorsAndUpdateCompilation(compilation, out var outputCompilation, out var diagnostics);
+
+ var errors = outputCompilation.GetDiagnostics()
+ .Where(d => d.Severity == DiagnosticSeverity.Error)
+ .ToImmutableArray();
+
+ return (diagnostics, errors);
+ }
+
+ ///
+ /// Runs multiple QueryGenerator instances concurrently to reproduce the race condition
+ /// caused by the static _classToMethods field.
+ ///
+ [Test]
+ [Repeat(10)]
+ public void ConcurrentGeneratorRuns_ShouldNotThrow()
+ {
+ var compilationA = CreateCompilation(SystemSourceA, "TestAssemblyA");
+ var compilationB = CreateCompilation(SystemSourceB, "TestAssemblyB");
+
+ var exceptions = new global::System.Collections.Concurrent.ConcurrentBag();
+
+ Parallel.For(0, 8, i =>
+ {
+ try
+ {
+ var compilation = i % 2 == 0 ? compilationA : compilationB;
+ var (diagnostics, errors) = RunGenerator(compilation);
+
+ // Each run should produce zero generator diagnostics and zero compilation errors
+ Assert.That(diagnostics, Is.Empty,
+ $"Iteration {i}: Generator diagnostics should be empty.\n{string.Join("\n", diagnostics)}");
+ Assert.That(errors, Is.Empty,
+ $"Iteration {i}: Output compilation should have no errors.\n{string.Join("\n", errors)}");
+ }
+ catch (Exception ex)
+ {
+ exceptions.Add(ex);
+ }
+ });
+
+ if (!exceptions.IsEmpty)
+ {
+ Assert.Fail(
+ $"Concurrent generator runs failed with {exceptions.Count} exception(s):\n" +
+ string.Join("\n---\n", exceptions.Select(e => e.ToString())));
+ }
+ }
+}
diff --git a/Arch.System.SourceGenerator/SourceGenerator.cs b/Arch.System.SourceGenerator/SourceGenerator.cs
index 2e4afc7..3c3669d 100644
--- a/Arch.System.SourceGenerator/SourceGenerator.cs
+++ b/Arch.System.SourceGenerator/SourceGenerator.cs
@@ -14,8 +14,6 @@ namespace Arch.System.SourceGenerator;
[Generator]
public class QueryGenerator : IIncrementalGenerator
{
- private static Dictionary> _classToMethods { get; set; } = null!;
-
///
public void Initialize(IncrementalGeneratorInitializationContext context)
{
@@ -34,15 +32,15 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
///
/// Adds a to its class.
- /// Stores them in .
///
+ /// The dictionary mapping classes to their methods.
/// The which will be added/mapped to its class.
- private static void AddMethodToClass(IMethodSymbol methodSymbol)
+ private static void AddMethodToClass(Dictionary> classToMethods, IMethodSymbol methodSymbol)
{
- if (!_classToMethods.TryGetValue(methodSymbol.ContainingSymbol, out var list))
+ if (!classToMethods.TryGetValue(methodSymbol.ContainingSymbol, out var list))
{
list = new List();
- _classToMethods[methodSymbol.ContainingSymbol] = list;
+ classToMethods[methodSymbol.ContainingSymbol] = list;
}
list.Add(methodSymbol);
}
@@ -89,7 +87,7 @@ private static void Generate(Compilation compilation, ImmutableArray>(512, SymbolEqualityComparer.Default);
foreach (var methodSyntax in methods)
{
IMethodSymbol? methodSymbol = null;
@@ -104,8 +102,8 @@ private static void Generate(Compilation compilation, ImmutableArray', '}');
@@ -113,7 +111,7 @@ private static void Generate(Compilation compilation, ImmutableArray