feat: wrap registered services in decorator chains via [Decorate<TDecorator, TService>]

Source/Awaiten.SourceGenerators/AnalyzerReleases.Unshipped.md

@@ -24,3 +24,5 @@
  AWT120  | Awaiten   | Error    | A synchronous Func/Lazy/Owned relationship reaches an async-tainted service transitively
  AWT121  | Awaiten   | Error    | An Owned<T> disposal handle is requested through a Lazy<Owned<T>> or Lazy<Task<Owned<T>>> relationship
  AWT122  | Awaiten   | Error    | A collection dependency has an async-tainted member but is materialized synchronously
+ AWT123  | Awaiten   | Error    | A [Decorate] names a service with no registration to decorate
+ AWT124  | Awaiten   | Error    | A decorator has no single constructor parameter assignable to the decorated service type

Source/Awaiten.SourceGenerators/ContainerRegistrations.cs

@@ -68,6 +68,46 @@ public static List<RawRegistration> Collect(INamedTypeSymbol containerSymbol)
 		return result;
 	}
 
+	/// <summary>
+	///     Reads the <c>[Decorate&lt;TDecorator, TService&gt;]</c> registrations declared on a container, in
+	///     declaration order. Each carries its declaration index so equal <c>Order</c> values fall back to
+	///     declaration order when the chain is built. Collected apart from the lifetime registrations because a
+	///     decorator wraps an existing registration after coalescing rather than introducing a new service.
+	/// </summary>
+	public static List<DecorateRegistration> CollectDecorators(INamedTypeSymbol containerSymbol)
+	{
+		List<DecorateRegistration> result = new();
+		foreach (AttributeData attribute in containerSymbol.GetAttributes())
+		{
+			if (attribute.AttributeClass is not { Name: "DecorateAttribute", IsGenericType: true, TypeArguments.Length: 2, } attributeClass
+			    || attributeClass.ContainingNamespace?.ToDisplayString() != AttributeNamespace
+			    || attributeClass.TypeArguments[0] is not INamedTypeSymbol decorator
+			    || attributeClass.TypeArguments[1] is not INamedTypeSymbol service)
+			{
+				continue;
+			}
+
+			int order = 0;
+			foreach (KeyValuePair<string, TypedConstant> argument in attribute.NamedArguments)
+			{
+				if (argument.Key == "Order" && argument.Value.Value is int value)
+				{
+					order = value;
+				}
+			}
+
+			result.Add(new DecorateRegistration(
+				service.ToDisplayString(FullyQualified),
+				service,
+				decorator,
+				order,
+				result.Count,
+				attribute.ApplicationSyntaxReference?.GetSyntax().GetLocation()));
+		}
+
+		return result;
+	}
+
 	/// <summary>
 	///     Resolves how a registration produces its instance. A <c>Factory</c> argument names a container
 	///     method that produces it; an <c>Instance</c> argument names a pre-built container member to

Source/Awaiten.SourceGenerators/DecorateRegistration.cs

@@ -0,0 +1,24 @@
+using Microsoft.CodeAnalysis;
+
+namespace Awaiten.SourceGenerators;
+
+/// <summary>
+///     A single <c>[Decorate&lt;TDecorator, TService&gt;]</c> registration read from a container: the
+///     decorated service type name, the service and decorator symbols, the requested chain
+///     <see cref="Order" />, and the declaration index used to break ties between equal orders (so
+///     decorators chain in declaration order by default). Collected apart from the lifetime registrations
+///     and expanded after coalescing by <c>AwaitenGenerator.BuildDecoratorChains</c> into synthetic-keyed
+///     chain links.
+/// </summary>
+/// <remarks>
+///     Like <see cref="RawRegistration" /> this is an intermediate type consumed within a single analysis
+///     pass, so it carries the live Roslyn <see cref="Location" /> (not an equatable
+///     <c>LocationInfo</c>) and never flows through the generator's incremental cache.
+/// </remarks>
+internal sealed record DecorateRegistration(
+	string Service,
+	INamedTypeSymbol ServiceSymbol,
+	INamedTypeSymbol Decorator,
+	int Order,
+	int DeclarationOrder,
+	Location? Location);

Source/Awaiten.SourceGenerators/Diagnostics.cs

@@ -331,4 +331,29 @@ internal static class Diagnostics
 		"Awaiten",
 		DiagnosticSeverity.Error,
 		isEnabledByDefault: true);
