Use post-processing step to cleanup unused component schemas

Author: bkoelman

src/JsonApiDotNetCore.OpenApi/ConfigureSwaggerGenOptions.cs

@@ -47,6 +47,7 @@ public void Configure(SwaggerGenOptions options)
         options.DocumentFilter<ServerDocumentFilter>();
         options.DocumentFilter<EndpointOrderingFilter>();
         options.OperationFilter<JsonApiOperationDocumentationFilter>();
+        options.DocumentFilter<UnusedComponentSchemaCleaner>();
     }
 
     private IEnumerable<Type> GetConstructedTypesForResourceData(Type baseType)

src/JsonApiDotNetCore.OpenApi/SwaggerComponents/DocumentSchemaGenerator.cs

@@ -163,7 +163,6 @@ private void SetJsonApiVersion(OpenApiSchema fullSchemaForDocument, SchemaReposi
             if (!_options.IncludeJsonApiVersion)
             {
                 fullSchemaForDocument.Properties.Remove(JsonApiPropertyName.Jsonapi);
-                schemaRepository.Schemas.Remove(jsonapiSchemaId);
             }
             else
             {

src/JsonApiDotNetCore.OpenApi/SwaggerComponents/ResourceDataSchemaGenerator.cs

@@ -125,7 +125,6 @@ private void SetResourceAttributes(OpenApiSchema fullSchemaForResourceData, Reso
         if (!fullSchemaForAttributes.Properties.Any())
         {
             fullSchemaForResourceData.Properties.Remove(JsonApiPropertyName.Attributes);
-            schemaRepository.Schemas.Remove(referenceSchemaForAttributes.Reference.Id);
         }
         else
         {
@@ -143,7 +142,6 @@ private void SetResourceRelationships(OpenApiSchema fullSchemaForResourceData, R
         if (!fullSchemaForRelationships.Properties.Any())
         {
             fullSchemaForResourceData.Properties.Remove(JsonApiPropertyName.Relationships);
-            schemaRepository.Schemas.Remove(referenceSchemaForRelationships.Reference.Id);
         }
         else
         {

src/JsonApiDotNetCore.OpenApi/UnusedComponentSchemaCleaner.cs

@@ -0,0 +1,60 @@
+using JetBrains.Annotations;
+using Microsoft.OpenApi.Interfaces;
+using Microsoft.OpenApi.Models;
+using Microsoft.OpenApi.Services;
+using Swashbuckle.AspNetCore.SwaggerGen;
+
+namespace JsonApiDotNetCore.OpenApi;
+
+/// <summary>
+/// Removes unreferenced component schemas from the OpenAPI document.
+/// </summary>
+[UsedImplicitly(ImplicitUseKindFlags.InstantiatedNoFixedConstructorSignature)]
+internal sealed class UnusedComponentSchemaCleaner : IDocumentFilter
+{
+    public void Apply(OpenApiDocument document, DocumentFilterContext context)
+    {
+        var visitor = new OpenApiReferenceVisitor();
+        var walker = new OpenApiWalker(visitor);
+        walker.Walk(document);
+
+        HashSet<string> unusedSchemaNames = [];
+
+        foreach (string schemaId in document.Components.Schemas.Select(schema => schema.Key).Where(schemaId => !visitor.UsedSchemaNames.Contains(schemaId)))
+        {
+            unusedSchemaNames.Add(schemaId);
+        }
+
+        foreach (string schemaId in unusedSchemaNames)
+        {
+            document.Components.Schemas.Remove(schemaId);
+        }
+    }
+
+    private sealed class OpenApiReferenceVisitor : OpenApiVisitorBase
+    {
+        private const string ComponentSchemaPrefix = "#/components/schemas/";
+
+        public HashSet<string> UsedSchemaNames { get; } = [];
+
+        public override void Visit(IOpenApiReferenceable referenceable)
+        {
+            UsedSchemaNames.Add(referenceable.Reference.Id);
+        }
+
+        public override void Visit(OpenApiSchema schema)
+        {
+            if (schema.Discriminator != null)
+            {
+                foreach (string discriminatorValue in schema.Discriminator.Mapping.Values)
+                {
+                    if (discriminatorValue.StartsWith(ComponentSchemaPrefix, StringComparison.Ordinal))
+                    {
+                        string schemaId = discriminatorValue[ComponentSchemaPrefix.Length..];
+                        UsedSchemaNames.Add(schemaId);
+                    }
+                }
+            }
+        }
+    }
+}