Support AsSelf registration (#6)

README.md

@@ -76,7 +76,8 @@ It adds MediatR handlers, which would work for simple cases, although you might
 | Property | Description |
 | --- | --- |
 | **FromAssemblyOf** |Set the assembly containing the given type as the source of types to register. If not specified, the assembly containing the method with this attribute will be used. |
-| **AssignableTo** | Set the type that the registered types must be assignable to. Types will be registered with this type as the service type. |
+| **AssignableTo** | Set the type that the registered types must be assignable to. Types will be registered with this type as the service type, unless `AsImplementedInterfaces` or `AsSelf` is set. |
 | **Lifetime** | Set the lifetime of the registered services. `ServiceLifetime.Transient` is used if not specified. |
-| **AsImplementedInterfaces** | If true, the registered types will be registered as implemented interfaces instead of their actual type. This option is ignored if `AssignableTo` is set. |
+| **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. |

ServiceScan.SourceGenerator.Tests/AddServicesTests.cs

@@ -200,6 +200,59 @@ public class MyService2 : AbstractService { }
         Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
     }
 
+    [Fact]
+    public void AddServicesAssignableToAbstractClassAsSelf()
+    {
+        var attribute = $"[GenerateServiceRegistrations(AssignableTo = typeof(AbstractService), AsSelf = true)]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public abstract class AbstractService { }
+            public class MyService1 : AbstractService { }
+            public class MyService2 : AbstractService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var registrations = $"""
+            return services
+                .AddTransient<GeneratorTests.MyService1, GeneratorTests.MyService1>()
+                .AddTransient<GeneratorTests.MyService2, GeneratorTests.MyService2>();
+            """;
+        Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
+    }
+
+    [Fact]
+    public void AddServiceAssignableToSelf()
+    {
+        var attribute = $"[GenerateServiceRegistrations(AssignableTo = typeof(MyService))]";
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public class MyService { }
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var registrations = $"""
+            return services
+                .AddTransient<GeneratorTests.MyService, GeneratorTests.MyService>();
+            """;
+        Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
+    }
+
     [Fact]
     public void AddServicesAssignableToOpenGenericAbstractClass()
     {
@@ -374,6 +427,41 @@ public class InterfacelessService {}
         Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
     }
 
+    [Fact]
+    public void AddServicesBothAsSelfAndAsImplementedInterfaces()
+    {
+        var attribute = """
+            [GenerateServiceRegistrations(
+                TypeNameFilter = "*Service", 
+                AsImplementedInterfaces = true, 
+                AsSelf = true, 
+                Lifetime = ServiceLifetime.Singleton))]
+            """;
+
+        var compilation = CreateCompilation(
+            Sources.MethodWithAttribute(attribute),
+            """
+            namespace GeneratorTests;
+
+            public interface IServiceA {}
+            public interface IServiceB {}
+            public class MyService: IServiceA, IServiceB {}
+            """);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var registrations = $"""
+            return services
+                .AddSingleton<GeneratorTests.MyService, GeneratorTests.MyService>()
+                .AddSingleton<GeneratorTests.IServiceA>(s => s.GetRequiredService<GeneratorTests.MyService>())
+                .AddSingleton<GeneratorTests.IServiceB>(s => s.GetRequiredService<GeneratorTests.MyService>());
+            """;
+        Assert.Equal(Sources.GetMethodImplementation(registrations), results.GeneratedTrees[1].ToString());
+    }
+
     [Fact]
     public void DontGenerateAnythingIfTypeIsInvalid()
     {

ServiceScan.SourceGenerator/DependencyInjectionGenerator.cs

@@ -50,7 +50,10 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                     }
                     else
                     {
-                        sb.AppendLine($"            .Add{registration.Lifetime}<{registration.ServiceTypeName}, {registration.ImplementationTypeName}>()");
+                        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}>()");
                     }
                 }
 
@@ -122,18 +125,13 @@ private static DiagnosticModel<MethodImplementationModel> FindServicesToRegister
                 if (assignableToType != null && !IsAssignableTo(implementationType, assignableToType, out matchedType))
                     continue;
 
