Refactor UnconstrainedUnionGenerator to use `ValidationExceptionCon…

…versionGenerator` (#3733) `UnconstrainedUnionGenerator` should use `ValidationExceptionConversionGenerator` to generate the `as_validation_exception` method. --------- Co-authored-by: Fahad Zubair <[email protected]>
smithy-lang · Jul 2, 2024 · dc1ffb8 · dc1ffb8
1 parent 9af72f5
commit dc1ffb8
Show file tree

Hide file tree

Showing 7 changed files with 62 additions and 24 deletions.
diff --git a/...in/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCodegenVisitor.kt b/...in/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCodegenVisitor.kt
@@ -212,6 +212,7 @@ class PythonServerCodegenVisitor(
                     rustCrate.createInlineModuleCreator(),
                     this@modelsModuleWriter,
                     shape,
+                    validationExceptionConversionGenerator,
                 ).render()
             }
         }

diff --git a/...src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt b/...src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt
@@ -596,6 +596,7 @@ open class ServerCodegenVisitor(
                     rustCrate.createInlineModuleCreator(),
                     this@modelsModuleWriter,
                     shape,
+                    validationExceptionConversionGenerator,
                 ).render()
             }
         }

diff --git a/...rust/codegen/server/smithy/customizations/CustomValidationExceptionWithReasonDecorator.kt b/...rust/codegen/server/smithy/customizations/CustomValidationExceptionWithReasonDecorator.kt
@@ -18,6 +18,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.rustlang.withBlock
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
@@ -31,6 +32,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.generators.Length
 import software.amazon.smithy.rust.codegen.server.smithy.generators.Pattern
 import software.amazon.smithy.rust.codegen.server.smithy.generators.Range
 import software.amazon.smithy.rust.codegen.server.smithy.generators.StringTraitInfo
+import software.amazon.smithy.rust.codegen.server.smithy.generators.UnionConstraintTraitInfo
 import software.amazon.smithy.rust.codegen.server.smithy.generators.ValidationExceptionConversionGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.generators.isKeyConstrained
 import software.amazon.smithy.rust.codegen.server.smithy.generators.isValueConstrained
@@ -320,4 +322,19 @@ class ValidationExceptionWithReasonConversionGenerator(private val codegenContex
             "AsValidationExceptionFields" to validationExceptionFields.join("\n"),
         )
     }
+
+    override fun unionShapeConstraintViolationImplBlock(
+        unionConstraintTraitInfo: Collection<UnionConstraintTraitInfo>,
+    ) = writable {
+        rustBlockTemplate(
+            "pub(crate) fn as_validation_exception_field(self, path: #{String}) -> crate::model::ValidationExceptionField",
+            "String" to RuntimeType.String,
+        ) {
+            withBlock("match self {", "}") {
+                for (constraintViolation in unionConstraintTraitInfo) {
+                    rust("""Self::${constraintViolation.name()}(inner) => inner.as_validation_exception_field(path + "/${constraintViolation.forMember.memberName}"),""")
+                }
+            }
+        }
+    }
 }
diff --git a/...on/smithy/rust/codegen/server/smithy/customizations/SmithyValidationExceptionDecorator.kt b/...on/smithy/rust/codegen/server/smithy/customizations/SmithyValidationExceptionDecorator.kt
@@ -18,6 +18,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.rustlang.withBlock
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
@@ -30,6 +31,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.generators.ConstraintVi
 import software.amazon.smithy.rust.codegen.server.smithy.generators.Range
 import software.amazon.smithy.rust.codegen.server.smithy.generators.StringTraitInfo
 import software.amazon.smithy.rust.codegen.server.smithy.generators.TraitInfo
+import software.amazon.smithy.rust.codegen.server.smithy.generators.UnionConstraintTraitInfo
 import software.amazon.smithy.rust.codegen.server.smithy.generators.ValidationExceptionConversionGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.generators.isKeyConstrained
 import software.amazon.smithy.rust.codegen.server.smithy.generators.isValueConstrained
@@ -244,4 +246,19 @@ class SmithyValidationExceptionConversionGenerator(private val codegenContext: S
             "AsValidationExceptionFields" to validationExceptionFields.join(""),
         )
     }
+
+    override fun unionShapeConstraintViolationImplBlock(
+        unionConstraintTraitInfo: Collection<UnionConstraintTraitInfo>,
+    ) = writable {
+        rustBlockTemplate(
+            "pub(crate) fn as_validation_exception_field(self, path: #{String}) -> crate::model::ValidationExceptionField",
+            "String" to RuntimeType.String,
+        ) {
+            withBlock("match self {", "}") {
+                for (constraintViolation in unionConstraintTraitInfo) {
+                    rust("""Self::${constraintViolation.name()}(inner) => inner.as_validation_exception_field(path + "/${constraintViolation.forMember.memberName}"),""")
+                }
+            }
+        }
+    }
 }
diff --git a/...ftware/amazon/smithy/rust/codegen/server/smithy/generators/UnconstrainedUnionGenerator.kt b/...ftware/amazon/smithy/rust/codegen/server/smithy/generators/UnconstrainedUnionGenerator.kt
@@ -15,7 +15,6 @@ import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.Visibility
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
-import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.withBlock
 import software.amazon.smithy.rust.codegen.core.rustlang.withBlockTemplate
