refactor(codegen): restructure operation serialization and deserialization (#579)

Author: aajtodd

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

@@ -124,20 +124,18 @@ class CodegenVisitor(context: PluginContext) : ShapeVisitor.Default<Unit>() {
                 writers
             )
 
-            LOGGER.info("[${service.id}] Generating serde for protocol $protocol")
-            generateSerializers(ctx)
-            generateDeserializers(ctx)
-
             LOGGER.info("[${service.id}] Generating unit tests for protocol $protocol")
             generateProtocolUnitTests(ctx)
 
             LOGGER.info("[${service.id}] Generating service client for protocol $protocol")
             generateProtocolClient(ctx)
         }
 
+        writers.finalize()
+
         if (settings.build.generateDefaultBuildFiles) {
             val dependencies = writers.dependencies
-                .map { it.properties["dependency"] as KotlinDependency }
+                .mapNotNull { it.properties["dependency"] as? KotlinDependency }
                 .distinct()
             writeGradleBuild(settings, fileManifest, dependencies)
         }

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

@@ -41,7 +41,12 @@ data class KotlinSettings(
     /**
      * Configuration elements specific to the service's package namespace, version, and description.
      */
-    data class PackageSettings(val name: String, val version: String, val description: String? = null)
+    data class PackageSettings(val name: String, val version: String, val description: String? = null) {
+        /**
+         * Derive a subpackage namespace from the root package name
+         */
+        fun subpackage(subpackageName: String): String = "$name.$subpackageName"
+    }
 
     /**
      * Get the corresponding [ServiceShape] from a model.

smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/KotlinDelegator.kt

@@ -5,12 +5,10 @@
 package software.amazon.smithy.kotlin.codegen.core
 
 import software.amazon.smithy.build.FileManifest
-import software.amazon.smithy.codegen.core.Symbol
-import software.amazon.smithy.codegen.core.SymbolDependency
-import software.amazon.smithy.codegen.core.SymbolProvider
-import software.amazon.smithy.codegen.core.SymbolReference
+import software.amazon.smithy.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.SymbolProperty
 import software.amazon.smithy.kotlin.codegen.utils.namespaceToPath
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.Shape
@@ -34,6 +32,30 @@ class KotlinDelegator(
     // Tracks dependencies for source not provided by codegen that may reside in the service source tree.
     val runtimeDependencies: MutableList<SymbolDependency> = mutableListOf()
 
+    /**
+     * Finalize renders any generated dependencies that were used and should be called just before [flushWriters]
+     */
+    fun finalize() {
+        // render generated "dependencies"
+        val writtenDependencies = mutableSetOf<String>()
+
+        // since generated dependencies can further mutate the set of generated dependencies this has
+        // to be unwound until we have nothing left to process
+        while (unprocessedDependencies(writtenDependencies).isNotEmpty()) {
+            unprocessedDependencies(writtenDependencies).forEach { generated ->
+                writtenDependencies.add(generated.fullName)
+                val writer = checkoutWriter(generated.definitionFile, generated.namespace)
+                writer.apply(generated.renderer)
+            }
+        }
+    }
+
+    private fun unprocessedDependencies(writtenDependencies: Set<String>) =
+        dependencies
+            .mapNotNull { it.properties[SymbolProperty.GENERATED_DEPENDENCY] as? GeneratedDependency }
+            .filterNot { writtenDependencies.contains(it.fullName) }
+            .distinctBy { it.fullName }
+
     /**
      * Writes all pending writers to disk and then clears them out.
      */
@@ -155,3 +177,37 @@ class KotlinDelegator(
         return writer
     }
 }
+
+/**
+ * A pseudo dependency on a snippet of code. A generated dependency is usually a symbol that is required
+ * by some other piece of code and must be generated.
+ *
+ * GeneratedDependency should not be instantiated directly, rather, it should be constructed by setting
+ * the [software.amazon.smithy.kotlin.codegen.model.SymbolBuilder.renderBy] field when creating a [Symbol]
+ *
+ * Generated dependencies are created in the definition file of the [Symbol] they generate. They are
+ * deduplicated by their fully qualified name in [KotlinDelegator] during codegen when writers are finalized.
+ */
+internal data class GeneratedDependency(
+    val name: String,
+    val namespace: String,
+    val definitionFile: String,
+    val renderer: SymbolRenderer
+) : SymbolDependencyContainer {
+    /**
+     * Fully qualified name
+     */
+    val fullName: String
+        get() = "$namespace.$name"
+
+    override fun getDependencies(): List<SymbolDependency> {
+        val symbolDep = SymbolDependency.builder()
+            .dependencyType("generated")
+            .version("n/a")
+            .packageName(fullName)
+            .putProperty(SymbolProperty.GENERATED_DEPENDENCY, this)
+            .build()
+
+        return listOf(symbolDep)
+    }
+}

smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/KotlinWriter.kt

@@ -25,6 +25,11 @@ import software.amazon.smithy.model.traits.EnumDefinition
 import software.amazon.smithy.utils.CodeWriter
 import java.util.function.BiFunction
 
+/**
+ * A function that renders a symbol to the given writer
+ */
+typealias SymbolRenderer = (KotlinWriter) -> Unit
+
 /**
  * Provides capability of writing Kotlin source code. An instance of a KotlinWriter corresponds to a file emitted
  * from codegen.
@@ -45,33 +50,39 @@ class KotlinWriter(
         expressionStart = '#'
 
         // type name: `Foo`
-        putFormatter('T', KotlinSymbolFormatter { symbol -> fullyQualifiedSymbols.contains(symbol.toFullyQualifiedSymbolName()) })
+        putFormatter('T', KotlinSymbolFormatter(this) { symbol -> fullyQualifiedSymbols.contains(symbol.toFullyQualifiedSymbolName()) })
         // fully qualified type: `aws.sdk.kotlin.model.Foo`
-        putFormatter('Q', KotlinSymbolFormatter { true })
+        putFormatter('Q', KotlinSymbolFormatter(this) { true })
 
         // like `T` but with nullability information: `aws.sdk.kotlin.model.Foo?`. This is mostly useful
         // when formatting properties
-        putFormatter('P', KotlinPropertyFormatter())
+        putFormatter('P', KotlinPropertyFormatter(this))
 
         // like `P` but fully qualified
-        putFormatter('F', KotlinPropertyFormatter(fullyQualifiedNames = true))
+        putFormatter('F', KotlinPropertyFormatter(this, fullyQualifiedNames = true))
 
         // like `P` but with default set (if applicable): `aws.sdk.kotlin.model.Foo = 1`
-        putFormatter('D', KotlinPropertyFormatter(setDefault = true))
+        putFormatter('D', KotlinPropertyFormatter(this, setDefault = true))
 
         // like `D` but fully qualified
-        putFormatter('E', KotlinPropertyFormatter(setDefault = true, fullyQualifiedNames = true))
+        putFormatter('E', KotlinPropertyFormatter(this, setDefault = true, fullyQualifiedNames = true))
 
         // Pass a function receiving a [KotlinWriter] to generate an inline value
         putFormatter('W', InlineKotlinWriterFormatter(this))
     }
 
+    /**
+     * Import [symbol] into the current file. If the symbol resides in the same package as the current writer
+     * then only the dependencies will be processed and no import statement generated.
+     *
+     * @param symbol the symbol to generate an import statement for
+     * @param alias an alias name to give to the imported symbol (e.g. `import foo.bar.baz as Quux`)
+     */
     fun addImport(symbol: Symbol, alias: String = symbol.name): KotlinWriter {
         // don't import built-in symbols
         if (symbol.isBuiltIn) return this
 
-        // always add dependencies
-        dependencies.addAll(symbol.dependencies)
+        addDepsRecursively(symbol)
 
         // only add imports for symbols in a different namespace
         if (symbol.namespace.isNotEmpty() && symbol.namespace != fullPackageName) {
@@ -108,6 +119,12 @@ class KotlinWriter(
         }
     }
 
+    private fun addDepsRecursively(symbol: Symbol) {
+        // always add dependencies
+        dependencies.addAll(symbol.dependencies)
+        symbol.references.forEach { addDepsRecursively(it.symbol) }
+    }
+
     /**
      * Directly add an import
      */
@@ -251,13 +268,17 @@ fun KotlinWriter.addImport(vararg imports: Iterable<Symbol>): KotlinWriter {
 
 /**
  * Implements Kotlin symbol formatting for the `#T` and `#Q` formatter(s)
+ * NOTE: That the symbol will automatically be imported.
  */
 private class KotlinSymbolFormatter(
+    private val writer: KotlinWriter,
     private val fullyQualifiedPredicate: (Symbol) -> Boolean = { false },
 ) : BiFunction<Any, String, String> {
     override fun apply(type: Any, indent: String): String {
         when (type) {
             is Symbol -> {
+                // writer will omit unnecessary same package imports and dedupe
+                writer.addImport(type)
                 return if (fullyQualifiedPredicate(type)) type.fullName else type.name
             }
             else -> throw CodegenException("Invalid type provided for #T. Expected a Symbol, but found `$type`")
@@ -269,6 +290,7 @@ private class KotlinSymbolFormatter(
  * Implements Kotlin symbol formatting for the `#D` and `#P` formatter(s)
  */
 class KotlinPropertyFormatter(
+    private val writer: KotlinWriter,
     // set defaults
     private val setDefault: Boolean = false,
     // format with nullability `?`
@@ -279,6 +301,7 @@ class KotlinPropertyFormatter(
     override fun apply(type: Any, indent: String): String {
         when (type) {
             is Symbol -> {
+                writer.addImport(type)
                 var formatted = if (fullyQualifiedNames) type.fullName else type.name
                 if (includeNullability && type.isBoxed) {
                     formatted += "?"

smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/SymbolBuilder.kt

@@ -7,7 +7,9 @@ package software.amazon.smithy.kotlin.codegen.model
 import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.codegen.core.SymbolDependencyContainer
 import software.amazon.smithy.codegen.core.SymbolReference
+import software.amazon.smithy.kotlin.codegen.core.GeneratedDependency
 import software.amazon.smithy.kotlin.codegen.core.KotlinDependency
+import software.amazon.smithy.kotlin.codegen.core.SymbolRenderer
 
 @DslMarker
 annotation class SymbolDsl
@@ -26,6 +28,11 @@ open class SymbolBuilder {
     var declarationFile: String? = null
     var defaultValue: String? = null
 
+    /**
+     * Register the function (dependency) responsible for rendering this symbol.
+     */
+    var renderBy: SymbolRenderer? = null
+
     val dependencies: MutableSet<SymbolDependencyContainer> = mutableSetOf()
     val references: MutableList<SymbolReference> = mutableListOf()
 
@@ -76,6 +83,15 @@ open class SymbolBuilder {
         declarationFile?.let { builder.declarationFile(it) }
         definitionFile?.let { builder.definitionFile(it) }
         defaultValue?.let { builder.defaultValue(it) }
+
+        if (renderBy != null) {
+            checkNotNull(definitionFile) { "a rendered dependency must declare a definition file!" }
+            checkNotNull(namespace) { "a rendered dependency must declare a namespace" }
+            // abuse dependencies to get the delegator to eventually render this
+            val generatedDep = GeneratedDependency(name!!, namespace!!, definitionFile!!, renderBy!!)
+            dependency(generatedDep)
+        }
+
         dependencies.forEach { builder.addDependency(it) }
         references.forEach { builder.addReference(it) }
 

smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/SymbolExt.kt

@@ -32,6 +32,9 @@ object SymbolProperty {
 
     // Entry type for Maps
     const val ENTRY_EXPRESSION: String = "entryExpression"
+
+    // Pseudo dependency on a snippet of code
+    const val GENERATED_DEPENDENCY: String = "generatedDependency"
 }
 
 /**

smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/knowledge/SerdeIndex.kt

@@ -47,6 +47,22 @@ class SerdeIndex(private val model: Model) : KnowledgeIndex {
         return walkNestedShapesRequiringSerde(model, topLevelMembers)
     }
 
+    /**
+     * Find and return the set of shapes reachable from the given shape that would require a "document" serializer.
+     * @return The set of shapes that require a serializer implementation
+     */
+    fun requiresDocumentSerializer(shape: Shape): Set<Shape> =
+        when (shape) {
+            is OperationShape -> requiresDocumentSerializer(listOf(shape))
+            else -> {
+                val topLevelMembers = shape.members()
+                    .map { model.expectShape(it.target) }
+                    .filter { it.isStructureShape || it.isUnionShape || it is CollectionShape || it.isMapShape }
+                    .toSet()
+                walkNestedShapesRequiringSerde(model, topLevelMembers)
+            }
+        }
+
     /**
      * Find and return the set of shapes that are not operation outputs but do require a deserializer
      *
@@ -83,6 +99,22 @@ class SerdeIndex(private val model: Model) : KnowledgeIndex {
 
         return walkNestedShapesRequiringSerde(model, topLevelMembers)
     }
+
+    /**
+     * Find and return the set of shapes reachable from the given shape that would require a "document" deserializer.
+     * @return The set of shapes that require a deserializer implementation
+     */
+    fun requiresDocumentDeserializer(shape: Shape): Set<Shape> =
+        when (shape) {
+            is OperationShape -> requiresDocumentDeserializer(listOf(shape))
+            else -> {
+                val topLevelMembers = shape.members()
+                    .map { model.expectShape(it.target) }
+                    .filter { it.isStructureShape || it.isUnionShape || it is CollectionShape || it.isMapShape }
+                    .toMutableSet()
+                walkNestedShapesRequiringSerde(model, topLevelMembers)
+            }
+        }
 }
 
 private fun walkNestedShapesRequiringSerde(model: Model, shapes: Set<Shape>): Set<Shape> {