Add support for keyed services

README.md

@@ -84,3 +84,4 @@ private static partial IServiceCollection AddRepositories(this IServiceCollectio
 | **AsImplementedInterfaces** | If true, the registered types will be registered as implemented interfaces instead of their actual type. |
 | **AsSelf** | If true, types will be registered with their actual type. It can be combined with `AsImplementedInterfaces`. In that case implemeted interfaces will be "forwarded" to an actual implementation type |
 | **TypeNameFilter** | Set this value to filter the types to register by their full name. You can use '*' wildcards. You can also use ',' to separate multiple filters. |
+| **KeySelector** | Set this value to a static method name returning string. Returned value will be used as a key for the registration. Method should either be generic, or have a single parameter of type `Type`. |

ServiceScan.SourceGenerator.Tests/AddServicesTests.cs

@@ -462,6 +462,68 @@ public class MyService: IServiceA, IServiceB {}
         Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
     }
 
+    [Fact]
+    public void AddAsKeyedServices_GenericMethod()
+    {
+        var attribute = @"
+            private static string GetName<T>() => typeof(T).Name.Replace(""Service"", """");
+
+            [GenerateServiceRegistrations(AssignableTo = typeof(IService), KeySelector = nameof(GetName))]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public interface IService { }
+            public class MyService1 : IService { }
+            public class MyService2 : IService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var registrations = $"""
+            return services
+                .AddKeyedTransient<GeneratorTests.IService, GeneratorTests.MyService1>(GetName<GeneratorTests.MyService1>())
+                .AddKeyedTransient<GeneratorTests.IService, GeneratorTests.MyService2>(GetName<GeneratorTests.MyService2>());
+            """;
+        Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
+    }
+
+    [Fact]
+    public void AddAsKeyedServices_MethodWithTypeParameter()
+    {
+        var attribute = @"
+            private static string GetName(Type type) => type.Name.Replace(""Service"", """");
+
+            [GenerateServiceRegistrations(AssignableTo = typeof(IService), KeySelector = nameof(GetName))]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public interface IService { }
+            public class MyService1 : IService { }
+            public class MyService2 : IService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var registrations = $"""
+            return services
+                .AddKeyedTransient<GeneratorTests.IService, GeneratorTests.MyService1>(GetName(typeof(GeneratorTests.MyService1)))
+                .AddKeyedTransient<GeneratorTests.IService, GeneratorTests.MyService2>(GetName(typeof(GeneratorTests.MyService2)));
+            """;
+        Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
+    }
+
     [Fact]
     public void DontGenerateAnythingIfTypeIsInvalid()
     {

ServiceScan.SourceGenerator.Tests/DiagnosticTests.cs

@@ -201,4 +201,111 @@ public static partial class ServicesExtensions
 
         Assert.Equal(results.Diagnostics.Single().Descriptor, DiagnosticDescriptors.MissingSearchCriteria);
     }
+
+    [Fact]
+    public void KeySelectorMethodDoesNotExist()
+    {
+        var attribute = @"
+            private static string GetName<T>() => typeof(T).Name.Replace(""Service"", """");
+
+            [GenerateServiceRegistrations(AssignableTo = typeof(IService), KeySelector = ""NoSuchMethodHere"")]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public interface IService { }
+            public class MyService1 : IService { }
+            public class MyService2 : IService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        Assert.Equal(results.Diagnostics.Single().Descriptor, DiagnosticDescriptors.KeySelectorMethodNotFound);
+    }
+
+    [Fact]
+    public void KeySelectorMethod_GenericButHasParameters()
+    {
+        var attribute = @"
+            private static string GetName<T>(string name) => typeof(T).Name.Replace(""Service"", name);
+
+            [GenerateServiceRegistrations(AssignableTo = typeof(IService), KeySelector = nameof(GetName))]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public interface IService { }
+            public class MyService1 : IService { }
+            public class MyService2 : IService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        Assert.Equal(results.Diagnostics.Single().Descriptor, DiagnosticDescriptors.KeySelectorMethodHasIncorrectSignature);
+    }
+
+    [Fact]
+    public void KeySelectorMethod_NonGenericWithoutParameters()
+    {
+        var attribute = @"
+            private static string GetName() => ""const"";
+
+            [GenerateServiceRegistrations(AssignableTo = typeof(IService), KeySelector = nameof(GetName))]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public interface IService { }
+            public class MyService1 : IService { }
+            public class MyService2 : IService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        Assert.Equal(results.Diagnostics.Single().Descriptor, DiagnosticDescriptors.KeySelectorMethodHasIncorrectSignature);
+    }
+
+    [Fact]
+    public void KeySelectorMethod_Void()
+    {
+        var attribute = @"
+            private static void GetName(Type type)
+            {
+                type.Name.ToString();
+            }
+
+            [GenerateServiceRegistrations(AssignableTo = typeof(IService), KeySelector = nameof(GetName))]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public interface IService { }
+            public class MyService1 : IService { }
+            public class MyService2 : IService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        Assert.Equal(results.Diagnostics.Single().Descriptor, DiagnosticDescriptors.KeySelectorMethodHasIncorrectSignature);
+    }
 }

ServiceScan.SourceGenerator/DependencyInjectionGenerator.cs

@@ -1,4 +1,5 @@
-using System.Collections.Generic;
+using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Text;
 using System.Text.RegularExpressions;
@@ -51,9 +52,23 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                     else
                     {
                         if (registration.ResolveImplementation)
+                        {
                             sb.AppendLine($"            .Add{registration.Lifetime}<{registration.ServiceTypeName}>(s => s.GetRequiredService<{registration.ImplementationTypeName}>())");
+                        }
                         else
-                            sb.AppendLine($"            .Add{registration.Lifetime}<{registration.ServiceTypeName}, {registration.ImplementationTypeName}>()");
+                        {
+                            var addMethod = registration.KeySelectorMethodName != null
+                                ? $"AddKeyed{registration.Lifetime}"
+                                : $"Add{registration.Lifetime}";
+
+                            var keyMethodInvocation = registration.KeySelectorMethodGeneric switch
+                            {
+                                true => $"{registration.KeySelectorMethodName}<{registration.ImplementationTypeName}>()",
+                                false => $"{registration.KeySelectorMethodName}(typeof({registration.ImplementationTypeName}))",
+                                null => null
+                            };
+                            sb.AppendLine($"            .{addMethod}<{registration.ServiceTypeName}, {registration.ImplementationTypeName}>({keyMethodInvocation})");
+                        }
                     }
                 }
 
@@ -96,12 +111,36 @@ private static DiagnosticModel<MethodImplementationModel> FindServicesToRegister
         {
             bool typesFound = false;
 
-            var assembly = compilation.GetTypeByMetadataName(attribute.AssemblyOfTypeName ?? method.TypeMetadataName).ContainingAssembly;
+            var containingType = compilation.GetTypeByMetadataName(method.TypeMetadataName);
+
+            var assembly = (attribute.AssemblyOfTypeName is null
+                ? containingType
+                : compilation.GetTypeByMetadataName(attribute.AssemblyOfTypeName)).ContainingAssembly;
 
             var assignableToType = attribute.AssignableToTypeName is null
                 ? null
                 : compilation.GetTypeByMetadataName(attribute.AssignableToTypeName);
 
+            var keySelectorMethod = attribute.KeySelector is null
+                ? null
+                : containingType.GetMembers().OfType<IMethodSymbol>().FirstOrDefault(m =>
+                    m.IsStatic && m.Name == attribute.KeySelector);
+
+            if (attribute.KeySelector != null)
+            {
+                if (keySelectorMethod is null)
+                    return Diagnostic.Create(KeySelectorMethodNotFound, attribute.Location);
+
+                if (keySelectorMethod.ReturnsVoid)
+                    return Diagnostic.Create(KeySelectorMethodHasIncorrectSignature, attribute.Location);
+
+                var validGenericKeySelector = keySelectorMethod.TypeArguments.Length == 1 && keySelectorMethod.Parameters.Length == 0;
+                var validNonGenericKeySelector = !keySelectorMethod.IsGenericMethod && keySelectorMethod.Parameters is [{ Type.Name: nameof(Type) }];
+
+                if (!validGenericKeySelector && !validNonGenericKeySelector)
+                    return Diagnostic.Create(KeySelectorMethodHasIncorrectSignature, attribute.Location);
+            }
+
             if (assignableToType != null && attribute.AssignableToGenericArguments != null)
             {
                 var typeArguments = attribute.AssignableToGenericArguments.Value.Select(t => compilation.GetTypeByMetadataName(t)).ToArray();
@@ -142,13 +181,28 @@ private static DiagnosticModel<MethodImplementationModel> FindServicesToRegister
                             ? serviceType.ConstructUnboundGenericType().ToDisplayString()
                             : serviceType.ToDisplayString();
 
-                        var registration = new ServiceRegistrationModel(attribute.Lifetime, serviceTypeName, implementationTypeName, false, true);
+                        var registration = new ServiceRegistrationModel(
+                            attribute.Lifetime,
+                            serviceTypeName,
+                            implementationTypeName,
+                            false,
+                            true,
+                            keySelectorMethod?.Name,
+                            keySelectorMethod?.IsGenericMethod);
+
                         registrations.Add(registration);
                     }
                     else
                     {
                         var shouldResolve = attribute.AsSelf && attribute.AsImplementedInterfaces && !SymbolEqualityComparer.Default.Equals(implementationType, serviceType);
-                        var registration = new ServiceRegistrationModel(attribute.Lifetime, serviceType.ToDisplayString(), implementationType.ToDisplayString(), shouldResolve, false);
+                        var registration = new ServiceRegistrationModel(
+                            attribute.Lifetime,
+                            serviceType.ToDisplayString(),
+                            implementationType.ToDisplayString(),
+                            shouldResolve,
+                            false,
+                            keySelectorMethod?.Name,
+                            keySelectorMethod?.IsGenericMethod);
                         registrations.Add(registration);
                     }
 

ServiceScan.SourceGenerator/DiagnosticDescriptors.cs

@@ -4,38 +4,52 @@
 
 public static class DiagnosticDescriptors
 {
     public static readonly DiagnosticDescriptor NotPartialDefinition = new("DI0001",
         "Method is not partial",
         "Method with GenerateServiceRegistrations attribute must have partial modifier",
         "Usage",
         DiagnosticSeverity.Error,
         true);
 
     public static readonly DiagnosticDescriptor WrongReturnType = new("DI0002",
         "Wrong return type",
         "Method with GenerateServiceRegistrations attribute must return void or IServiceCollection",
         "Usage",
         DiagnosticSeverity.Error,
         true);
 
     public static readonly DiagnosticDescriptor WrongMethodParameters = new("DI0003",
         "Wrong method parameters",
         "Method with GenerateServiceRegistrations attribute must have a single IServiceCollection parameter",
         "Usage",
         DiagnosticSeverity.Error,
         true);
 
     public static readonly DiagnosticDescriptor MissingSearchCriteria = new("DI0004",
         "Missing search criteria",
         "GenerateServiceRegistrations must have at least one search criteria",
         "Usage",
         DiagnosticSeverity.Error,
         true);
 
     public static readonly DiagnosticDescriptor NoMatchingTypesFound = new("DI0005",
         "No matching types found",
         "There are no types matching attribute's search criteria",
         "Usage",
         DiagnosticSeverity.Warning,
         true);
+
+    public static readonly DiagnosticDescriptor KeySelectorMethodNotFound = new("DI0006",
+        "Provided KeySelector method is not found",
+        "KeySelector parameter should point to a static method in the class",
+        "Usage",
+        DiagnosticSeverity.Error,
+        true);
+
+    public static readonly DiagnosticDescriptor KeySelectorMethodHasIncorrectSignature = new("DI0007",
+        "Provided KeySelector method has incorrect signature",
+        "KeySelector should have non-void return type, and either be generic with no parameters, or non-generic with a single Type parameter",
+        "Usage",
+        DiagnosticSeverity.Error,
+        true);
 }

ServiceScan.SourceGenerator/GenerateAttributeSource.cs

@@ -45,7 +45,7 @@ internal class GenerateServiceRegistrationsAttribute : Attribute
             /// "forwarded" to "self" implementation.
             /// </summary>
             public bool AsSelf { get; set; }
-
+        
             /// <summary>
             /// Set this value to filter the types to register by their full name. 
             /// You can use '*' wildcards.
@@ -54,6 +54,14 @@ internal class GenerateServiceRegistrationsAttribute : Attribute
             /// <example>Namespace.With.Services.*</example>
             /// <example>*Service,*Factory</example>
             public string? TypeNameFilter { get; set; }
+
+            /// <summary>
+            /// Set this value to a static method name returning string.
+            /// Returned value will be used as a key for the registration.
+            /// Method should either be generic, or have a single parameter of type <see cref="Type"/>.
+            /// </summary>
+            /// <example>nameof(GetKey)</example>
+            public string? KeySelector { get; set; }
         }
         """;
 }

ServiceScan.SourceGenerator/Model/AttributeModel.cs

@@ -9,6 +9,7 @@ record AttributeModel(
     string? AssemblyOfTypeName,
     string Lifetime,
     string? TypeNameFilter,
+    string? KeySelector,
     bool AsImplementedInterfaces,
     bool AsSelf,
     Location Location,
@@ -23,6 +24,7 @@ public static AttributeModel Create(AttributeData attribute)
         var asImplementedInterfaces = attribute.NamedArguments.FirstOrDefault(a => a.Key == "AsImplementedInterfaces").Value.Value is true;
         var asSelf = attribute.NamedArguments.FirstOrDefault(a => a.Key == "AsSelf").Value.Value is true;
         var typeNameFilter = attribute.NamedArguments.FirstOrDefault(a => a.Key == "TypeNameFilter").Value.Value as string;
+        var keySelector = attribute.NamedArguments.FirstOrDefault(a => a.Key == "KeySelector").Value.Value as string;
 
         if (string.IsNullOrWhiteSpace(typeNameFilter))
             typeNameFilter = null;
@@ -46,6 +48,15 @@ public static AttributeModel Create(AttributeData attribute)
 
         var hasError = assemblyType is { TypeKind: TypeKind.Error } || assignableTo is { TypeKind: TypeKind.Error };
 
-        return new(assignableToTypeName, assignableToGenericArguments, assemblyOfTypeName, lifetime, typeNameFilter, asImplementedInterfaces, asSelf, location, hasError);
+        return new(assignableToTypeName,
+            assignableToGenericArguments,
+            assemblyOfTypeName,
+            lifetime,
+            typeNameFilter,
+            keySelector,
+            asImplementedInterfaces,
+            asSelf,
+            location,
+            hasError);
     }
 }

ServiceScan.SourceGenerator/Model/ServiceRegistrationModel.cs

@@ -5,4 +5,6 @@ record ServiceRegistrationModel(
     string ServiceTypeName,
     string ImplementationTypeName,
     bool ResolveImplementation,
-    bool IsOpenGeneric);
+    bool IsOpenGeneric,
+    string? KeySelectorMethodName,
+    bool? KeySelectorMethodGeneric);