@@ -53,6 +52,7 @@ class UnconstrainedUnionGenerator(
     private val inlineModuleCreator: InlineModuleCreator,
     private val modelsModuleWriter: RustWriter,
     val shape: UnionShape,
+    private val validationExceptionConversionGenerator: ValidationExceptionConversionGenerator,
 ) {
     private val model = codegenContext.model
     private val symbolProvider = codegenContext.symbolProvider
@@ -172,18 +172,15 @@ class UnconstrainedUnionGenerator(
             )
 
             if (shape.isReachableFromOperationInput()) {
-                rustBlock("impl $constraintViolationName") {
-                    rustBlockTemplate(
-                        "pub(crate) fn as_validation_exception_field(self, path: #{String}) -> crate::model::ValidationExceptionField",
-                        "String" to RuntimeType.String,
-                    ) {
-                        withBlock("match self {", "}") {
-                            for (constraintViolation in constraintViolations()) {
-                                rust("""Self::${constraintViolation.name()}(inner) => inner.as_validation_exception_field(path + "/${constraintViolation.forMember.memberName}"),""")
-                            }
-                        }
+                rustTemplate(
+                    """
+                    impl $constraintViolationName {
+                        #{UnionShapeConstraintViolationImplBlock:W}
                     }
-                }
+                    """,
+                    "UnionShapeConstraintViolationImplBlock" to
+                        validationExceptionConversionGenerator.unionShapeConstraintViolationImplBlock(constraintViolations()),
+                )
             }
         }
     }
@@ -199,30 +196,26 @@ class UnconstrainedUnionGenerator(
             }
         }
 
-    data class ConstraintViolation(val forMember: MemberShape) {
-        fun name() = forMember.memberName.toPascalCase()
-    }
-
     private fun constraintViolations() =
         sortedMembers
             .filter { it.targetCanReachConstrainedShape(model, symbolProvider) }
-            .map { ConstraintViolation(it) }
+            .map { UnionConstraintTraitInfo(it) }
 
     private fun renderConstraintViolation(
         writer: RustWriter,
-        constraintViolation: ConstraintViolation,
+        unionConstraintTraitInfo: UnionConstraintTraitInfo,
     ) {
-        val targetShape = model.expectShape(constraintViolation.forMember.target)
+        val targetShape = model.expectShape(unionConstraintTraitInfo.forMember.target)
 
         val constraintViolationSymbol =
             constraintViolationSymbolProvider.toSymbol(targetShape)
                 // Box this constraint violation symbol if necessary.
-                .letIf(constraintViolation.forMember.hasTrait<ConstraintViolationRustBoxTrait>()) {
+                .letIf(unionConstraintTraitInfo.forMember.hasTrait<ConstraintViolationRustBoxTrait>()) {
                     it.makeRustBoxed()
                 }
 
         writer.rust(
-            "${constraintViolation.name()}(#T),",
+            "${unionConstraintTraitInfo.name()}(#T),",
             constraintViolationSymbol,
         )
     }
@@ -291,7 +284,7 @@ class UnconstrainedUnionGenerator(
                     {
                         let constrained: #{ConstrainedSymbol} = $unconstrainedVar
                             .try_into()$boxIt$boxErr
-                            .map_err(Self::Error::${ConstraintViolation(member).name()})?;
+                            .map_err(Self::Error::${UnionConstraintTraitInfo(member).name()})?;
                         constrained.into()
                     }
                     """,
@@ -304,9 +297,13 @@ class UnconstrainedUnionGenerator(
                         .try_into()
                         $boxIt
                         $boxErr
-                        .map_err(Self::Error::${ConstraintViolation(member).name()})?
+                        .map_err(Self::Error::${UnionConstraintTraitInfo(member).name()})?
                     """,
                 )
             }
         }
 }
+
+data class UnionConstraintTraitInfo(val forMember: MemberShape) {
+    fun name() = forMember.memberName.toPascalCase()
+}
diff --git a/...on/smithy/rust/codegen/server/smithy/generators/ValidationExceptionConversionGenerator.kt b/...on/smithy/rust/codegen/server/smithy/generators/ValidationExceptionConversionGenerator.kt
@@ -53,4 +53,8 @@ interface ValidationExceptionConversionGenerator {
         collectionConstraintsInfo: Collection<CollectionTraitInfo>,
         isMemberConstrained: Boolean,
     ): Writable
+
+    fun unionShapeConstraintViolationImplBlock(
+        unionConstraintTraitInfo: Collection<UnionConstraintTraitInfo>,
+    ): Writable
 }
diff --git a/...re/amazon/smithy/rust/codegen/server/smithy/generators/UnconstrainedUnionGeneratorTest.kt b/...re/amazon/smithy/rust/codegen/server/smithy/generators/UnconstrainedUnionGeneratorTest.kt
@@ -16,6 +16,7 @@ import software.amazon.smithy.rust.codegen.core.testutil.unitTest
 import software.amazon.smithy.rust.codegen.core.util.lookup
 import software.amazon.smithy.rust.codegen.server.smithy.ServerRustModule
 import software.amazon.smithy.rust.codegen.server.smithy.createInlineModuleCreator
+import software.amazon.smithy.rust.codegen.server.smithy.customizations.SmithyValidationExceptionConversionGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.generators.protocol.ServerRestJsonProtocol
 import software.amazon.smithy.rust.codegen.server.smithy.renderInlineMemoryModules
 import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverRenderWithModelBuilder
@@ -63,7 +64,7 @@ class UnconstrainedUnionGeneratorTest {
             TestUtility.generateIsError().invoke(this)
 
             project.withModule(ServerRustModule.Model) modelsModuleWriter@{
-                UnconstrainedUnionGenerator(codegenContext, project.createInlineModuleCreator(), this@modelsModuleWriter, unionShape).render()
+                UnconstrainedUnionGenerator(codegenContext, project.createInlineModuleCreator(), this@modelsModuleWriter, unionShape, SmithyValidationExceptionConversionGenerator(codegenContext)).render()
 
                 this@unconstrainedModuleWriter.unitTest(
                     name = "unconstrained_union_fail_to_constrain",