Create a more generic form for registration enrichers

src/main/Yardarm.Client/Serialization/Literals/LiteralConverterRegistry.cs

@@ -147,10 +147,10 @@ public LiteralConverterRegistry Add<T>(LiteralConverter<T?> converter)
     }
 
     /// <summary>
-    /// Returns a new <see cref="LiteralConverterRegistry"/> with the all default converters registered.
+    /// Returns a new <see cref="LiteralConverterRegistry"/> with the built-in converters registered.
     /// </summary>
     /// <returns>A new <see cref="LiteralConverterRegistry"/>.</returns>
-    public static LiteralConverterRegistry CreateDefaultRegistry() =>
+    public static LiteralConverterRegistry CreateBasicRegistry() =>
         new LiteralConverterRegistry()
             .Add(new BooleanLiteralConverter())
             .Add(new DateTimeLiteralConverter())
@@ -175,4 +175,15 @@ public static LiteralConverterRegistry CreateDefaultRegistry() =>
             .Add(new TimeOnlyLiteralConverter())
 #endif
             ;
+
+    /// <summary>
+    /// Returns a new <see cref="LiteralConverterRegistry"/> with the all default converters registered.
+    /// </summary>
+    /// <returns>A new <see cref="LiteralConverterRegistry"/>.</returns>
+    public static LiteralConverterRegistry CreateDefaultRegistry()
+    {
+        LiteralConverterRegistry registry = CreateBasicRegistry();
+        // Enrichment point
+        return registry;
+    }
 }

src/main/Yardarm.Client/Serialization/TypeSerializerRegistry.cs

@@ -114,7 +114,11 @@ public static ITypeSerializerRegistry CreateBasicRegistry() =>
         /// Returns a new <see cref="ITypeSerializerRegistry"/> with the all default serializers registered.
         /// </summary>
         /// <returns>A new <see cref="ITypeSerializerRegistry"/>.</returns>
-        public static ITypeSerializerRegistry CreateDefaultRegistry() =>
-            CreateBasicRegistry();
+        public static ITypeSerializerRegistry CreateDefaultRegistry()
+        {
+            ITypeSerializerRegistry registry = CreateBasicRegistry();
+            // Enrichment point
+            return registry;
+        }
     }
 }

src/main/Yardarm.NewtonsoftJson/JsonCreateDefaultRegistryEnricher.cs

@@ -1,13 +1,12 @@
 using System;