-                IEnumerable<INamedTypeSymbol> serviceTypes = null;
-
-                if (matchedType != null)
-                {
-                    serviceTypes = [matchedType];
-                }
-                else
+                IEnumerable<INamedTypeSymbol> serviceTypes = (attribute.AsSelf, attribute.AsImplementedInterfaces) switch
                 {
-                    serviceTypes = attribute.AsImplementedInterfaces
-                        ? implementationType.AllInterfaces
-                        : [implementationType];
-                }
+                    (true, true) => new[] { implementationType }.Concat(implementationType.AllInterfaces),
+                    (false, true) => implementationType.AllInterfaces,
+                    (true, false) => [implementationType],
+                    _ => [matchedType ?? implementationType]
+                };
 
                 foreach (var serviceType in serviceTypes)
                 {
@@ -144,13 +142,13 @@ private static DiagnosticModel<MethodImplementationModel> FindServicesToRegister
                             ? serviceType.ConstructUnboundGenericType().ToDisplayString()
                             : serviceType.ToDisplayString();
 
-                        var registration = new ServiceRegistrationModel(attribute.Lifetime, serviceTypeName, implementationTypeName, true);
+                        var registration = new ServiceRegistrationModel(attribute.Lifetime, serviceTypeName, implementationTypeName, false, true);
                         registrations.Add(registration);
                     }
                     else
                     {
-
-                        var registration = new ServiceRegistrationModel(attribute.Lifetime, serviceType.ToDisplayString(), implementationType.ToDisplayString(), false);
+                        var shouldResolve = attribute.AsSelf && attribute.AsImplementedInterfaces && !SymbolEqualityComparer.Default.Equals(implementationType, serviceType);
+                        var registration = new ServiceRegistrationModel(attribute.Lifetime, serviceType.ToDisplayString(), implementationType.ToDisplayString(), shouldResolve, false);
                         registrations.Add(registration);
                     }
 

ServiceScan.SourceGenerator/GenerateAttributeSource.cs

@@ -23,7 +23,8 @@ internal class GenerateServiceRegistrationsAttribute : Attribute
 
             /// <summary>
             /// Set the type that the registered types must be assignable to.
-            /// Types will be registered with this type as the service type.
+            /// Types will be registered with this type as the service type,
+            /// unless <see cref="AsImplementedInterfaces"/> or <see cref="AsSelf"/> is set.
             /// </summary>
             public Type? AssignableTo { get; set; }
 
@@ -32,12 +33,18 @@ internal class GenerateServiceRegistrationsAttribute : Attribute
             /// <see cref="ServiceLifetime.Transient"/> is used if not specified.
             /// </summary>
             public ServiceLifetime Lifetime { get; set; }
-
+        
             /// <summary>
-            /// If set to true, the registered types will be registered as implemented interfaces instead of their actual type.
-            /// This option is ignored if <see cref="AssignableTo"/> is set.
+            /// If set to true, types will be registered as implemented interfaces instead of their actual type.
             /// </summary>
             public bool AsImplementedInterfaces { get; set; }
+        
+            /// <summary>
+            /// If set to true, types will be registered with their actual type.
+            /// It can be combined with <see cref="AsImplementedInterfaces"/>, in that case implemeted interfaces will be
+            /// "forwarded" to "self" implementation.
+            /// </summary>
+            public bool AsSelf { get; set; }
 
             /// <summary>
             /// Set this value to filter the types to register by their full name. 

ServiceScan.SourceGenerator/Model/AttributeModel.cs

@@ -10,6 +10,7 @@ record AttributeModel(
     string Lifetime,
     string? TypeNameFilter,
     bool AsImplementedInterfaces,
+    bool AsSelf,
     Location Location,
     bool HasErrors)
 {
@@ -20,6 +21,7 @@ public static AttributeModel Create(AttributeData attribute)
         var assemblyType = attribute.NamedArguments.FirstOrDefault(a => a.Key == "FromAssemblyOf").Value.Value as INamedTypeSymbol;
         var assignableTo = attribute.NamedArguments.FirstOrDefault(a => a.Key == "AssignableTo").Value.Value as INamedTypeSymbol;
         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;
 
         if (string.IsNullOrWhiteSpace(typeNameFilter))
@@ -44,6 +46,6 @@ 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, location, hasError);
+        return new(assignableToTypeName, assignableToGenericArguments, assemblyOfTypeName, lifetime, typeNameFilter, asImplementedInterfaces, asSelf, location, hasError);
     }
 }

ServiceScan.SourceGenerator/Model/ServiceRegistrationModel.cs

@@ -4,4 +4,5 @@ record ServiceRegistrationModel(
     string Lifetime,
     string ServiceTypeName,
     string ImplementationTypeName,
+    bool ResolveImplementation,
     bool IsOpenGeneric);