fix: correctly codegen maps with enum keys (#1052)

Author: ianbotsf

.changes/a7f82a33-11f1-4184-97af-ff713e922dfc.json

@@ -0,0 +1,9 @@
+{
+    "id": "a7f82a33-11f1-4184-97af-ff713e922dfc",
+    "type": "bugfix",
+    "description": "⚠️ **IMPORTANT**: Fix codegen for map shapes which use string enums as map keys. See the [**Map key changes** breaking change announcement](https://github.com/awslabs/aws-sdk-kotlin/discussions/1258) for more details",
+    "issues": [
+        "awslabs/smithy-kotlin#1045"
+    ],
+    "requiresMinorVersionBump": true
+}

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

@@ -43,6 +43,7 @@ data class KotlinSettings(
     val sdkId: String,
     val build: BuildSettings = BuildSettings.Default,
     val api: ApiSettings = ApiSettings.Default,
+    val debug: Boolean = false,
 ) {
 
     /**
@@ -104,12 +105,14 @@ data class KotlinSettings(
             val sdkId = config.getStringMemberOrDefault(SDK_ID, serviceId.name)
             val build = config.getObjectMember(BUILD_SETTINGS)
             val api = config.getObjectMember(API_SETTINGS)
+            val debug = config.getBooleanMemberOrDefault("debug", false)
             return KotlinSettings(
                 serviceId,
                 PackageSettings(packageName, version, desc),
                 sdkId,
                 BuildSettings.fromNode(build),
                 ApiSettings.fromNode(api),
+                debug,
             )
         }
     }

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

@@ -150,6 +150,7 @@ class KotlinDelegator(
         val needsNewline = writers.containsKey(formattedFilename)
         val writer = writers.getOrPut(formattedFilename) {
             val kotlinWriter = KotlinWriter(namespace)
+            if (settings.debug) kotlinWriter.enableStackTraceComments(true)
 
             // Register all integrations [SectionWriterBindings] on the writer.
             integrations.forEach { integration ->

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

@@ -6,7 +6,6 @@ package software.amazon.smithy.kotlin.codegen.core
 
 import software.amazon.smithy.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
-import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
 import software.amazon.smithy.kotlin.codegen.lang.kotlinReservedWords
 import software.amazon.smithy.kotlin.codegen.model.*
 import software.amazon.smithy.kotlin.codegen.utils.dq
@@ -162,15 +161,18 @@ class KotlinSymbolProvider(private val model: Model, private val settings: Kotli
     }
 
     override fun mapShape(shape: MapShape): Symbol {
-        val reference = toSymbol(shape.value)
-        val valueSuffix = if (reference.isNullable) "?" else ""
-        val valueType = "${reference.name}$valueSuffix"
-        val fullyQualifiedValueType = "${reference.fullName}$valueSuffix"
+        val keyReference = toSymbol(shape.key)
+        val keyType = keyReference.name
+        val fullyQualifiedKeyType = keyReference.fullName
+
+        val valueReference = toSymbol(shape.value)
+        val valueSuffix = if (valueReference.isNullable) "?" else ""
+        val valueType = "${valueReference.name}$valueSuffix"
+        val fullyQualifiedValueType = "${valueReference.fullName}$valueSuffix"
 
-        val keyType = KotlinTypes.String.name
-        val fullyQualifiedKeyType = KotlinTypes.String.fullName
         return createSymbolBuilder(shape, "Map<$keyType, $valueType>")
-            .addReferences(reference)
+            .addReferences(keyReference)
+            .addReferences(valueReference)
             .putProperty(SymbolProperty.FULLY_QUALIFIED_NAME_HINT, "Map<$fullyQualifiedKeyType, $fullyQualifiedValueType>")
             .putProperty(SymbolProperty.MUTABLE_COLLECTION_FUNCTION, "mutableMapOf<$keyType, $valueType>")
             .putProperty(SymbolProperty.IMMUTABLE_COLLECTION_FUNCTION, "mapOf<$keyType, $valueType>")

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

@@ -188,7 +188,16 @@ class ShapeValueGenerator(
                     }
                     is MapShape -> {
                         memberShape = generator.model.expectShape(currShape.value.target)
-                        writer.writeInline("#S to ", keyNode.value)
+
+                        val keyTarget = generator.model.expectShape(currShape.key.target)
+                        if (keyTarget.isEnum) {
+                            val keySymbol = generator.symbolProvider.toSymbol(currShape.key)
+                            writer.writeInline("#T.fromValue(#S)", keySymbol, keyNode.value)
+                        } else {
+                            writer.writeInline("#S", keyNode.value)
+                        }
+
+                        writer.writeInline(" to ")
 
                         if (valueNode is NullNode) {
                             writer.write("null")

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGenerator.kt

@@ -5,6 +5,7 @@
 package software.amazon.smithy.kotlin.codegen.rendering.serde
 
 import software.amazon.smithy.codegen.core.CodegenException
+import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.kotlin.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
 import software.amazon.smithy.kotlin.codegen.model.*
@@ -145,9 +146,10 @@ open class DeserializeStructGenerator(
             .indent()
             .withBlock("deserializer.#T($descriptorName) {", "}", RuntimeTypes.Serde.deserializeMap) {
                 write(
-                    "val #L = #T<String, #T#L>()",
+                    "val #L = #T<#T, #T#L>()",
                     mutableCollectionName,
                     KotlinTypes.Collections.mutableMapOf,
+                    ctx.symbolProvider.toSymbol(targetShape.key),
                     ctx.symbolProvider.toSymbol(targetShape.value),
                     nullabilitySuffix(targetShape.isSparse),
                 )
@@ -168,6 +170,8 @@ open class DeserializeStructGenerator(
         nestingLevel: Int,
         parentMemberName: String,
     ) {
+        val keyShape = ctx.model.expectShape(mapShape.key.target)
+        val keySymbol = ctx.symbolProvider.toSymbol(keyShape)
         val elementShape = ctx.model.expectShape(mapShape.value.target)
         val isSparse = mapShape.isSparse
 
@@ -187,21 +191,47 @@ open class DeserializeStructGenerator(
             ShapeType.TIMESTAMP,
             ShapeType.ENUM,
             ShapeType.INT_ENUM,
-            -> renderEntry(elementShape, nestingLevel, isSparse, parentMemberName)
+            -> renderEntry(keyShape, keySymbol, elementShape, nestingLevel, isSparse, parentMemberName)
 
             ShapeType.SET,
             ShapeType.LIST,
-            -> renderListEntry(rootMemberShape, elementShape as CollectionShape, nestingLevel, isSparse, parentMemberName)
+            -> renderListEntry(
+                rootMemberShape,
+                keyShape,
+                keySymbol,
+                elementShape as CollectionShape,
+                nestingLevel,
+                isSparse,
+                parentMemberName,
+            )
+
+            ShapeType.MAP -> renderMapEntry(
+                rootMemberShape,
+                keyShape,
+                keySymbol,
+                elementShape as MapShape,
+                nestingLevel,
+                isSparse,
+                parentMemberName,
+            )
 
-            ShapeType.MAP -> renderMapEntry(rootMemberShape, elementShape as MapShape, nestingLevel, isSparse, parentMemberName)
             ShapeType.UNION,
             ShapeType.STRUCTURE,
-            -> renderNestedStructureEntry(elementShape, nestingLevel, isSparse, parentMemberName)
+            -> renderNestedStructureEntry(keyShape, keySymbol, elementShape, nestingLevel, isSparse, parentMemberName)
 
             else -> error("Unhandled type ${elementShape.type}")
         }
     }
 
+    private fun writeKeyVal(keyShape: Shape, keySymbol: Symbol, keyName: String) {
+        writer.writeInline("val $keyName = ")
+        if (keyShape.isEnum) {
+            writer.write("#T.fromValue(key())", keySymbol)
+        } else {
+            writer.write("key()")
+        }
+    }
+
     /**
      * Renders the deserialization of a nested structure contained in a map.  Example:
      *
@@ -212,6 +242,8 @@ open class DeserializeStructGenerator(
      * ```
      */
     private fun renderNestedStructureEntry(
+        keyShape: Shape,
+        keySymbol: Symbol,
         elementShape: Shape,
         nestingLevel: Int,
         isSparse: Boolean,
@@ -226,7 +258,7 @@ open class DeserializeStructGenerator(
             writer.addImport(symbol)
         }
 
-        writer.write("val $keyName = key()")
+        writeKeyVal(keyShape, keySymbol, keyName)
         writer.write("val $valueName = if (nextHasValue()) { $deserializerFn } else { deserializeNull()$populateNullValuePostfix }")
         writer.write("$parentMemberName[$keyName] = $valueName")
     }
@@ -247,6 +279,8 @@ open class DeserializeStructGenerator(
      */
     private fun renderMapEntry(
         rootMemberShape: MemberShape,
+        keyShape: Shape,
+        keySymbol: Symbol,
         mapShape: MapShape,
         nestingLevel: Int,
         isSparse: Boolean,
@@ -260,14 +294,15 @@ open class DeserializeStructGenerator(
         val memberName = nextNestingLevel.variableNameFor(NestedIdentifierType.MAP)
         val collectionReturnExpression = collectionReturnExpression(rootMemberShape, memberName)
 
-        writer.write("val $keyName = key()")
+        writeKeyVal(keyShape, keySymbol, keyName)
         writer.withBlock("val $valueName =", "") {
             withBlock("if (nextHasValue()) {", "} else { deserializeNull()$populateNullValuePostfix }") {
                 withBlock("deserializer.#T($descriptorName) {", "}", RuntimeTypes.Serde.deserializeMap) {
                     write(
-                        "val #L = #T<String, #T#L>()",
+                        "val #L = #T<#T, #T#L>()",
                         memberName,
                         KotlinTypes.Collections.mutableMapOf,
+                        keySymbol,
                         ctx.symbolProvider.toSymbol(mapShape.value),
                         nullabilitySuffix(mapShape.isSparse),
                     )
@@ -298,6 +333,8 @@ open class DeserializeStructGenerator(
      */
     private fun renderListEntry(
         rootMemberShape: MemberShape,
+        keyShape: Shape,
+        keySymbol: Symbol,
         collectionShape: CollectionShape,
         nestingLevel: Int,
         isSparse: Boolean,
@@ -311,7 +348,7 @@ open class DeserializeStructGenerator(
         val memberName = nextNestingLevel.variableNameFor(NestedIdentifierType.COLLECTION)
         val collectionReturnExpression = collectionReturnExpression(rootMemberShape, memberName)
 
-        writer.write("val $keyName = key()")
+        writeKeyVal(keyShape, keySymbol, keyName)
         writer.withBlock("val $valueName =", "") {
             withBlock("if (nextHasValue()) {", "} else { deserializeNull()$populateNullValuePostfix }") {
                 withBlock("deserializer.#T($descriptorName) {", "}", RuntimeTypes.Serde.deserializeList) {
@@ -340,13 +377,20 @@ open class DeserializeStructGenerator(
      * map0[k0] = el0
      * ```
      */
-    private fun renderEntry(elementShape: Shape, nestingLevel: Int, isSparse: Boolean, parentMemberName: String) {
+    private fun renderEntry(
+        keyShape: Shape,
+        keySymbol: Symbol,
+        elementShape: Shape,
+        nestingLevel: Int,
+        isSparse: Boolean,
+        parentMemberName: String,
+    ) {
         val deserializerFn = deserializerForShape(elementShape)
         val keyName = nestingLevel.variableNameFor(NestedIdentifierType.KEY)
         val valueName = nestingLevel.variableNameFor(NestedIdentifierType.VALUE)
         val populateNullValuePostfix = if (isSparse) "" else "; continue"
 
-        writer.write("val $keyName = key()")
+        writeKeyVal(keyShape, keySymbol, keyName)
         writer.write("val $valueName = if (nextHasValue()) { $deserializerFn } else { deserializeNull()$populateNullValuePostfix }")
         writer.write("$parentMemberName[$keyName] = $valueName")
     }
@@ -476,9 +520,10 @@ open class DeserializeStructGenerator(
 
         writer.withBlock("val $elementName = deserializer.#T($descriptorName) {", "}", RuntimeTypes.Serde.deserializeMap) {
             write(
-                "val #L = #T<String, #T#L>()",
+                "val #L = #T<#T, #T#L>()",
                 mapName,
                 KotlinTypes.Collections.mutableMapOf,
+                ctx.symbolProvider.toSymbol(mapShape.key),
                 ctx.symbolProvider.toSymbol(mapShape.value),
                 nullabilitySuffix(mapShape.isSparse),
             )