-using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
-using Yardarm.Enrichment;
+using Yardarm.Enrichment.Registration;
 using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
 
 namespace Yardarm.NewtonsoftJson
 {
-    public class JsonCreateDefaultRegistryEnricher : ICreateDefaultRegistryEnricher
+    public class JsonCreateDefaultRegistryEnricher : ReturnValueRegistrationEnricher, ICreateDefaultRegistryEnricher
     {
         private readonly IJsonSerializationNamespace _jsonSerializationNamespace;
 
@@ -18,10 +17,10 @@ public JsonCreateDefaultRegistryEnricher(IJsonSerializationNamespace jsonSeriali
             _jsonSerializationNamespace = jsonSerializationNamespace;
         }
 
-        public ExpressionSyntax Enrich(ExpressionSyntax target) =>
+        protected override ExpressionSyntax EnrichReturnValue(ExpressionSyntax expression) =>
             InvocationExpression(
                 MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
-                    target.WithTrailingTrivia(TriviaList(CarriageReturnLineFeed, Whitespace("                "))),
+                    expression,
                     GenericName(
                         Identifier("Add"),
                         TypeArgumentList(SingletonSeparatedList<TypeSyntax>(_jsonSerializationNamespace.JsonTypeSerializer)))),

src/main/Yardarm.SystemTextJson/JsonCreateDefaultRegistryEnricher.cs

@@ -1,14 +1,13 @@
 using System;
-using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
-using Yardarm.Enrichment;
+using Yardarm.Enrichment.Registration;
 using Yardarm.SystemTextJson.Internal;
 using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
 
 namespace Yardarm.SystemTextJson
 {
-    public class JsonCreateDefaultRegistryEnricher : ICreateDefaultRegistryEnricher
+    public class JsonCreateDefaultRegistryEnricher : ReturnValueRegistrationEnricher, ICreateDefaultRegistryEnricher
     {
         private readonly IJsonSerializationNamespace _jsonSerializationNamespace;
 
@@ -19,16 +18,16 @@ public JsonCreateDefaultRegistryEnricher(IJsonSerializationNamespace jsonSeriali
             _jsonSerializationNamespace = jsonSerializationNamespace;
         }
 
-        public ExpressionSyntax Enrich(ExpressionSyntax target) =>
+        protected override ExpressionSyntax EnrichReturnValue(ExpressionSyntax target) =>
             // Don't use the Add<T> overload here because it will cause trimming to retain
             // all constructors. This will then cause IL2026 warnings if trimming is enabled.
             // Instead create a new instance directly using the default constructor and add it.
             InvocationExpression(
                 MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
-                    target.WithTrailingTrivia(TriviaList(CarriageReturnLineFeed, Whitespace("                "))),
+                    target,
                     IdentifierName("Add")),
-                ArgumentList(SeparatedList(new[]
-                {
+                ArgumentList(SeparatedList(
+                [
                     Argument(MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
                         _jsonSerializationNamespace.JsonTypeSerializer,
                         IdentifierName("SupportedMediaTypes"))),
@@ -38,6 +37,6 @@ public ExpressionSyntax Enrich(ExpressionSyntax target) =>
                                 QualifiedName(_jsonSerializationNamespace.Name, IdentifierName(JsonSerializerContextGenerator.TypeName)),
                                 IdentifierName("Default"))))),
                             null))
-                })));
+                ])));
     }
 }

src/main/Yardarm/Enrichment/Compilation/DefaultLiteralConvertersEnricher.cs

@@ -1,15 +1,18 @@
-using System;
-using System.Collections.Generic;
+using System.Collections.Generic;
 using System.Linq;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
