Remove toEnumVariantName from RustSymbolProvider (smithy-lang#2377)

The `toEnumVariantName` function existed on symbol provider to work around enum definitions not being shapes. In the future when we refactor to use `EnumShape` instead of `EnumTrait`, there will be `MemberShape`s for each enum member. This change incrementally moves us to that future by creating fake `MemberShape`s in the enum generator from the enum definition.
thomas-k-cameron · Feb 22, 2023 · 7cffe14 · 7cffe14
1 parent da26405
commit 7cffe14
Show file tree

Hide file tree

Showing 11 changed files with 144 additions and 107 deletions.
diff --git a/...lin/software/amazon/smithy/rust/codegen/client/smithy/StreamingShapeSymbolProviderTest.kt b/...lin/software/amazon/smithy/rust/codegen/client/smithy/StreamingShapeSymbolProviderTest.kt
@@ -9,8 +9,10 @@ import io.kotest.matchers.shouldBe
 import org.junit.jupiter.api.Test
 import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.rust.codegen.client.testutil.testSymbolProvider
+import software.amazon.smithy.rust.codegen.core.rustlang.RustType
 import software.amazon.smithy.rust.codegen.core.smithy.Default
 import software.amazon.smithy.rust.codegen.core.smithy.defaultValue
+import software.amazon.smithy.rust.codegen.core.smithy.rustType
 import software.amazon.smithy.rust.codegen.core.smithy.transformers.OperationNormalizer
 import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
 import software.amazon.smithy.rust.codegen.core.util.lookup
@@ -38,8 +40,18 @@ internal class StreamingShapeSymbolProviderTest {
         // "doing the right thing"
         val modelWithOperationTraits = OperationNormalizer.transform(model)
         val symbolProvider = testSymbolProvider(modelWithOperationTraits)
-        symbolProvider.toSymbol(modelWithOperationTraits.lookup<MemberShape>("test.synthetic#GenerateSpeechOutput\$data")).name shouldBe ("ByteStream")
-        symbolProvider.toSymbol(modelWithOperationTraits.lookup<MemberShape>("test.synthetic#GenerateSpeechInput\$data")).name shouldBe ("ByteStream")
+        modelWithOperationTraits.lookup<MemberShape>("test.synthetic#GenerateSpeechOutput\$data").also { shape ->
+            symbolProvider.toSymbol(shape).also { symbol ->
+                symbol.name shouldBe "data"
+                symbol.rustType() shouldBe RustType.Opaque("ByteStream", "aws_smithy_http::byte_stream")
+            }
+        }
+        modelWithOperationTraits.lookup<MemberShape>("test.synthetic#GenerateSpeechInput\$data").also { shape ->
+            symbolProvider.toSymbol(shape).also { symbol ->
+                symbol.name shouldBe "data"
+                symbol.rustType() shouldBe RustType.Opaque("ByteStream", "aws_smithy_http::byte_stream")
+            }
+        }
     }
 
     @Test

diff --git a/...re/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustReservedWords.kt b/...re/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustReservedWords.kt
@@ -8,28 +8,29 @@ package software.amazon.smithy.rust.codegen.core.rustlang
 import software.amazon.smithy.codegen.core.ReservedWordSymbolProvider
 import software.amazon.smithy.codegen.core.ReservedWords
 import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.model.shapes.EnumShape
 import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.shapes.UnionShape
-import software.amazon.smithy.model.traits.EnumDefinition
-import software.amazon.smithy.rust.codegen.core.smithy.MaybeRenamed
+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
-import software.amazon.smithy.rust.codegen.core.util.orNull
-import software.amazon.smithy.rust.codegen.core.util.toPascalCase
 
 class RustReservedWordSymbolProvider(private val base: RustSymbolProvider) : WrappingSymbolProvider(base) {
     private val internal =
         ReservedWordSymbolProvider.builder().symbolProvider(base).memberReservedWords(RustReservedWords).build()
 
     override fun toMemberName(shape: MemberShape): String {
-        val baseName = internal.toMemberName(shape)
-        return when (val container = model.expectShape(shape.container)) {
-            is StructureShape -> when (baseName) {
+        val baseName = super.toMemberName(shape)
+        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"
@@ -40,10 +41,10 @@ class RustReservedWordSymbolProvider(private val base: RustSymbolProvider) : Wra
                 "customize" -> "customize_value"
                 // To avoid conflicts with the error metadata `meta` field
                 "meta" -> "meta_value"
-                else -> baseName
+                else -> reservedWordReplacedName
             }
 
-            is UnionShape -> when (baseName) {
+            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"
@@ -53,7 +54,20 @@ class RustReservedWordSymbolProvider(private val base: RustSymbolProvider) : Wra
                 "Self" -> "SelfValue"
                 // Real models won't end in `_` so it's safe to stop here
                 "SelfValue" -> "SelfValue_"
-                else -> baseName
+                else -> reservedWordReplacedName
+            }
+
+            container is EnumShape || container.hasTrait<EnumTrait>() -> when (baseName) {
+                // Self cannot be used as a raw identifier, so we can't use the normal escaping strategy
+                // https://internals.rust-lang.org/t/raw-identifiers-dont-work-for-all-identifiers/9094/4
+                "Self" -> "SelfValue"
+                // Real models won't end in `_` so it's safe to stop here
+                "SelfValue" -> "SelfValue_"
+                // 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
             }
 
             else -> error("unexpected container: $container")
@@ -67,46 +81,31 @@ class RustReservedWordSymbolProvider(private val base: RustSymbolProvider) : Wra
      * code generators to generate special docs.
      */
     override fun toSymbol(shape: Shape): Symbol {
+        // Sanity check that the symbol provider stack is set up correctly
+        check(super.toSymbol(shape).renamedFrom() == null) {
+            "RustReservedWordSymbolProvider should only run once"
+        }
+
+        var renamedSymbol = internal.toSymbol(shape)
         return when (shape) {
             is MemberShape -> {
                 val container = model.expectShape(shape.container)
-                if (!(container is StructureShape || container is UnionShape)) {
+                val containerIsEnum = container is EnumShape || container.hasTrait<EnumTrait>()
+                if (container !is StructureShape && container !is UnionShape && !containerIsEnum) {
                     return base.toSymbol(shape)
                 }
                 val previousName = base.toMemberName(shape)
                 val escapedName = this.toMemberName(shape)
-                val baseSymbol = base.toSymbol(shape)
                 // if the names don't match and it isn't a simple escaping with `r#`, record a rename
-                baseSymbol.letIf(escapedName != previousName && !escapedName.contains("r#")) {
-                    it.toBuilder().renamedFrom(previousName).build()
-                }
+                renamedSymbol.toBuilder().name(escapedName)
+                    .letIf(escapedName != previousName && !escapedName.contains("r#")) {
+                        it.renamedFrom(previousName)
+                    }.build()
             }
 
             else -> base.toSymbol(shape)
         }
     }
-
-    override fun toEnumVariantName(definition: EnumDefinition): MaybeRenamed? {
-        val baseName = base.toEnumVariantName(definition) ?: return null
-        check(definition.name.orNull()?.toPascalCase() == baseName.name) {
-            "Enum variants must already be in pascal case ${baseName.name} differed from ${baseName.name.toPascalCase()}. Definition: ${definition.name}"
-        }
-        check(baseName.renamedFrom == null) {
-            "definitions should only pass through the renamer once"
-        }
-        return when (baseName.name) {
-            // Self cannot be used as a raw identifier, so we can't use the normal escaping strategy
-            // https://internals.rust-lang.org/t/raw-identifiers-dont-work-for-all-identifiers/9094/4
-            "Self" -> MaybeRenamed("SelfValue", "Self")
-            // Real models won't end in `_` so it's safe to stop here
-            "SelfValue" -> MaybeRenamed("SelfValue_", "SelfValue")
-            // Unknown is used as the name of the variant containing unexpected values
-            "Unknown" -> MaybeRenamed("UnknownValue", "Unknown")
-            // Real models won't end in `_` so it's safe to stop here
-            "UnknownValue" -> MaybeRenamed("UnknownValue_", "UnknownValue")
-            else -> baseName
-        }
-    }
 }
 
 object RustReservedWords : ReservedWords {

diff --git a/...ore/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RustSymbolProvider.kt b/...ore/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RustSymbolProvider.kt
@@ -14,19 +14,16 @@ import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.UnionShape
-import software.amazon.smithy.model.traits.EnumDefinition
 import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 
 /**
- * SymbolProvider interface that carries both the inner configuration and a function to produce an enum variant name.
+ * SymbolProvider interface that carries additional configuration and module/symbol resolution.
  */
 interface RustSymbolProvider : SymbolProvider {
     val model: Model
     val moduleProviderContext: ModuleProviderContext
     val config: RustSymbolProviderConfig
 
-    fun toEnumVariantName(definition: EnumDefinition): MaybeRenamed?
-
     fun moduleForShape(shape: Shape): RustModule.LeafModule =
         config.moduleProvider.moduleForShape(moduleProviderContext, shape)
     fun moduleForOperationError(operation: OperationShape): RustModule.LeafModule =
@@ -84,7 +81,6 @@ open class WrappingSymbolProvider(private val base: RustSymbolProvider) : RustSy
     override val moduleProviderContext: ModuleProviderContext get() = base.moduleProviderContext
     override val config: RustSymbolProviderConfig get() = base.config
 
-    override fun toEnumVariantName(definition: EnumDefinition): MaybeRenamed? = base.toEnumVariantName(definition)
     override fun toSymbol(shape: Shape): Symbol = base.toSymbol(shape)
     override fun toMemberName(shape: MemberShape): String = base.toMemberName(shape)
     override fun symbolForOperationError(operation: OperationShape): Symbol = base.symbolForOperationError(operation)

diff --git a/...src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/SymbolMetadataProvider.kt b/...src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/SymbolMetadataProvider.kt
@@ -114,6 +114,11 @@ class BaseSymbolMetadataProvider(
             }
 
             is UnionShape, is CollectionShape, is MapShape -> RustMetadata(visibility = Visibility.PUBLIC)
+
+            // This covers strings with the enum trait for now, and can be removed once we're fully on EnumShape
+            // TODO(https://github.com/awslabs/smithy-rs/issues/1700): Remove this `is StringShape` match arm
+            is StringShape -> RustMetadata(visibility = Visibility.PUBLIC)
+
             else -> TODO("Unrecognized container type: $container")
         }
 

diff --git a/...gen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/SymbolVisitor.kt b/...gen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/SymbolVisitor.kt
@@ -17,6 +17,7 @@ import software.amazon.smithy.model.shapes.BooleanShape
 import software.amazon.smithy.model.shapes.ByteShape
 import software.amazon.smithy.model.shapes.DocumentShape
 import software.amazon.smithy.model.shapes.DoubleShape
+import software.amazon.smithy.model.shapes.EnumShape
 import software.amazon.smithy.model.shapes.FloatShape
 import software.amazon.smithy.model.shapes.IntegerShape
 import software.amazon.smithy.model.shapes.ListShape
@@ -35,7 +36,6 @@ import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.shapes.TimestampShape
 import software.amazon.smithy.model.shapes.UnionShape
-import software.amazon.smithy.model.traits.EnumDefinition
 import software.amazon.smithy.model.traits.EnumTrait
 import software.amazon.smithy.model.traits.ErrorTrait
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
@@ -47,7 +47,6 @@ import software.amazon.smithy.rust.codegen.core.smithy.traits.RustBoxTrait
 import software.amazon.smithy.rust.codegen.core.util.PANIC
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
 import software.amazon.smithy.rust.codegen.core.util.letIf
-import software.amazon.smithy.rust.codegen.core.util.orNull
 import software.amazon.smithy.rust.codegen.core.util.toPascalCase
 import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 import kotlin.reflect.KClass
@@ -132,21 +131,13 @@ open class SymbolVisitor(
             module.toType().resolve("${symbol.name}Error").toSymbol().toBuilder().locatedIn(module).build()
         }
 
-    /**
-     * Return the name of a given `enum` variant. Note that this refers to `enum` in the Smithy context
-     * where enum is a trait that can be applied to [StringShape] and not in the Rust context of an algebraic data type.
-     *
-     * Because enum variants are not member shape, a separate handler is required.
-     */
-    override fun toEnumVariantName(definition: EnumDefinition): MaybeRenamed? {
-        val baseName = definition.name.orNull()?.toPascalCase() ?: return null
-        return MaybeRenamed(baseName, null)
-    }
-
-    override fun toMemberName(shape: MemberShape): String = when (val container = model.expectShape(shape.container)) {
-        is StructureShape -> shape.memberName.toSnakeCase()
-        is UnionShape -> shape.memberName.toPascalCase()
-        else -> error("unexpected container shape: $container")
+    override fun toMemberName(shape: MemberShape): String {
+        val container = model.expectShape(shape.container)
+        return when {
+            container is StructureShape -> shape.memberName.toSnakeCase()
+            container is UnionShape || container is EnumShape || container.hasTrait<EnumTrait>() -> shape.memberName.toPascalCase()
+            else -> error("unexpected container shape: $container")
+        }
     }
 
     override fun blobShape(shape: BlobShape?): Symbol {

diff --git a/...c/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt b/...c/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt
@@ -7,6 +7,8 @@ package software.amazon.smithy.rust.codegen.core.smithy.generators
 
 import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.MemberShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.traits.DocumentationTrait
 import software.amazon.smithy.model.traits.EnumDefinition
@@ -27,12 +29,14 @@ import software.amazon.smithy.rust.codegen.core.smithy.MaybeRenamed
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.expectRustMetadata
+import software.amazon.smithy.rust.codegen.core.smithy.renamedFrom
 import software.amazon.smithy.rust.codegen.core.util.REDACTION
 import software.amazon.smithy.rust.codegen.core.util.dq
 import software.amazon.smithy.rust.codegen.core.util.expectTrait
 import software.amazon.smithy.rust.codegen.core.util.getTrait
 import software.amazon.smithy.rust.codegen.core.util.orNull
 import software.amazon.smithy.rust.codegen.core.util.shouldRedact
+import software.amazon.smithy.rust.codegen.core.util.toPascalCase
 
 data class EnumGeneratorContext(
     val enumName: String,
@@ -71,14 +75,41 @@ abstract class EnumType {
 }
 
 /** Model that wraps [EnumDefinition] to calculate and cache values required to generate the Rust enum source. */
-class EnumMemberModel(private val definition: EnumDefinition, private val symbolProvider: RustSymbolProvider) {
+class EnumMemberModel(
+    private val parentShape: Shape,
+    private val definition: EnumDefinition,
+    private val symbolProvider: RustSymbolProvider,
+) {
+    companion object {
+        /**
+         * Return the name of a given `enum` variant. Note that this refers to `enum` in the Smithy context
+         * where enum is a trait that can be applied to [StringShape] and not in the Rust context of an algebraic data type.
+         *
+         * Ordinarily, the symbol provider would determine this name, but the enum trait doesn't allow for this.
+         *
+         * TODO(https://github.com/awslabs/smithy-rs/issues/1700): Remove this function when refactoring to EnumShape.
+         */
+        @Deprecated("This function will go away when we handle EnumShape instead of EnumTrait")
+        fun toEnumVariantName(
+            symbolProvider: RustSymbolProvider,
+            parentShape: Shape,
+            definition: EnumDefinition,
+        ): MaybeRenamed? {
+            val name = definition.name.orNull()?.toPascalCase() ?: return null
+            // Create a fake member shape for symbol look up until we refactor to use EnumShape
+            val fakeMemberShape =
+                MemberShape.builder().id(parentShape.id.withMember(name)).target("smithy.api#String").build()
+            val symbol = symbolProvider.toSymbol(fakeMemberShape)
+            return MaybeRenamed(symbol.name, symbol.renamedFrom())
+        }
+    }
     // Because enum variants always start with an upper case letter, they will never
     // conflict with reserved words (which are always lower case), therefore, we never need
     // to fall back to raw identifiers
 
     val value: String get() = definition.value
 
-    fun name(): MaybeRenamed? = symbolProvider.toEnumVariantName(definition)
+    fun name(): MaybeRenamed? = toEnumVariantName(symbolProvider, parentShape, definition)
 
     private fun renderDocumentation(writer: RustWriter) {
         val name =
@@ -97,7 +128,7 @@ class EnumMemberModel(private val definition: EnumDefinition, private val symbol
         }
     }
 
-    fun derivedName() = checkNotNull(symbolProvider.toEnumVariantName(definition)).name
+    fun derivedName() = checkNotNull(toEnumVariantName(symbolProvider, parentShape, definition)).name
 
     fun render(writer: RustWriter) {
         renderDocumentation(writer)
@@ -138,7 +169,7 @@ open class EnumGenerator(
         enumName = symbol.name,
         enumMeta = symbol.expectRustMetadata(),
         enumTrait = enumTrait,
-        sortedMembers = enumTrait.values.sortedBy { it.value }.map { EnumMemberModel(it, symbolProvider) },
+        sortedMembers = enumTrait.values.sortedBy { it.value }.map { EnumMemberModel(shape, it, symbolProvider) },
     )
 
     fun render(writer: RustWriter) {