Add Decorate<TDecorated, TDecorator> support

readme.md

@@ -174,6 +174,63 @@ Note you can also register the same service using multiple keys, as shown in the
 > [!IMPORTANT]
 > Keyed services are a feature of version 8.0+ of Microsoft.Extensions.DependencyInjection
 
+### Decorating Services
+
+After services are registered with `AddServices`, you can wrap existing registrations 
+with a decorator using the `Decorate<TDecorated, TDecorator>()` extension method. 
+The decorated service must already be registered, and the source generator replaces 
+matching registrations in-place while preserving each registration's lifetime.
+
+The decorator type must implement `TDecorated` and provide a constructor that accepts 
+the decorated service as one of its parameters (additional dependencies are resolved 
+from the container as usual). Annotating the decorator with `[Service]` is optional 
+(when present, lifetime compatibility with the decorated service is validated at compile time):
+
+```csharp
+public interface INotificationService
+{
+    void Send(string message);
+}
+
+[Service(ServiceLifetime.Scoped)]
+public class EmailNotificationService : INotificationService
+{
+    public void Send(string message) => Console.WriteLine($"[Email] {message}");
+}
+
+[Service(ServiceLifetime.Scoped)]
+public class LoggingNotificationService(INotificationService inner) : INotificationService
+{
+    public void Send(string message)
+    {
+        Console.WriteLine("Sending notification...");
+        inner.Send(message);
+    }
+}
+
+var builder = WebApplication.CreateBuilder(args);
+
+builder.Services.AddServices();
+builder.Services.Decorate<INotificationService, LoggingNotificationService>();
+```
+
+When resolving `INotificationService`, the container returns a `LoggingNotificationService` 
+that wraps the original `EmailNotificationService`. `Func<T>` and `Lazy<T>` registrations 
+for the same service type also resolve the decorated instance.
+
+If multiple registrations exist for the same service type, all of them are decorated. 
+For keyed services, pass the key to decorate a specific registration:
+
+```csharp
+builder.Services.AddServices();
+builder.Services.Decorate<INotificationService, LoggingNotificationService>("email");
+```
+
+The generator validates decorations at compile-time: the decorator must have a constructor 
+that accepts the decorated service type (plus any additional dependencies). If the decorator 
+is annotated with `[Service]`, its lifetime is also validated for compatibility with the 
+decorated registration(s).
+
 ## How It Works
 
 In all cases, the generated code that implements the registration looks like the following:

src/CodeAnalysis.Tests/CodeAnalysis.Tests.csproj

@@ -12,9 +12,9 @@
     <PackageReference Include="xunit" Version="2.9.3" />
     <PackageReference Include="xunit.runner.visualstudio" Version="3.1.5" PrivateAssets="all" />
     <PackageReference Include="Microsoft.CodeAnalysis.CSharp.Workspaces" Version="5.0.0" />
-    <PackageReference Include="Microsoft.CodeAnalysis.CSharp.Analyzer.Testing" Version="1.1.3" />
-    <PackageReference Include="Microsoft.CodeAnalysis.CSharp.SourceGenerators.Testing" Version="1.1.2" />
-    <PackageReference Include="Microsoft.CodeAnalysis.CSharp.CodeFix.Testing" Version="1.1.2" />
+    <PackageReference Include="Microsoft.CodeAnalysis.CSharp.Analyzer.Testing" Version="1.1.4" />
+    <PackageReference Include="Microsoft.CodeAnalysis.CSharp.SourceGenerators.Testing" Version="1.1.4" />
+    <PackageReference Include="Microsoft.CodeAnalysis.CSharp.CodeFix.Testing" Version="1.1.4" />
     <PackageReference Include="ThisAssembly.Resources" Version="2.1.2" PrivateAssets="all" />
   </ItemGroup>
 

src/CodeAnalysis.Tests/DecorateGeneratorTests.cs

