Incorporate feedback

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/ClientCodegenDecorator.kt

@@ -6,20 +6,16 @@
 package software.amazon.smithy.rust.codegen.client.smithy.customize
 
 import software.amazon.smithy.build.PluginContext
-import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.OperationShape
-import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.ClientProtocolGenerator
-import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
+import software.amazon.smithy.rust.codegen.core.smithy.customize.CombinedCoreCodegenDecorator
+import software.amazon.smithy.rust.codegen.core.smithy.customize.CoreCodegenDecorator
 import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization
-import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization
-import software.amazon.smithy.rust.codegen.core.smithy.generators.ManifestCustomizations
 import software.amazon.smithy.rust.codegen.core.smithy.protocols.ProtocolMap
-import software.amazon.smithy.rust.codegen.core.util.deepMergeWith
 import java.util.ServiceLoader
 import java.util.logging.Logger
 
@@ -32,23 +28,7 @@ typealias ClientProtocolMap = ProtocolMap<ClientProtocolGenerator, ClientCodegen
  * AWS services. A different downstream customer may wish to add a different set of derive
  * attributes to the generated classes.
  */
-interface ClientCodegenDecorator {
-    /**
-     * The name of this [ClientCodegenDecorator], used for logging and debug information
-     */
-    val name: String
-
-    /**
-     * Enable a deterministic ordering to be applied, with the lowest numbered integrations being applied first
-     */
-    val order: Byte
-
-    /**
-     * Whether this decorator can be discovered on the classpath (defaults to true).
-     * This is intended to only be overridden for decorators written specifically for tests.
-     */
-    fun classpathDiscoverable(): Boolean = true
-
+interface ClientCodegenDecorator : CoreCodegenDecorator<ClientCodegenContext> {
     fun configCustomizations(
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<ConfigCustomization>,
@@ -60,24 +40,8 @@ interface ClientCodegenDecorator {
         baseCustomizations: List<OperationCustomization>,
     ): List<OperationCustomization> = baseCustomizations
 
-    fun libRsCustomizations(
-        codegenContext: ClientCodegenContext,
-        baseCustomizations: List<LibRsCustomization>,
-    ): List<LibRsCustomization> = baseCustomizations
-
-    /**
-     * Returns a map of Cargo.toml properties to change. For example, if a `homepage` needs to be
-     * added to the Cargo.toml `[package]` section, a `mapOf("package" to mapOf("homepage", "https://example.com"))`
-     * could be returned. Properties here overwrite the default properties.
-     */
-    fun crateManifestCustomizations(codegenContext: ClientCodegenContext): ManifestCustomizations = emptyMap()
-
-    fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {}
-
     fun protocols(serviceId: ShapeId, currentProtocols: ClientProtocolMap): ClientProtocolMap = currentProtocols
 
-    fun transformModel(service: ServiceShape, model: Model): Model = model
-
     fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> = listOf()
 }
 
@@ -86,8 +50,8 @@ interface ClientCodegenDecorator {
  *
  * This makes the actual concrete codegen simpler by not needing to deal with multiple separate decorators.
  */
-open class CombinedClientCodegenDecorator(decorators: List<ClientCodegenDecorator>) : ClientCodegenDecorator {
-    private val orderedDecorators = decorators.sortedBy { it.order }
+open class CombinedClientCodegenDecorator(decorators: List<ClientCodegenDecorator>) :
+    CombinedCoreCodegenDecorator<ClientCodegenContext, ClientCodegenDecorator>(decorators), ClientCodegenDecorator {
     override val name: String
         get() = "CombinedClientCodegenDecorator"
     override val order: Byte
@@ -96,59 +60,25 @@ open class CombinedClientCodegenDecorator(decorators: List<ClientCodegenDecorato
     override fun configCustomizations(
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<ConfigCustomization>,
-    ): List<ConfigCustomization> {
-        return orderedDecorators.foldRight(baseCustomizations) { decorator: ClientCodegenDecorator, customizations ->
-            decorator.configCustomizations(codegenContext, customizations)
-        }
+    ): List<ConfigCustomization> = combineCustomizations(baseCustomizations) { decorator, customizations ->
+        decorator.configCustomizations(codegenContext, customizations)
     }
 
     override fun operationCustomizations(
         codegenContext: ClientCodegenContext,
         operation: OperationShape,
         baseCustomizations: List<OperationCustomization>,
-    ): List<OperationCustomization> {
-        return orderedDecorators.foldRight(baseCustomizations) { decorator: ClientCodegenDecorator, customizations ->
-            decorator.operationCustomizations(codegenContext, operation, customizations)
-        }
+    ): List<OperationCustomization> = combineCustomizations(baseCustomizations) { decorator, customizations ->
+        decorator.operationCustomizations(codegenContext, operation, customizations)
     }
 
-    override fun libRsCustomizations(
-        codegenContext: ClientCodegenContext,
-        baseCustomizations: List<LibRsCustomization>,
-    ): List<LibRsCustomization> {
-        return orderedDecorators.foldRight(baseCustomizations) { decorator, customizations ->
-            decorator.libRsCustomizations(
-                codegenContext,
-                customizations,
-            )
-        }
-    }
-
-    override fun protocols(serviceId: ShapeId, currentProtocols: ClientProtocolMap): ClientProtocolMap {
-        return orderedDecorators.foldRight(currentProtocols) { decorator, protocolMap ->
+    override fun protocols(serviceId: ShapeId, currentProtocols: ClientProtocolMap): ClientProtocolMap =
+        combineCustomizations(currentProtocols) { decorator, protocolMap ->
             decorator.protocols(serviceId, protocolMap)
         }
-    }
-
-    override fun crateManifestCustomizations(codegenContext: ClientCodegenContext): ManifestCustomizations {
-        return orderedDecorators.foldRight(emptyMap()) { decorator, customizations ->
-            customizations.deepMergeWith(decorator.crateManifestCustomizations(codegenContext))
-        }
-    }
-
-    override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
-        return orderedDecorators.forEach { it.extras(codegenContext, rustCrate) }
-    }
-
-    override fun transformModel(service: ServiceShape, model: Model): Model {
-        return orderedDecorators.foldRight(model) { decorator, otherModel ->
-            decorator.transformModel(service, otherModel)
-        }
-    }
 
-    override fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> {
-        return orderedDecorators.flatMap { it.endpointCustomizations(codegenContext) }
-    }
+    override fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> =
+        addCustomizations { decorator -> decorator.endpointCustomizations(codegenContext) }
 
     companion object {
         fun fromClasspath(

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/CoreCodegenDecorator.kt

@@ -0,0 +1,147 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.core.smithy.customize
+
+import software.amazon.smithy.build.PluginContext
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
+import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization
+import software.amazon.smithy.rust.codegen.core.smithy.generators.ManifestCustomizations
+import software.amazon.smithy.rust.codegen.core.util.deepMergeWith
+import java.util.ServiceLoader
+import java.util.logging.Logger
+
+/**
+ * Represents the bare minimum for codegen plugin customization.
+ */
+interface CoreCodegenDecorator<CodegenContext> {
+    /**
+     * The name of this decorator, used for logging and debug information
+     */
+    val name: String
+
+    /**
+     * Enable a deterministic ordering to be applied, with the lowest numbered integrations being applied first
+     */
+    val order: Byte
+
+    /**
+     * Whether this decorator can be discovered on the classpath (defaults to true).
+     * This is intended to only be overridden for decorators written specifically for tests.
+     */
+    fun classpathDiscoverable(): Boolean = true
+
+    /**
+     * Hook to transform the Smithy model before codegen takes place.
+     */
+    fun transformModel(service: ServiceShape, model: Model): Model = model
+
+    /**
+     * Hook to add additional modules to the generated crate.
+     */
+    fun extras(codegenContext: CodegenContext, rustCrate: RustCrate) {}
+
+    /**
+     * Hook to customize the generated `Cargo.toml` file.
+     *
+     * Returns a map of Cargo.toml properties to change. For example, if a `homepage` needs to be
+     * added to the Cargo.toml `[package]` section, a `mapOf("package" to mapOf("homepage", "https://example.com"))`
+     * could be returned. Properties here overwrite the default properties.
+     */
+    fun crateManifestCustomizations(codegenContext: CodegenContext): ManifestCustomizations = emptyMap()
+
+    /**
+     * Hook to customize the generated `lib.rs` file.
+     */
+    fun libRsCustomizations(
+        codegenContext: CodegenContext,
+        baseCustomizations: List<LibRsCustomization>,
+    ): List<LibRsCustomization> = baseCustomizations
+}
+
+/**
+ * Implementations for combining decorators for the core customizations.
+ */
+abstract class CombinedCoreCodegenDecorator<CodegenContext, Decorator : CoreCodegenDecorator<CodegenContext>>(
+    decorators: List<Decorator>,
+) : CoreCodegenDecorator<CodegenContext> {
+    private val orderedDecorators = decorators.sortedBy { it.order }
+
+    final override fun libRsCustomizations(
+        codegenContext: CodegenContext,
+        baseCustomizations: List<LibRsCustomization>,
+    ): List<LibRsCustomization> =
+        combineCustomizations(baseCustomizations) { decorator, customizations ->
+            decorator.libRsCustomizations(codegenContext, customizations)
+        }
+
+    final override fun crateManifestCustomizations(codegenContext: CodegenContext): ManifestCustomizations =
+        combineCustomizations(emptyMap()) { decorator, customizations ->
+            customizations.deepMergeWith(decorator.crateManifestCustomizations(codegenContext))
+        }
+
+    final override fun extras(codegenContext: CodegenContext, rustCrate: RustCrate) {
+        return orderedDecorators.forEach { it.extras(codegenContext, rustCrate) }
+    }
+
+    final override fun transformModel(service: ServiceShape, model: Model): Model =
+        combineCustomizations(model) { decorator, otherModel ->
+            decorator.transformModel(service, otherModel)
+        }
+
+    /**
+     * Combines customizations from multiple ordered codegen decorators.
+     *
+     * Using this combinator allows for customizations to remove other customizations since the `mergeCustomizations`
+     * function can mutate the entire returned customization list.
+     */
+    protected fun <CombinedCustomizations> combineCustomizations(
+        /** Initial customizations. These will remain at the front of the list unless `mergeCustomizations` changes order. */
+        baseCustomizations: CombinedCustomizations,
+        /** A function that retrieves customizations from a decorator and combines them with the given customizations. */
+        mergeCustomizations: (Decorator, CombinedCustomizations) -> CombinedCustomizations,
+    ): CombinedCustomizations =
+        orderedDecorators.foldRight(baseCustomizations) { decorator, customizations ->
+            mergeCustomizations(decorator, customizations)
+        }
+
+    /**
+     * Combines customizations from multiple ordered codegen decorators in a purely additive way.
+     *
+     * Unlike `combineCustomizations`, customizations combined in this way cannot remove other customizations.
+     */
+    protected fun <Customization> addCustomizations(
+        /** Returns customizations from a decorator. */
+        getCustomizations: (Decorator) -> List<Customization>,
+    ): List<Customization> = orderedDecorators.flatMap(getCustomizations)
+
+    companion object {
+        /**
+         * Loads decorators of the given class from the classpath and filters them by the `classpathDiscoverable`
+         * method on `CoreCodegenDecorator`.
+         */
+        @JvmStatic
+        protected fun <CodegenContext, Decorator : CoreCodegenDecorator<CodegenContext>> decoratorsFromClasspath(
+            context: PluginContext,
+            decoratorClass: Class<Decorator>,
+            logger: Logger,
+            vararg extras: Decorator,
+        ): List<Decorator> {
+            val decorators = ServiceLoader.load(
+                decoratorClass,
+                context.pluginClassLoader.orElse(decoratorClass.classLoader),
+            )
+
+            val filteredDecorators = decorators.asSequence()
+                .onEach { logger.info("Discovered Codegen Decorator: ${it!!::class.java.name}") }
+                .filter { it!!.classpathDiscoverable() }
+                .onEach { logger.info("Adding Codegen Decorator: ${it!!::class.java.name}") }
+                .toList()
+            return filteredDecorators + extras
+        }
+    }
+}