From 8c1fd787bf17720856502365e522d4f84488a783 Mon Sep 17 00:00:00 2001
From: Ian Botsford <83236726+ianbotsf@users.noreply.github.com>
Date: Mon, 18 Mar 2024 11:21:03 -0700
Subject: [PATCH] fix: correctly codegen maps which use nested maps that
 contain enum keys (#1054)

---
 .../serde/DeserializeStructGenerator.kt       |  2 +-
 .../serde/DeserializeStructGeneratorTest.kt   | 78 +++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGenerator.kt
index ea91afc9f..7f1074525 100644
--- a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGenerator.kt
+++ b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGenerator.kt
@@ -302,7 +302,7 @@ open class DeserializeStructGenerator(
                         "val #L = #T<#T, #T#L>()",
                         memberName,
                         KotlinTypes.Collections.mutableMapOf,
-                        keySymbol,
+                        ctx.symbolProvider.toSymbol(mapShape.key),
                         ctx.symbolProvider.toSymbol(mapShape.value),
                         nullabilitySuffix(mapShape.isSparse),
                     )
diff --git a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGeneratorTest.kt b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGeneratorTest.kt
index abd959dda..f59165ed9 100644
--- a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGeneratorTest.kt
+++ b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/DeserializeStructGeneratorTest.kt
@@ -1077,6 +1077,84 @@ class DeserializeStructGeneratorTest {
         actual.shouldContainOnlyOnceWithDiff(expected)
     }
 
+    @Test
+    fun `it deserializes a structure containing a map with an enum key and a value that is a map with an enum key`() {
+        val model = (
+            modelPrefix + """
+            structure FooResponse {
+                payload: OuterKeyValuedMap
+            }
+
+            map OuterKeyValuedMap {
+                key: OuterKeyType,
+                value: InnerKeyValuedMap
+            }
+
+            @enum([
+                {
+                    value: "FOO",
+                },
+                {
+                    value: "BAR",
+                },
+            ])
+            string OuterKeyType
+
+            map InnerKeyValuedMap {
+                key: InnerKeyType,
+                value: String
+            }
+
+            @enum([
+                {
+                    value: "BAZ",
+                },
+                {
+                    value: "QUX",
+                },
+            ])
+            string InnerKeyType
+        """
+            ).toSmithyModel()
+
+        val expected = """
+            deserializer.deserializeStruct(OBJ_DESCRIPTOR) {
+                loop@while (true) {
+                    when (findNextFieldIndex()) {
+                        PAYLOAD_DESCRIPTOR.index -> builder.payload =
+                            deserializer.deserializeMap(PAYLOAD_DESCRIPTOR) {
+                                val map0 = mutableMapOf<OuterKeyType, Map<InnerKeyType, String>>()
+                                while (hasNextEntry()) {
+                                    val k0 = OuterKeyType.fromValue(key())
+                                    val v0 =
+                                        if (nextHasValue()) {
+                                            deserializer.deserializeMap(PAYLOAD_C0_DESCRIPTOR) {
+                                                val map1 = mutableMapOf<InnerKeyType, String>()
+                                                while (hasNextEntry()) {
+                                                    val k1 = InnerKeyType.fromValue(key())
+                                                    val v1 = if (nextHasValue()) { deserializeString() } else { deserializeNull(); continue }
+                                                    map1[k1] = v1
+                                                }
+                                                map1
+                                            }
+                                        } else { deserializeNull(); continue }
+
+                                    map0[k0] = v0
+                                }
+                                map0
+                            }
+                        null -> break@loop
+                        else -> skipValue()
+                    }
+                }
+            }
+        """.trimIndent()
+
+        val actual = codegenDeserializerForShape(model, "com.test#Foo")
+
+        actual.shouldContainOnlyOnceWithDiff(expected)
+    }
+
     @Test
     fun `it deserializes a structure containing a map of a union of primitive values`() {
         val model = (