Make RustReservedWordsSymbolProvider configurable

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy
+
+import software.amazon.smithy.rust.codegen.core.rustlang.RustReservedWordConfig
+import software.amazon.smithy.rust.codegen.core.smithy.generators.StructureGenerator
+import software.amazon.smithy.rust.codegen.core.smithy.generators.UnionGenerator
+
+val ClientReservedWords = RustReservedWordConfig(
+    structMemberMap = StructureGenerator.structMemberMap +
+        mapOf(
+            "send" to "send_value",
+            // To avoid conflicts with the `make_operation` and `presigned` functions on generated inputs
+            "make_operation" to "make_operation_value",
+            "presigned" to "presigned_value",
+            "customize" to "customize_value",
+            // To avoid conflicts with the error metadata `meta` field
+            "meta" to "meta_value",
+        ),
+    unionMemberMap = mapOf(
+        // Unions contain an `Unknown` variant. This exists to support parsing data returned from the server
+        // that represent union variants that have been added since this SDK was generated.
+        UnionGenerator.UnknownVariantName to "${UnionGenerator.UnknownVariantName}Value",
+        "${UnionGenerator.UnknownVariantName}Value" to "${UnionGenerator.UnknownVariantName}Value_",
+    ),
+    enumMemberMap = mapOf(
+        // Unknown is used as the name of the variant containing unexpected values
+        "Unknown" to "UnknownValue",
+        // Real models won't end in `_` so it's safe to stop here
+        "UnknownValue" to "UnknownValue_",
+    ),
+)

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

@@ -84,6 +84,6 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() {
                 .let { StreamingShapeMetadataProvider(it, model) }
                 // Rename shapes that clash with Rust reserved words & and other SDK specific features e.g. `send()` cannot
                 // be the name of an operation input
-                .let { RustReservedWordSymbolProvider(it, model) }
+                .let { RustReservedWordSymbolProvider(it, model, ClientReservedWords) }
     }
 }

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustReservedWords.kt

@@ -17,13 +17,24 @@ import software.amazon.smithy.model.shapes.UnionShape
 import software.amazon.smithy.model.traits.EnumTrait
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.WrappingSymbolProvider
-import software.amazon.smithy.rust.codegen.core.smithy.generators.UnionGenerator
 import software.amazon.smithy.rust.codegen.core.smithy.renamedFrom
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
 import software.amazon.smithy.rust.codegen.core.util.letIf
 