+
+	/// <summary>
+	///     A <c>[Decorate&lt;TDecorator, TService&gt;]</c> names a service that has no registration to
+	///     decorate, so there is no inner implementation to wrap.
+	/// </summary>
+	public static readonly DiagnosticDescriptor DecoratedServiceNotRegistered = new(
+		"AWT123",
+		"Decorated service not registered",
+		"'{0}' is decorated by '{1}' but has no registration to decorate; register the service before decorating it",
+		"Awaiten",
+		DiagnosticSeverity.Error,
+		isEnabledByDefault: true);
+
+	/// <summary>
+	///     A decorator's constructor has no — or more than one ambiguous — parameter assignable to the
+	///     decorated service type, so Awaiten cannot tell which parameter receives the inner instance
+	///     (e.g. <c>LoggingDecorator(IService a, IService b)</c> or one that takes no <c>IService</c> at all).
+	/// </summary>
+	public static readonly DiagnosticDescriptor DecoratorMissingInnerParameter = new(
+		"AWT124",
+		"Decorator has no single inner parameter",
+		"The decorator '{0}' must have exactly one constructor parameter assignable to the decorated service '{1}'; it has none or several",
+		"Awaiten",
+		DiagnosticSeverity.Error,
+		isEnabledByDefault: true);
 }

Source/Awaiten.SourceGenerators/Emitter.cs

@@ -619,15 +619,15 @@ private static void EmitCacheFields(StringBuilder builder, int depth, InstanceMo
 			if (!instance.IsAsyncTainted)
 			{
 				string modifier = instance.IsReferenceType ? "private volatile " : "private ";
-				Indent(builder, depth).Append(modifier).Append(instance.ImplementationType)
+				Indent(builder, depth).Append(modifier).Append(instance.ConstructedType)
 					.Append("? ").Append(names.Field(i)).AppendLine(";");
 			}
 
 			// The async cache memoizes the construction-and-initialization Task so it is awaited exactly
 			// once; it is volatile so the lock-free read cannot observe it before its writes are visible.
 			if (instance.IsAsyncTainted)
 			{
-				Indent(builder, depth).Append("private volatile global::System.Threading.Tasks.Task<").Append(instance.ImplementationType)
+				Indent(builder, depth).Append("private volatile global::System.Threading.Tasks.Task<").Append(instance.ConstructedType)
 					.Append(">? ").Append(names.AsyncField(i)).AppendLine(";");
 			}
 		}