+using Microsoft.Extensions.DependencyInjection;
+using Yardarm.Enrichment.Registration;
 
 namespace Yardarm.Enrichment.Compilation;
 
 public class DefaultLiteralConvertersEnricher(
-    IEnumerable<IDefaultLiteralConverterEnricher> createDefaultRegistryEnrichers)
+    [FromKeyedServices(DefaultLiteralConvertersEnricher.RegistrationEnricherKey)] IEnumerable<IRegistrationEnricher> createDefaultRegistryEnrichers)
     : IResourceFileEnricher
 {
+    public const string RegistrationEnricherKey = "DefaultLiteralConverters";
+
     public bool ShouldEnrich(string resourceName) =>
         resourceName == "Yardarm.Client.Serialization.Literals.LiteralConverterRegistry.cs";
 
@@ -25,11 +28,10 @@ public CompilationUnitSyntax Enrich(CompilationUnitSyntax target, ResourceFileEn
             .OfType<MethodDeclarationSyntax>()
             .FirstOrDefault(p => p.Identifier.ValueText == "CreateDefaultRegistry");
 
-        if (methodDeclaration?.ExpressionBody != null)
+        if (methodDeclaration?.Body is { } body)
         {
-            MethodDeclarationSyntax newMethodDeclaration = methodDeclaration.WithExpressionBody(
-                methodDeclaration.ExpressionBody.WithExpression(
-                    methodDeclaration.ExpressionBody.Expression.Enrich(createDefaultRegistryEnrichers)));
+            MethodDeclarationSyntax newMethodDeclaration = methodDeclaration.WithBody(
+                body.Enrich(createDefaultRegistryEnrichers));
 
             target = target.ReplaceNode(methodDeclaration, newMethodDeclaration);
         }

src/main/Yardarm/Enrichment/Compilation/DefaultTypeSerializersEnricher.cs

@@ -2,13 +2,17 @@
 using System.Linq;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
+using Microsoft.Extensions.DependencyInjection;
+using Yardarm.Enrichment.Registration;
 
 namespace Yardarm.Enrichment.Compilation;
 
 public class DefaultTypeSerializersEnricher(
-    IEnumerable<ICreateDefaultRegistryEnricher> createDefaultRegistryEnrichers) :
+    [FromKeyedServices(DefaultTypeSerializersEnricher.RegistrationEnricherKey)] IEnumerable<IRegistrationEnricher> createDefaultRegistryEnrichers) :
     IResourceFileEnricher
 {
+    public const string RegistrationEnricherKey = "DefaultTypeSerializers";
+
     public bool ShouldEnrich(string resourceName) =>
         resourceName == "Yardarm.Client.Serialization.TypeSerializerRegistry.cs";
 
@@ -24,11 +28,10 @@ public CompilationUnitSyntax Enrich(CompilationUnitSyntax target, ResourceFileEn
             .OfType<MethodDeclarationSyntax>()
             .FirstOrDefault(p => p.Identifier.ValueText == "CreateDefaultRegistry");
 
-        if (methodDeclaration?.ExpressionBody != null)
+        if (methodDeclaration?.Body is { } body)
         {
-            MethodDeclarationSyntax newMethodDeclaration = methodDeclaration.WithExpressionBody(
-                methodDeclaration.ExpressionBody.WithExpression(
-                    methodDeclaration.ExpressionBody.Expression.Enrich(createDefaultRegistryEnrichers)));
+            MethodDeclarationSyntax newMethodDeclaration = methodDeclaration.WithBody(
+                body.Enrich(createDefaultRegistryEnrichers));
 
             target = target.ReplaceNode(methodDeclaration, newMethodDeclaration);
         }

src/main/Yardarm/Enrichment/EnricherServiceCollectionExtensions.cs

@@ -1,7 +1,9 @@
-using Microsoft.Extensions.DependencyInjection;
+using System;
+using Microsoft.Extensions.DependencyInjection;
 using Yardarm.Enrichment.Authentication;
 using Yardarm.Enrichment.Compilation;
 using Yardarm.Enrichment.Packaging;
+using Yardarm.Enrichment.Registration;
 using Yardarm.Enrichment.Requests;
 using Yardarm.Enrichment.Responses;
 using Yardarm.Enrichment.Schema;
@@ -21,13 +23,22 @@ public static IServiceCollection AddDefaultEnrichers(this IServiceCollection ser
                 .AddDefaultResponseEnrichers()
                 .AddDefaultTagEnrichers();
 
+        public static IServiceCollection AddRegistrationEnricher<T>(this IServiceCollection services, string registrationType)
+            where T : class, IRegistrationEnricher
+        {
+            ArgumentNullException.ThrowIfNull(services);
+            ArgumentNullException.ThrowIfNull(registrationType);
+
+            return services.AddKeyedTransient<IRegistrationEnricher, T>(registrationType);
+        }
+
         public static IServiceCollection AddCreateDefaultRegistryEnricher<T>(this IServiceCollection services)
             where T : class, ICreateDefaultRegistryEnricher =>
-            services.AddTransient<ICreateDefaultRegistryEnricher, T>();
+            services.AddRegistrationEnricher<T>(DefaultTypeSerializersEnricher.RegistrationEnricherKey);
 
         public static IServiceCollection AddDefaultLiteralConverterEnricher<T>(this IServiceCollection services)
             where T : class, IDefaultLiteralConverterEnricher =>
-            services.AddTransient<IDefaultLiteralConverterEnricher, T>();
+            services.AddRegistrationEnricher<T>(DefaultLiteralConvertersEnricher.RegistrationEnricherKey);
 
         public static IServiceCollection AddOpenApiSyntaxNodeEnricher<T>(this IServiceCollection services)
             where T : class, IOpenApiSyntaxNodeEnricher =>

src/main/Yardarm/Enrichment/Registration/ICreateDefaultRegistryEnricher.cs

@@ -0,0 +1,6 @@
+namespace Yardarm.Enrichment.Registration
+{
+    public interface ICreateDefaultRegistryEnricher : IRegistrationEnricher
+    {
+    }
+}

src/main/Yardarm/Enrichment/Registration/IDefaultLiteralConverterEnricher.cs

@@ -0,0 +1,8 @@
+namespace Yardarm.Enrichment.Registration;
+
+/// <summary>
+/// Enriches the body of the CreateDefaultRegistry method in the LiteralConverterRegistry class.
+/// </summary>
+public interface IDefaultLiteralConverterEnricher : IRegistrationEnricher
+{
+}

src/main/Yardarm/Enrichment/Registration/IRegistrationEnricher.cs

@@ -0,0 +1,10 @@
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
+namespace Yardarm.Enrichment.Registration;
+
+/// <summary>
+/// An expression enricher used by extensions to register themselves in the client.
+/// </summary>
+public interface IRegistrationEnricher : IEnricher<BlockSyntax>
+{
+}

src/main/Yardarm/Enrichment/Registration/ReturnValueRegistrationEnricher.cs

@@ -0,0 +1,35 @@
+using System.Collections.Generic;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
+
+namespace Yardarm.Enrichment.Registration;
+
+/// <summary>
+/// Common implementation for an <see cref="IRegistrationEnricher"> that modifies the return value of a method.
+/// </summary>
+public abstract class ReturnValueRegistrationEnricher : StatementInsertingRegistrationEnricher
+{
+    protected override sealed IEnumerable<StatementSyntax> GenerateStatements(ExpressionSyntax? returnExpression)
+    {
+        if (returnExpression is not IdentifierNameSyntax identifier)
+        {
+            return [];
+        }
+
+        ExpressionSyntax newReturnValue = EnrichReturnValue(identifier);
+        if (newReturnValue == returnExpression)
+        {
+            // Was not changed
+            return [];
+        }
+
+        ExpressionStatementSyntax statement = ExpressionStatement(AssignmentExpression(
+            Microsoft.CodeAnalysis.CSharp.SyntaxKind.SimpleAssignmentExpression,
+            identifier,
+            newReturnValue));
+
+        return [statement];
+    }
+
+    protected abstract ExpressionSyntax EnrichReturnValue(ExpressionSyntax expression);
+}

src/main/Yardarm/Enrichment/Registration/StatementInsertingRegistrationEnricher.cs

@@ -0,0 +1,38 @@
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
+namespace Yardarm.Enrichment.Registration;
+
+/// <summary>
+/// Common implementation for an <see cref="IRegistrationEnricher"> that inserts statements before the return
+/// at the end of the method.
+/// </summary>
+public abstract class StatementInsertingRegistrationEnricher : IRegistrationEnricher
+{
+    public BlockSyntax Enrich(BlockSyntax target)
+    {
+        int returnStatementIndex = target.Statements.LastIndexOf(p => p is ReturnStatementSyntax);
+        ExpressionSyntax? returnExpression = returnStatementIndex >= 0
+            ? ((ReturnStatementSyntax)target.Statements[returnStatementIndex]).Expression
+            : null;
+
+        IEnumerable<StatementSyntax> newStatements = GenerateStatements(returnExpression);
+        if (newStatements == Enumerable.Empty<StatementSyntax>() ||
+            newStatements is ICollection<StatementSyntax> { Count: 0 })
+        {
+            // No statements added, short-circuit
+            return target;
+        }
+
+        return target.WithStatements(
+            target.Statements.InsertRange(returnStatementIndex, newStatements));
+    }
+
+    /// <summary>
+    /// Generates the statements to insert before the return statement.
+    /// </summary>
+    /// <param name="returnExpression">The expression being returned in the return statement.</param>
+    /// <returns></returns>
+    protected abstract IEnumerable<StatementSyntax> GenerateStatements(ExpressionSyntax? returnExpression);
+}