-class RustReservedWordSymbolProvider(private val base: RustSymbolProvider, private val model: Model) :
-    WrappingSymbolProvider(base) {
+data class RustReservedWordConfig(
+    /** Map of struct member names that should get renamed */
+    val structMemberMap: Map<String, String>,
+    /** Map of union member names that should get renamed */
+    val unionMemberMap: Map<String, String>,
+    /** Map of enum member names that should get renamed */
+    val enumMemberMap: Map<String, String>,
+)
+
+class RustReservedWordSymbolProvider(
+    private val base: RustSymbolProvider,
+    private val model: Model,
+    private val config: RustReservedWordConfig,
+) : WrappingSymbolProvider(base) {
     private val internal =
         ReservedWordSymbolProvider.builder().symbolProvider(base)
             .nameReservedWords(RustReservedWords)
@@ -35,34 +46,19 @@ class RustReservedWordSymbolProvider(private val base: RustSymbolProvider, priva
         val reservedWordReplacedName = internal.toMemberName(shape)
         val container = model.expectShape(shape.container)
         return when {
-            container is StructureShape -> when (baseName) {
-                "build" -> "build_value"
-                "builder" -> "builder_value"
-                "default" -> "default_value"
-                "send" -> "send_value"
-                // To avoid conflicts with the `make_operation` and `presigned` functions on generated inputs
-                "make_operation" -> "make_operation_value"
-                "presigned" -> "presigned_value"
-                "customize" -> "customize_value"
-                // To avoid conflicts with the error metadata `meta` field
-                "meta" -> "meta_value"
-                else -> reservedWordReplacedName
+            container is StructureShape -> when (val mapped = config.structMemberMap[baseName]) {
+                null -> reservedWordReplacedName
+                else -> mapped
             }
 
-            container is UnionShape -> when (baseName) {
-                // Unions contain an `Unknown` variant. This exists to support parsing data returned from the server
-                // that represent union variants that have been added since this SDK was generated.
-                UnionGenerator.UnknownVariantName -> "${UnionGenerator.UnknownVariantName}Value"
-                "${UnionGenerator.UnknownVariantName}Value" -> "${UnionGenerator.UnknownVariantName}Value_"
-                else -> reservedWordReplacedName
+            container is UnionShape -> when (val mapped = config.unionMemberMap[baseName]) {
+                null -> reservedWordReplacedName
+                else -> mapped
             }
 
-            container is EnumShape || container.hasTrait<EnumTrait>() -> when (baseName) {
-                // Unknown is used as the name of the variant containing unexpected values
-                "Unknown" -> "UnknownValue"
-                // Real models won't end in `_` so it's safe to stop here
-                "UnknownValue" -> "UnknownValue_"
-                else -> reservedWordReplacedName
+            container is EnumShape || container.hasTrait<EnumTrait>() -> when (val mapped = config.enumMemberMap[baseName]) {
+                null -> reservedWordReplacedName
+                else -> mapped
             }
 
             else -> error("unexpected container: $container")

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/StructureGenerator.kt

@@ -60,6 +60,15 @@ open class StructureGenerator(
     private val shape: StructureShape,
     private val customizations: List<StructureCustomization>,
 ) {
+    companion object {
+        /** Reserved struct member names */
+        val structMemberMap: Map<String, String> = mapOf(
+            "build" to "build_value",
+            "builder" to "builder_value",
+            "default" to "default_value",
+        )
+    }
+
     private val errorTrait = shape.getTrait<ErrorTrait>()
     protected val members: List<MemberShape> = shape.allMembers.values.toList()
     private val accessorMembers: List<MemberShape> = when (errorTrait) {

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/testutil/TestHelpers.kt

@@ -16,6 +16,7 @@ import software.amazon.smithy.model.shapes.UnionShape
 import software.amazon.smithy.model.traits.ErrorTrait
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
 import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
+import software.amazon.smithy.rust.codegen.core.rustlang.RustReservedWordConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.RustReservedWordSymbolProvider
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.implBlock
@@ -113,12 +114,21 @@ fun String.asSmithyModel(sourceLocation: String? = null, smithyVersion: String =
 }
 
 // Intentionally only visible to codegen-core since the other modules have their own symbol providers
-internal fun testSymbolProvider(model: Model): RustSymbolProvider = SymbolVisitor(
+internal fun testSymbolProvider(
+    model: Model,
+    rustReservedWordConfig: RustReservedWordConfig? = null,
+): RustSymbolProvider = SymbolVisitor(
     model,
     ServiceShape.builder().version("test").id("test#Service").build(),
     TestSymbolVisitorConfig,
 ).let { BaseSymbolMetadataProvider(it, model, additionalAttributes = listOf(Attribute.NonExhaustive)) }
-    .let { RustReservedWordSymbolProvider(it, model) }
+    .let {
+        RustReservedWordSymbolProvider(
+            it,
+            model,
+            rustReservedWordConfig ?: RustReservedWordConfig(emptyMap(), emptyMap(), emptyMap()),
+        )
+    }
 
 // Intentionally only visible to codegen-core since the other modules have their own contexts
 internal fun testCodegenContext(

codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustReservedWordsTest.kt

@@ -21,18 +21,108 @@ import software.amazon.smithy.rust.codegen.core.util.lookup
 internal class RustReservedWordSymbolProviderTest {
     private class TestSymbolProvider(model: Model) :
         WrappingSymbolProvider(SymbolVisitor(model, null, TestSymbolVisitorConfig))
+    private val emptyConfig = RustReservedWordConfig(emptyMap(), emptyMap(), emptyMap())
 
     @Test
     fun `structs are escaped`() {
         val model = """
             namespace test
             structure Self {}
         """.asSmithyModel()
-        val provider = RustReservedWordSymbolProvider(TestSymbolProvider(model), model)
+        val provider = RustReservedWordSymbolProvider(TestSymbolProvider(model), model, emptyConfig)
         val symbol = provider.toSymbol(model.lookup("test#Self"))
         symbol.name shouldBe "SelfValue"
     }
 
+    private fun mappingTest(config: RustReservedWordConfig, model: Model, id: String, test: (String) -> Unit) {
+        val provider = RustReservedWordSymbolProvider(TestSymbolProvider(model), model, config)
+        val symbol = provider.toMemberName(model.lookup("test#Container\$$id"))
+        test(symbol)
+    }
+
+    @Test
+    fun `structs member names are mapped via config`() {
+        val config = emptyConfig.copy(
+            structMemberMap = mapOf(
+                "name_to_map" to "mapped_name",
+                "NameToMap" to "MappedName",
+            ),
+        )
+        var model = """
+            namespace test
+            structure Container {
+                name_to_map: String
+            }
+        """.asSmithyModel()
+        mappingTest(config, model, "name_to_map") { memberName ->
+            memberName shouldBe "mapped_name"
+        }
+
+        model = """
+            namespace test
+            enum Container {
+                NameToMap = "NameToMap"
+            }
+        """.asSmithyModel(smithyVersion = "2.0")
+        mappingTest(config, model, "NameToMap") { memberName ->
+            // Container was not a struct, so the field keeps its old name
+            memberName shouldBe "NameToMap"
+        }
+
+        model = """
+            namespace test
+            union Container {
+                NameToMap: String
+            }
+        """.asSmithyModel()
+        mappingTest(config, model, "NameToMap") { memberName ->
+            // Container was not a struct, so the field keeps its old name
+            memberName shouldBe "NameToMap"
+        }
+    }
+
+    @Test
+    fun `union member names are mapped via config`() {
+        val config = emptyConfig.copy(
+            unionMemberMap = mapOf(
+                "name_to_map" to "mapped_name",
+                "NameToMap" to "MappedName",
+            ),
+        )
+
+        var model = """
+            namespace test
+            union Container {
+                NameToMap: String
+            }
+        """.asSmithyModel()
+        mappingTest(config, model, "NameToMap") { memberName ->
+            memberName shouldBe "MappedName"
+        }
+
+        model = """
+            namespace test
+            structure Container {
+                name_to_map: String
+            }
+        """.asSmithyModel()
+        mappingTest(config, model, "name_to_map") { memberName ->
+            // Container was not a union, so the field keeps its old name
+            memberName shouldBe "name_to_map"
+        }
+
+        model = """
+            namespace test
+            enum Container {
+                NameToMap = "NameToMap"
+            }
+        """.asSmithyModel(smithyVersion = "2.0")
+        mappingTest(config, model, "NameToMap") { memberName ->
+            // Container was not a union, so the field keeps its old name
+            memberName shouldBe "NameToMap"
+        }
+    }
+
     @Test
     fun `member names are escaped`() {
         val model = """
@@ -41,7 +131,7 @@ internal class RustReservedWordSymbolProviderTest {
                 async: String
             }
         """.asSmithyModel()
-        val provider = RustReservedWordSymbolProvider(TestSymbolProvider(model), model)
+        val provider = RustReservedWordSymbolProvider(TestSymbolProvider(model), model, emptyConfig)
         provider.toMemberName(
             MemberShape.builder().id("namespace#container\$async").target("namespace#Integer").build(),
         ) shouldBe "r##async"
@@ -57,7 +147,15 @@ internal class RustReservedWordSymbolProviderTest {
             namespace foo
             @enum([{ name: "dontcare", value: "dontcare" }]) string Container
         """.asSmithyModel()
-        val provider = RustReservedWordSymbolProvider(TestSymbolProvider(model), model)
+        val provider = RustReservedWordSymbolProvider(
+            TestSymbolProvider(model), model,
+            config = emptyConfig.copy(
+                enumMemberMap = mapOf(
+                    "Unknown" to "UnknownValue",
+                    "UnknownValue" to "UnknownValue_",
+                ),
+            ),
+        )
 
         fun expectEnumRename(original: String, expected: MaybeRenamed) {
             val symbol = provider.toSymbol(

codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGeneratorTest.kt

@@ -14,6 +14,7 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.traits.EnumTrait
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute.Companion.AllowDeprecated
+import software.amazon.smithy.rust.codegen.core.rustlang.RustReservedWordConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
@@ -30,6 +31,12 @@ import software.amazon.smithy.rust.codegen.core.util.lookup
 import software.amazon.smithy.rust.codegen.core.util.orNull
 
 class EnumGeneratorTest {
+    private val rustReservedWordConfig = RustReservedWordConfig(
+        enumMemberMap = mapOf("Unknown" to "UnknownValue"),
+        structMemberMap = emptyMap(),
+        unionMemberMap = emptyMap(),
+    )
+
     @Nested
     inner class EnumMemberModelTests {
         private val testModel = """
@@ -47,7 +54,7 @@ class EnumGeneratorTest {
             ])
             string EnumWithUnknown
         """.asSmithyModel()
-        private val symbolProvider = testSymbolProvider(testModel)
+        private val symbolProvider = testSymbolProvider(testModel, rustReservedWordConfig = rustReservedWordConfig)
 
         private val enumTrait = testModel.lookup<StringShape>("test#EnumWithUnknown").expectTrait<EnumTrait>()
 
@@ -276,7 +283,7 @@ class EnumGeneratorTest {
             """.asSmithyModel()
 
             val shape = model.lookup<StringShape>("test#SomeEnum")
-            val provider = testSymbolProvider(model)
+            val provider = testSymbolProvider(model, rustReservedWordConfig = rustReservedWordConfig)
             val project = TestWorkspace.testProject(provider)
             project.moduleFor(shape) {
                 renderEnum(model, provider, shape)