@@ -1158,7 +1158,7 @@ private static void EmitScopeResolver(StringBuilder builder, int depth, int inde
 		InstanceModel instance = context.Instances[index];
 		Names names = context.Names;
 		bool asyncDisposal = context.AsyncDisposal;
-		string type = instance.ImplementationType;
+		string type = instance.ConstructedType;
 		string resolver = names.Resolver(index);
 
 		if (instance.IsParameterized)
@@ -1314,7 +1314,7 @@ private static void EmitRootResolver(StringBuilder builder, int depth, int index
 	{
 		InstanceModel instance = context.Instances[index];
 		Names names = context.Names;
-		string type = instance.ImplementationType;
+		string type = instance.ConstructedType;
 		string resolver = names.Resolver(index);
 
 		if (instance.Production == ProductionKind.Instance)
@@ -1433,7 +1433,7 @@ private static void EmitAsyncScopeResolver(StringBuilder builder, int depth, int
 
 		if (instance.Lifetime == Lifetime.Singleton)
 		{
-			Indent(builder, depth).Append("protected virtual ").Append(task).Append('<').Append(instance.ImplementationType)
+			Indent(builder, depth).Append("protected virtual ").Append(task).Append('<').Append(instance.ConstructedType)
 				.Append("> ").Append(names.AsyncResolver(index)).Append('(').Append(ct).Append(") => __root.")
 				.Append(names.AsyncResolver(index)).AppendLine("(cancellationToken);");
 			return;
@@ -1467,7 +1467,7 @@ private static void EmitDelegatingSyncResolver(StringBuilder builder, int depth,
 		string signature = string.Join(", ", argTypes.Select((t, i) => $"{t} a{i}"));
 		string forward = string.Join("", argTypes.Select((_, i) => "a" + i + ", "));
 
-		Indent(builder, depth).Append("internal ").Append(instance.ImplementationType).Append(' ')
+		Indent(builder, depth).Append("internal ").Append(instance.ConstructedType).Append(' ')
 			.Append(names.Resolver(index)).Append('(').Append(signature).AppendLine(")");
 		Indent(builder, depth).AppendLine("{");
 		EmitDisposedGuard(builder, depth + 1);
@@ -1499,7 +1499,7 @@ private static void EmitAsyncCachingResolver(StringBuilder builder, int depth, i
 	{
 		InstanceModel instance = context.Instances[index];
 		Names names = context.Names;
-		string type = instance.ImplementationType;
+		string type = instance.ConstructedType;
 		string asyncResolver = names.AsyncResolver(index);
 		string asyncField = names.AsyncField(index);
 		string creator = names.AsyncCreator(index);
@@ -1553,7 +1553,7 @@ private static void EmitAsyncFreshResolver(StringBuilder builder, int depth, int
 	{
 		InstanceModel instance = context.Instances[index];
 		Names names = context.Names;
-		string type = instance.ImplementationType;
+		string type = instance.ConstructedType;
 		const string task = "global::System.Threading.Tasks.Task";
 		const string ct = "global::System.Threading.CancellationToken cancellationToken";
 
@@ -1827,7 +1827,7 @@ private static string EmitConstruction(InstanceModel instance, InstanceModel[] i
 			return $"{instance.ProductionMember}({arguments})";
 		}
 
-		return $"new {instance.ImplementationType}({arguments})";
+		return $"new {instance.ConstructedType}({arguments})";
 	}
 
 	/// <summary>

Source/Awaiten.SourceGenerators/Internals/InstanceModel.cs

@@ -41,8 +41,19 @@ internal sealed record InstanceModel(
 	bool IsAsyncTainted = false,
 	bool IsAsyncFactory = false,
 	bool RuntimeDisposalCheck = false,
-	bool IsAsyncDisposable = false)
+	bool IsAsyncDisposable = false,
+	string? EmitType = null)
 {
+	/// <summary>
+	///     The concrete type to construct (<c>new …</c>) and to use for cache fields and resolver return
+	///     types. Normally the same as <see cref="ImplementationType" />, but a decorator chain link reuses
+	///     one decorator type across several distinct instances, so it carries a synthetic
+	///     <see cref="ImplementationType" /> identity while <see cref="EmitType" /> holds the real type. Every
+	///     other use of <see cref="ImplementationType" /> stays the instance identity (keying, caching and
+	///     collection tracking), so ordinary registrations - where <see cref="EmitType" /> is null - are unaffected.
+	/// </summary>
+	public string ConstructedType => EmitType ?? ImplementationType;
+
 	/// <summary>
 	///     Whether the container owns this instance for disposal in either sense - its declared type implements
 	///     <c>IDisposable</c> or <c>IAsyncDisposable</c> - so it is tracked for teardown and its on-demand

Source/Awaiten/DecorateAttribute.cs

@@ -0,0 +1,59 @@
+using System;
+
+namespace Awaiten;
+
+// S2326: the type parameters are the Awaiten source generator's input — it reads the service and
+// decorator types from the attribute's type arguments via Roslyn symbols, so TDecorator is
+// intentionally not referenced in the attribute body (only in the type constraint).
+#pragma warning disable S2326
+
+/// <summary>
+///     Wraps the registration of <typeparamref name="TService" /> in
+///     <typeparamref name="TDecorator" />, so every consumer of <typeparamref name="TService" />
+///     receives <c>new TDecorator(inner)</c> instead of the bare implementation. The decorator's
+///     single constructor parameter of type <typeparamref name="TService" /> is supplied the
+///     decorated registration (its other constructor parameters resolve from the graph as usual),
+///     and the decorator inherits the decorated registration's lifetime.
+/// </summary>
+/// <remarks>
+///     Multiple decorators of the same service chain in declaration order: the last
+///     <c>[Decorate]</c> declared is the outermost. With <c>[Decorate&lt;D1, IService&gt;]</c>
+///     followed by <c>[Decorate&lt;D2, IService&gt;]</c>, resolving <c>IService</c> yields
+///     <c>D2(D1(Real))</c>. Use <see cref="Order" /> to position a decorator explicitly rather than
+///     by declaration order. Decorating a service also decorates every collection
+///     (<c>IEnumerable&lt;IService&gt;</c>, <c>IService[]</c>, …) view of it, so the decorator cannot
+///     be bypassed: a service with several registrations is decorated member by member. Only the unkeyed
+///     registration of <typeparamref name="TService" /> is decorated; a service registered solely under a
+///     <c>[FromKey]</c> key has no unkeyed registration to wrap and reports AWT123.
+///     <para>
+///         The type parameters are ordered <typeparamref name="TDecorator" /> then
+///         <typeparamref name="TService" /> - the concrete type first, then the service it is exposed
+///         as - to match the lifetime attributes (<c>[Transient&lt;TImplementation, TService&gt;]</c>
+///         and friends).
+///     </para>
+///     <para>
+///         The decorator and the implementation it wraps are <em>each</em> owned by the container and
+///         disposed independently, outermost first (a decorator is built after its inner, so it is
+///         disposed before it). A decorator that also disposes the inner instance handed to it would
+///         therefore dispose it twice, so leave the inner's disposal to the container - or make the
+///         decorator's <c>Dispose</c> idempotent if it must forward.
+///     </para>
+/// </remarks>
+/// <typeparam name="TDecorator">
+///     The decorator implementation; it must have exactly one constructor parameter assignable to
+///     <typeparamref name="TService" />.
+/// </typeparam>
+/// <typeparam name="TService">The service type to decorate.</typeparam>
+/// <example><c>[Decorate&lt;LoggingDecorator, IService&gt;]</c></example>
+[AttributeUsage(AttributeTargets.Class, AllowMultiple = true, Inherited = false)]
+public sealed class DecorateAttribute<TDecorator, TService> : Attribute
+	where TDecorator : class, TService
+{
+	/// <summary>
+	///     The decorator's position in the chain (ascending, outermost last). Decorators are ordered
+	///     by <see cref="Order" /> first, then by declaration order, so a lower value sits closer to the
+	///     decorated implementation and a higher value closer to the consumer. Leave unset to chain
+	///     purely by declaration order.
+	/// </summary>
+	public int Order { get; set; }
+}

Tests/Awaiten.Api.Tests/Expected/Awaiten_net10.0.txt

@@ -26,6 +26,13 @@ namespace Awaiten
         public Awaiten.LifetimeSafety LifetimeSafety { get; set; }
         public bool SyncResolveAfterInit { get; set; }
     }
+    [System.AttributeUsage(System.AttributeTargets.Class, AllowMultiple=true, Inherited=false)]
+    public sealed class DecorateAttribute<TDecorator, TService> : System.Attribute
+        where TDecorator :  class, TService
+    {
+        public DecorateAttribute() { }
+        public int Order { get; set; }
+    }
     [System.AttributeUsage(System.AttributeTargets.Parameter, AllowMultiple=false, Inherited=false)]
     public sealed class FromKeyAttribute : System.Attribute
     {

Tests/Awaiten.Api.Tests/Expected/Awaiten_net8.0.txt

@@ -26,6 +26,13 @@ namespace Awaiten
         public Awaiten.LifetimeSafety LifetimeSafety { get; set; }
         public bool SyncResolveAfterInit { get; set; }
     }
+    [System.AttributeUsage(System.AttributeTargets.Class, AllowMultiple=true, Inherited=false)]
+    public sealed class DecorateAttribute<TDecorator, TService> : System.Attribute
+        where TDecorator :  class, TService
+    {
+        public DecorateAttribute() { }
+        public int Order { get; set; }
+    }
     [System.AttributeUsage(System.AttributeTargets.Parameter, AllowMultiple=false, Inherited=false)]
     public sealed class FromKeyAttribute : System.Attribute
     {

Tests/Awaiten.Api.Tests/Expected/Awaiten_netstandard2.0.txt

@@ -26,6 +26,13 @@ namespace Awaiten
         public Awaiten.LifetimeSafety LifetimeSafety { get; set; }
         public bool SyncResolveAfterInit { get; set; }
     }
+    [System.AttributeUsage(System.AttributeTargets.Class, AllowMultiple=true, Inherited=false)]
+    public sealed class DecorateAttribute<TDecorator, TService> : System.Attribute
+        where TDecorator :  class, TService
+    {
+        public DecorateAttribute() { }
+        public int Order { get; set; }
+    }
     [System.AttributeUsage(System.AttributeTargets.Parameter, AllowMultiple=false, Inherited=false)]
     public sealed class FromKeyAttribute : System.Attribute
     {