fix: codegen visitor filtering bug (#1049)

Author: 0marperez

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/CodegenVisitor.kt

@@ -40,7 +40,7 @@ class CodegenVisitor(context: PluginContext) : ShapeVisitor.Default<Unit>() {
     private val fileManifest: FileManifest = context.fileManifest
     private val symbolProvider: SymbolProvider
     private val writers: KotlinDelegator
-    private val integrations: List<KotlinIntegration>
+    private val integrations: MutableList<KotlinIntegration>
     private val protocolGenerator: ProtocolGenerator?
     private val applicationProtocol: ApplicationProtocol
     private val baseGenerationContext: GenerationContext
@@ -50,20 +50,28 @@ class CodegenVisitor(context: PluginContext) : ShapeVisitor.Default<Unit>() {
         LOGGER.info("Discovering KotlinIntegration providers...")
         integrations = ServiceLoader.load(KotlinIntegration::class.java, classLoader)
             .onEach { integration -> LOGGER.info("Loaded KotlinIntegration: ${integration.javaClass.name}") }
-            .filter { integration -> integration.enabledForService(context.model, settings) } // TODO: Change so we don't filter until previous integrations model modifications are complete
-            .onEach { integration -> LOGGER.info("Enabled KotlinIntegration: ${integration.javaClass.name}") }
             .sortedBy(KotlinIntegration::order)
+            .toMutableList()
+
+        var resolvedModel = context.model
+        val disabledIntegrations = mutableListOf<KotlinIntegration>()
 
         LOGGER.info("Preprocessing model")
+
         // Model pre-processing:
         // 1. Start with the model from the plugin context
         // 2. Apply integrations
         // 3. Flatten error shapes (see: https://github.com/awslabs/smithy/pull/919)
         // 4. Normalize the operations
-        var resolvedModel = context.model
         for (integration in integrations) {
-            resolvedModel = integration.preprocessModel(resolvedModel, settings)
+            if (integration.enabledForService(resolvedModel, settings)) {
+                LOGGER.info("Enabled KotlinIntegration: ${integration.javaClass.name}")
+                resolvedModel = integration.preprocessModel(resolvedModel, settings)
+            } else {
+                disabledIntegrations.add(integration)
+            }
         }
+        integrations.removeAll(disabledIntegrations)
 
         resolvedModel = ModelTransformer.create()
             .copyServiceErrorsToOperations(resolvedModel, settings.getService(resolvedModel))

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/auth/SigV4AsymmetricAuthSchemeIntegration.kt

@@ -25,15 +25,15 @@ import software.amazon.smithy.model.shapes.ShapeId
  * Register support for the `aws.auth#sigv4a` auth scheme.
  */
 class SigV4AsymmetricAuthSchemeIntegration : KotlinIntegration {
-    // Needs to happen after the `SigV4AsymmetricTraitCustomization` (-60).
-    override val order: Byte = -50
-
-    // Needs to be true due to the way integrations are filtered out before application and sigV4a customization.
-    // See 'CodegenVisitor' & 'SigV4AsymmetricTraitCustomization'
-    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean = true
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        ServiceIndex
+            .of(model)
+            .getAuthSchemes(settings.service)
+            .values
+            .any { it.javaClass == SigV4ATrait::class.java }
 
     override fun authSchemes(ctx: ProtocolGenerator.GenerationContext): List<AuthSchemeHandler> =
-        if (modelHasSigV4aTrait(ctx)) listOf(SigV4AsymmetricAuthSchemeHandler()) else emptyList()
+        listOf(SigV4AsymmetricAuthSchemeHandler())
 }
 
 private class SigV4AsymmetricAuthSchemeHandler : AuthSchemeHandler {
@@ -69,10 +69,3 @@ private class SigV4AsymmetricAuthSchemeHandler : AuthSchemeHandler {
         writer.write("#T(#T, #S)", RuntimeTypes.Auth.HttpAuthAws.SigV4AsymmetricAuthScheme, RuntimeTypes.Auth.Signing.AwsSigningStandard.DefaultAwsSigner, signingService)
     }
 }
-
-internal fun modelHasSigV4aTrait(ctx: ProtocolGenerator.GenerationContext): Boolean =
-    ServiceIndex
-        .of(ctx.model)
-        .getAuthSchemes(ctx.service)
-        .values
-        .any { it.javaClass == SigV4ATrait::class.java }