@@ -0,0 +1,244 @@
+using System.Collections.Immutable;
+using System.IO;
+using System.Threading.Tasks;
+using Devlooped.Extensions.DependencyInjection;
+using Microsoft.CodeAnalysis.CSharp.Testing;
+using Microsoft.CodeAnalysis.Testing;
+using Xunit;
+using Xunit.Abstractions;
+using Verifier = Microsoft.CodeAnalysis.CSharp.Testing.CSharpAnalyzerVerifier<Devlooped.Extensions.DependencyInjection.AddServicesAnalyzer, Microsoft.CodeAnalysis.Testing.DefaultVerifier>;
+
+namespace Tests.CodeAnalysis;
+
+public class DecorateGeneratorTests(ITestOutputHelper Output)
+{
+    [Fact]
+    public async Task ErrorIfDecoratorLifetimeIsIncompatible()
+    {
+        var test = CreateTest(
+            """
+            using Microsoft.Extensions.DependencyInjection;
+
+            public interface IFoo { }
+
+            [Service(ServiceLifetime.Scoped)]
+            public class Foo : IFoo { }
+
+            [Service(ServiceLifetime.Singleton)]
+            public class FooDecorator(IFoo inner) : IFoo { }
+
+            public static class Program
+            {
+                public static void Main()
+                {
+                    var services = new ServiceCollection();
+                    services.AddServices();
+                    {|#0:services.Decorate<IFoo, FooDecorator>()|};
+                }
+            }
+            """);
+
+        test.ExpectedDiagnostics.Add(
+            Verifier.Diagnostic(IncrementalGenerator.DecoratorLifetimeIncompatible)
+                .WithLocation(0)
+                .WithArguments("FooDecorator", "Singleton", "IFoo", "Scoped"));
+
+        await test.RunAsync();
+    }
+
+    [Fact]
+    public async Task ErrorIfDecoratorConstructorDoesNotAcceptDecoratedService()
+    {
+        var test = CreateTest(
+            """
+            using Microsoft.Extensions.DependencyInjection;
+
+            public interface IFoo { }
+
+            [Service]
+            public class Foo : IFoo { }
+
+            [Service]
+            public class FooDecorator : IFoo { }
+
+            public static class Program
+            {
+                public static void Main()
+                {
+                    var services = new ServiceCollection();
+                    services.AddServices();
+                    {|#0:services.Decorate<IFoo, FooDecorator>()|};
+                }
+            }
+            """);
+
+        test.ExpectedDiagnostics.Add(
+            Verifier.Diagnostic(IncrementalGenerator.DecoratorConstructorMissing)
+                .WithLocation(0)
+                .WithArguments("FooDecorator", "IFoo"));
+
+        await test.RunAsync();
+    }
+
+    [Fact]
+    public async Task NoErrorIfDecoratorConstructorHasOtherDependencies()
+    {
+        var test = CreateTest(
+            """
+            using Microsoft.Extensions.DependencyInjection;
+
+            public interface IFoo { }
+            public interface IOtherDependency { }
+
+            [Service]
+            public class Foo : IFoo { }
+
+            [Service]
+            public class OtherDependency : IOtherDependency { }
+
+            [Service]
+            public class FooDecorator(IFoo inner, IOtherDependency other) : IFoo { }
+
+            public static class Program
+            {
+                public static void Main()
+                {
+                    var services = new ServiceCollection();
+                    services.AddServices();
+                    services.Decorate<IFoo, FooDecorator>();
+                }
+            }
+            """);
+
+        await test.RunAsync();
+    }
+
+    [Fact]
+    public async Task NoErrorIfDecoratorHasNoServiceAttribute()
+    {
+        var test = CreateTest(
+            """
+            using Microsoft.Extensions.DependencyInjection;
+
+            public interface IFoo { }
+
+            [Service]
+            public class Foo : IFoo { }
+
+            // Decorator intentionally has NO [Service] attribute
+            public class FooDecorator(IFoo inner) : IFoo { }
+
+            public static class Program
+            {
+                public static void Main()
+                {
+                    var services = new ServiceCollection();
+                    services.AddServices();
+                    services.Decorate<IFoo, FooDecorator>();
+                }
+            }
+            """);
+
+        // No diagnostics expected — decorator no longer requires [Service]
+        await test.RunAsync();
+    }
+
+    [Fact]
+    public async Task NoErrorIfKeyedDecoratorLifetimeMatchesSelectedKey()
+    {
+        var test = CreateTest(
+            """
+            using Microsoft.Extensions.DependencyInjection;
+
+            public interface IFoo { }
+
+            [Service("foo", ServiceLifetime.Scoped)]
+            public class Foo : IFoo { }
+
+            [Service("bar", ServiceLifetime.Singleton)]
+            public class OtherFoo : IFoo { }
+
+            [Service("foo", ServiceLifetime.Scoped)]
+            public class FooDecorator(IFoo inner) : IFoo { }
+
+            public static class Program
+            {
+                public static void Main()
+                {
+                    var services = new ServiceCollection();
+                    services.AddServices();
+                    services.Decorate<IFoo, FooDecorator>("foo");
+                }
+            }
+            """);
+
+        await test.RunAsync();
+    }
+
+    [Fact]
+    public async Task ErrorIfKeyedDecoratorLifetimeIsIncompatible()
+    {
+        var test = CreateTest(
+            """
+            using Microsoft.Extensions.DependencyInjection;
+
+            public interface IFoo { }
+
+            [Service("foo", ServiceLifetime.Scoped)]
+            public class Foo : IFoo { }
+
+            [Service("foo", ServiceLifetime.Singleton)]
+            public class FooDecorator(IFoo inner) : IFoo { }
+
+            public static class Program
+            {
+                public static void Main()
+                {
+                    var services = new ServiceCollection();
+                    services.AddServices();
+                    {|#0:services.Decorate<IFoo, FooDecorator>("foo")|};
+                }
+            }
+            """);
+
+        test.ExpectedDiagnostics.Add(
+            Verifier.Diagnostic(IncrementalGenerator.DecoratorLifetimeIncompatible)
+                .WithLocation(0)
+                .WithArguments("FooDecorator", "Singleton", "IFoo", "Scoped"));
+
+        await test.RunAsync();
+    }
+
+    static CSharpSourceGeneratorTest<IncrementalGenerator, DefaultVerifier> CreateTest(string source)
+    {
+        return new CSharpSourceGeneratorTest<IncrementalGenerator, DefaultVerifier>
+        {
+            TestBehaviors = TestBehaviors.SkipGeneratedSourcesCheck,
+            TestCode = source,
+            TestState =
+            {
+                AnalyzerConfigFiles =
+                {
+                    ("/.editorconfig",
+                    """
+                    is_global = true
+                    build_property.AddServicesExtension = true
+                    """)
+                },
+                Sources =
+                {
+                    ThisAssembly.Resources.AddServicesNoReflectionExtension.Text,
+                    ThisAssembly.Resources.ServiceAttribute.Text,
+                    ThisAssembly.Resources.ServiceAttribute_1.Text
+                },
+                ReferenceAssemblies = new ReferenceAssemblies(
+                    "net8.0",
+                    new PackageIdentity(
+                        "Microsoft.NETCore.App.Ref", "8.0.0"),
+                        Path.Combine("ref", "net8.0"))
+                    .AddPackages(ImmutableArray.Create(
+                        new PackageIdentity("Microsoft.Extensions.DependencyInjection", "8.0.0")))
+            },
+        }.WithPreprocessorSymbols();
+    }
+}