Add support for multiple projects

Author: ejsmith

src/Foundatio.Mediator.SourceGenerator/DIRegistrationGenerator.cs

@@ -6,8 +6,11 @@ namespace Foundatio.Mediator;
 
 internal static class DIRegistrationGenerator
 {
-    public static void Execute(SourceProductionContext context, List<HandlerInfo> handlers)
+    public static void Execute(SourceProductionContext context, List<HandlerInfo> handlers, Compilation compilation)
     {
+        var assemblyName = compilation.AssemblyName?.ToIdentifier() ?? Guid.NewGuid().ToString("N").Substring(0, 10);
+        var className = $"{assemblyName}_MediatorHandlers";
+
         var source = new IndentedStringBuilder();
 
         source.AddGeneratedFileHeader();
@@ -19,18 +22,18 @@ public static void Execute(SourceProductionContext context, List<HandlerInfo> ha
         source.AppendLine("using System.Threading;");
         source.AppendLine("using System.Threading.Tasks;");
         source.AppendLine();
+        source.AppendLine("[assembly: Foundatio.Mediator.FoundatioHandlerModule]");
+        source.AppendLine();
         source.AppendLine("namespace Foundatio.Mediator;");
         source.AppendLine();
         source.AppendLine("[DebuggerStepThrough]");
         source.AppendLine("[DebuggerNonUserCode]");
         source.AppendLine("[ExcludeFromCodeCoverage]");
-        source.AppendLine("public static partial class ServiceCollectionExtensions");
+        source.AppendLine($"public static class {className}");
         source.AppendLine("{");
         source.AppendLine("    [DebuggerStepThrough]");
-        source.AppendLine("    public static IServiceCollection AddMediator(this IServiceCollection services)");
+        source.AppendLine("    public static void AddHandlers(this IServiceCollection services)");
         source.AppendLine("    {");
-        source.AppendLine("        services.AddSingleton<IMediator, Mediator>();");
-        source.AppendLine();
         source.AppendLine("        // Register HandlerRegistration instances keyed by message type name");
         source.AppendLine("        // Note: Handlers themselves are NOT auto-registered in DI");
         source.AppendLine("        // Users can register them manually if they want specific lifetimes");
@@ -59,13 +62,10 @@ public static void Execute(SourceProductionContext context, List<HandlerInfo> ha
             source.AppendLine($"        {handler.IsAsync.ToString().ToLower()}));");
         }
 
-        source.AppendLine();
-        source.AppendLine("return services;");
-
         source.DecrementIndent().DecrementIndent();
         source.AppendLine("    }");
         source.AppendLine("}");
 
-        context.AddSource("ServiceCollectionExtensions.g.cs", source.ToString());
+        context.AddSource($"{className}.g.cs", source.ToString());
     }
 }

src/Foundatio.Mediator.SourceGenerator/MediatorGenerator.cs

@@ -54,18 +54,20 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
             .Combine(middleware.Collect())
             .Combine(callSites.Collect())
             .Combine(interceptionEnabled)
+            .Combine(context.CompilationProvider)
             .Select(static (spc, _) => (
-                Handlers: spc.Left.Left.Left,
-                Middleware: spc.Left.Left.Right,
-                CallSites: spc.Left.Right,
-                InterceptorsEnabled: spc.Right
+                Handlers: spc.Left.Left.Left.Left,
+                Middleware: spc.Left.Left.Left.Right,
+                CallSites: spc.Left.Left.Right,
+                InterceptorsEnabled: spc.Left.Right,
+                Compilation: spc.Right
             ));
 
         context.RegisterImplementationSourceOutput(compilationAndData,
-            static (spc, source) => Execute(source.Handlers, source.Middleware, source.CallSites, source.InterceptorsEnabled, spc));
+            static (spc, source) => Execute(source.Handlers, source.Middleware, source.CallSites, source.InterceptorsEnabled, source.Compilation, spc));
     }
 
-    private static void Execute(ImmutableArray<HandlerInfo> handlers, ImmutableArray<MiddlewareInfo> middleware, ImmutableArray<CallSiteInfo> callSites, bool interceptorsEnabled, SourceProductionContext context)
+    private static void Execute(ImmutableArray<HandlerInfo> handlers, ImmutableArray<MiddlewareInfo> middleware, ImmutableArray<CallSiteInfo> callSites, bool interceptorsEnabled, Compilation compilation, SourceProductionContext context)
     {
         if (handlers.IsDefaultOrEmpty)
             return;
@@ -87,7 +89,7 @@ private static void Execute(ImmutableArray<HandlerInfo> handlers, ImmutableArray
 
         HandlerGenerator.Execute(context, handlersWithInfo, interceptorsEnabled);
 
-        DIRegistrationGenerator.Execute(context, handlersWithInfo);
+        DIRegistrationGenerator.Execute(context, handlersWithInfo, compilation);
     }
 
     private static EquatableArray<MiddlewareInfo> GetApplicableMiddlewares(ImmutableArray<MiddlewareInfo> middlewares, HandlerInfo handler)

src/Foundatio.Mediator/FoundatioHandlerModuleAttribute.cs

@@ -0,0 +1,4 @@
+namespace Foundatio.Mediator;
+
+[AttributeUsage(AttributeTargets.Assembly)]
+public sealed class FoundatioHandlerModuleAttribute : Attribute { }

src/Foundatio.Mediator/MediatorExtensions.cs

@@ -0,0 +1,43 @@
+using System.Reflection;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Foundatio.Mediator;
+
+public static class MediatorExtensions
+{
+    /// <summary>
+    /// Adds Foundatio.Mediator to the service collection.
+    /// </summary>
+    /// <param name="services"></param>
+    /// <returns>The updated service collection with Foundatio.Mediator registered.</returns>
+    public static IServiceCollection AddMediator(this IServiceCollection services)
+    {
+        services.AddSingleton<IMediator, Mediator>();
+
+        foreach (var assembly in AppDomain.CurrentDomain.GetAssemblies())
+        {
+            if (!IsAssemblyMarkedWithFoundatioHandlerModule(assembly))
+                continue;
+
+            var moduleType = assembly.GetTypes().FirstOrDefault(t =>
+                t.IsClass &&
+                t.IsAbstract &&
+                t.IsSealed &&
+                t.Name.EndsWith("_MediatorHandlers"));
+
+            var method = moduleType?.GetMethod("AddHandlers", BindingFlags.Public | BindingFlags.Static);
+
+            if (method != null)
+            {
+                method.Invoke(null, [services]);
+            }
+        }
+
+        return services;
+    }
+
+    private static bool IsAssemblyMarkedWithFoundatioHandlerModule(Assembly assembly)
+    {
+        return assembly.GetCustomAttributes(typeof(FoundatioHandlerModuleAttribute), false).Any();
+    }
+}