smithy-lang · david-perez · Feb 15, 2023 · Feb 14, 2023 · Feb 15, 2023 · Velfi
@@ -105,14 +105,14 @@ class ClientCodegenVisitor(
             // Add errors attached at the service level to the models
             .let { ModelTransformer.create().copyServiceErrorsToOperations(it, settings.getService(it)) }
             // Add `Box<T>` to recursive shapes as necessary
-            .let(RecursiveShapeBoxer::transform)
+            .let(RecursiveShapeBoxer()::transform)
             // Normalize the `message` field on errors when enabled in settings (default: true)
             .letIf(settings.codegenConfig.addMessageToErrors, AddErrorMessage::transform)
             // NormalizeOperations by ensuring every operation has an input & output shape
             .let(OperationNormalizer::transform)
             // Drop unsupported event stream operations from the model
             .let { RemoveEventStreamOperations.transform(it, settings) }
-            // - Normalize event stream operations
+            // Normalize event stream operations
             .let(EventStreamNormalizer::transform)
 
     /**

@@ -36,7 +36,7 @@ internal class ResiliencyConfigCustomizationTest {
 
     @Test
     fun `generates a valid config`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val project = TestWorkspace.testProject()
         val codegenContext = testCodegenContext(model, settings = project.rustSettings())
 

@@ -12,7 +12,8 @@ import software.amazon.smithy.model.traits.Trait
 /**
  * Trait indicating that this shape should be represented with `Box<T>` when converted into Rust
  *
- * This is used to handle recursive shapes. See RecursiveShapeBoxer.
+ * This is used to handle recursive shapes.
+ * See [software.amazon.smithy.rust.codegen.core.smithy.transformers.RecursiveShapeBoxer].
  *
  * This trait is synthetic, applied during code generation, and never used in actual models.
  */

@@ -7,21 +7,29 @@ package software.amazon.smithy.rust.codegen.core.smithy.transformers
 
 import software.amazon.smithy.codegen.core.TopologicalIndex
 import software.amazon.smithy.model.Model
-import software.amazon.smithy.model.shapes.ListShape
+import software.amazon.smithy.model.shapes.CollectionShape
 import software.amazon.smithy.model.shapes.MapShape
 import software.amazon.smithy.model.shapes.MemberShape
-import software.amazon.smithy.model.shapes.SetShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.transform.ModelTransformer
 import software.amazon.smithy.rust.codegen.core.smithy.traits.RustBoxTrait
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
 
-object RecursiveShapeBoxer {
+class RecursiveShapeBoxer(
+    private val containsIndirectionPredicate: (Collection<Shape>) -> Boolean = ::containsIndirection,
+    private val boxShapeFn: (MemberShape) -> MemberShape = ::addRustBoxTrait,
+) {
     /**
-     * Transform a model which may contain recursive shapes into a model annotated with [RustBoxTrait]
+     * Transform a model which may contain recursive shapes into a model annotated with [RustBoxTrait].
      *
-     * When recursive shapes do NOT go through a List, Map, or Set, they must be boxed in Rust. This function will
-     * iteratively find loops & add the `RustBox` trait in a deterministic way until it reaches a fixed point.
+     * When recursive shapes do NOT go through a `CollectionShape` or a `MapShape` shape, they must be boxed in Rust.
+     * This function will iteratively find loops and add the [RustBoxTrait] trait in a deterministic way until it
+     * reaches a fixed point.
+     *
+     * Why `CollectionShape`s and `MapShape`s? Note that `CollectionShape`s get rendered in Rust as `Vec<T>`, and
+     * `MapShape`s as `HashMap<String, T>`; they're the only Smithy shapes that "organically" introduce indirection
+     * (via a pointer to the heap) in the recursive path. For other recursive paths, we thus have to introduce the
+     * indirection artificially ourselves using `Box`.
      *
      * This function MUST be deterministic (always choose the same shapes to `Box`). If it is not, that is a bug. Even so
      * this function may cause backward compatibility issues in certain pathological cases where a changes to recursive
@@ -41,12 +49,12 @@ object RecursiveShapeBoxer {
      * If [model] contains no loops, return null.
      */
     private fun transformInner(model: Model): Model? {
-        // Execute 1-step of the boxing algorithm in the path to reaching a fixed point
-        // 1. Find all the shapes that are part of a cycle
-        // 2. Find all the loops that those shapes are part of
-        // 3. Filter out the loops that go through a layer of indirection
-        // 3. Pick _just one_ of the remaining loops to fix
-        // 4. Select the member shape in that loop with the earliest shape id
+        // Execute 1 step of the boxing algorithm in the path to reaching a fixed point:
+        // 1. Find all the shapes that are part of a cycle.
+        // 2. Find all the loops that those shapes are part of.
+        // 3. Filter out the loops that go through a layer of indirection.
+        // 3. Pick _just one_ of the remaining loops to fix.
+        // 4. Select the member shape in that loop with the earliest shape id.
         // 5. Box it.
         // (External to this function) Go back to 1.
         val index = TopologicalIndex.of(model)
@@ -58,34 +66,32 @@ object RecursiveShapeBoxer {
             // Flatten the connections into shapes.
             loops.map { it.shapes }
         }
-        val loopToFix = loops.firstOrNull { !containsIndirection(it) }
+        val loopToFix = loops.firstOrNull { !containsIndirectionPredicate(it) }
 
         return loopToFix?.let { loop: List<Shape> ->
             check(loop.isNotEmpty())
-            // pick the shape to box in a deterministic way
+            // Pick the shape to box in a deterministic way.
             val shapeToBox = loop.filterIsInstance<MemberShape>().minByOrNull { it.id }!!
             ModelTransformer.create().mapShapes(model) { shape ->
                 if (shape == shapeToBox) {
-                    shape.asMemberShape().get().toBuilder().addTrait(RustBoxTrait()).build()
+                    boxShapeFn(shape.asMemberShape().get())
                 } else {
                     shape
                 }
             }
         }
     }
+}
 
-    /**
-     * Check if a List<Shape> contains a shape which will use a pointer when represented in Rust, avoiding the
-     * need to add more Boxes
-     */
-    private fun containsIndirection(loop: List<Shape>): Boolean {
-        return loop.find {
-            when (it) {
-                is ListShape,
-                is MapShape,
-                is SetShape, -> true
-                else -> it.hasTrait<RustBoxTrait>()
-            }
-        } != null
+/**
+ * Check if a `List<Shape>` contains a shape which will use a pointer when represented in Rust, avoiding the
+ * need to add more `Box`es.
+ */
+private fun containsIndirection(loop: Collection<Shape>): Boolean = loop.find {
+    when (it) {
+        is CollectionShape, is MapShape -> true
+        else -> it.hasTrait<RustBoxTrait>()
     }
-}
+} != null
+
+private fun addRustBoxTrait(shape: MemberShape): MemberShape = shape.toBuilder().addTrait(RustBoxTrait()).build()
@@ -82,7 +82,7 @@ class InstantiatorTest {
             @required
             num: Integer
         }
-    """.asSmithyModel().let { RecursiveShapeBoxer.transform(it) }
+    """.asSmithyModel().let { RecursiveShapeBoxer().transform(it) }
 
     private val codegenContext = testCodegenContext(model)
     private val symbolProvider = codegenContext.symbolProvider

@@ -313,7 +313,7 @@ class StructureGeneratorTest {
     @Test
     fun `it generates accessor methods`() {
         val testModel =
-            RecursiveShapeBoxer.transform(
+            RecursiveShapeBoxer().transform(
                 """
                 namespace test
 

@@ -42,7 +42,7 @@ class AwsQueryParserGeneratorTest {
 
     @Test
     fun `it modifies operation parsing to include Response and Result tags`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model)
         val symbolProvider = codegenContext.symbolProvider
         val parserGenerator = AwsQueryParserGenerator(

@@ -42,7 +42,7 @@ class Ec2QueryParserGeneratorTest {
 
     @Test
     fun `it modifies operation parsing to include Response and Result tags`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model)
         val symbolProvider = codegenContext.symbolProvider
         val parserGenerator = Ec2QueryParserGenerator(

@@ -114,7 +114,7 @@ class JsonParserGeneratorTest {
 
     @Test
     fun `generates valid deserializers`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model)
         val symbolProvider = codegenContext.symbolProvider
         fun builderSymbol(shape: StructureShape): Symbol =

@@ -94,7 +94,7 @@ internal class XmlBindingTraitParserGeneratorTest {
 
     @Test
     fun `generates valid parsers`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model)
         val symbolProvider = codegenContext.symbolProvider
         val parserGenerator = XmlBindingTraitParserGenerator(

@@ -92,7 +92,7 @@ class AwsQuerySerializerGeneratorTest {
             true -> CodegenTarget.CLIENT
             false -> CodegenTarget.SERVER
         }
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model, codegenTarget = codegenTarget)
         val symbolProvider = codegenContext.symbolProvider
         val parserGenerator = AwsQuerySerializerGenerator(testCodegenContext(model, codegenTarget = codegenTarget))

@@ -85,7 +85,7 @@ class Ec2QuerySerializerGeneratorTest {
 
     @Test
     fun `generates valid serializers`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model)
         val symbolProvider = codegenContext.symbolProvider
         val parserGenerator = Ec2QuerySerializerGenerator(codegenContext)

@@ -100,7 +100,7 @@ class JsonSerializerGeneratorTest {
 
     @Test
     fun `generates valid serializers`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model)
         val symbolProvider = codegenContext.symbolProvider
         val parserSerializer = JsonSerializerGenerator(

@@ -105,7 +105,7 @@ internal class XmlBindingTraitSerializerGeneratorTest {
 
     @Test
     fun `generates valid serializers`() {
-        val model = RecursiveShapeBoxer.transform(OperationNormalizer.transform(baseModel))
+        val model = RecursiveShapeBoxer().transform(OperationNormalizer.transform(baseModel))
         val codegenContext = testCodegenContext(model)
         val symbolProvider = codegenContext.symbolProvider
         val parserGenerator = XmlBindingTraitSerializerGenerator(

@@ -31,7 +31,7 @@ internal class RecursiveShapeBoxerTest {
                 hello: Hello
             }
         """.asSmithyModel()
-        RecursiveShapeBoxer.transform(model) shouldBe model
+        RecursiveShapeBoxer().transform(model) shouldBe model
     }
 
     @Test
@@ -43,7 +43,7 @@ internal class RecursiveShapeBoxerTest {
                 anotherField: Boolean
             }
         """.asSmithyModel()
-        val transformed = RecursiveShapeBoxer.transform(model)
+        val transformed = RecursiveShapeBoxer().transform(model)
         val member: MemberShape = transformed.lookup("com.example#Recursive\$RecursiveStruct")
         member.expectTrait<RustBoxTrait>()
     }
@@ -70,7 +70,7 @@ internal class RecursiveShapeBoxerTest {
                  third: SecondTree
             }
         """.asSmithyModel()
-        val transformed = RecursiveShapeBoxer.transform(model)
+        val transformed = RecursiveShapeBoxer().transform(model)
         val boxed = transformed.shapes().filter { it.hasTrait<RustBoxTrait>() }.toList()
         boxed.map { it.id.toString().removePrefix("com.example#") }.toSet() shouldBe setOf(
             "Atom\$add",

@@ -66,7 +66,7 @@ class RecursiveShapesIntegrationTest {
         }
         output.message shouldContain "has infinite size"
 
-        val fixedProject = check(RecursiveShapeBoxer.transform(model))
+        val fixedProject = check(RecursiveShapeBoxer().transform(model))
         fixedProject.compileAndTest()
     }
 }
@@ -78,6 +78,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.generators.protocol.Ser
 import software.amazon.smithy.rust.codegen.server.smithy.protocols.ServerProtocolLoader
 import software.amazon.smithy.rust.codegen.server.smithy.traits.isReachableFromOperationInput
 import software.amazon.smithy.rust.codegen.server.smithy.transformers.AttachValidationExceptionToConstrainedOperationInputsInAllowList
+import software.amazon.smithy.rust.codegen.server.smithy.transformers.RecursiveConstraintViolationBoxer
 import software.amazon.smithy.rust.codegen.server.smithy.transformers.RemoveEbsModelValidationException
 import software.amazon.smithy.rust.codegen.server.smithy.transformers.ShapesReachableFromOperationInputTagger
 import java.util.logging.Logger
@@ -162,7 +163,9 @@ open class ServerCodegenVisitor(
             // Add errors attached at the service level to the models
             .let { ModelTransformer.create().copyServiceErrorsToOperations(it, settings.getService(it)) }
             // Add `Box<T>` to recursive shapes as necessary
-            .let(RecursiveShapeBoxer::transform)
+            .let(RecursiveShapeBoxer()::transform)
+            // Add `Box<T>` to recursive constraint violations as necessary
+            .let(RecursiveConstraintViolationBoxer::transform)
             // Normalize operations by adding synthetic input and output shapes to every operation
             .let(OperationNormalizer::transform)
             // Remove the EBS model's own `ValidationException`, which collides with `smithy.framework#ValidationException`

@@ -9,12 +9,15 @@ import software.amazon.smithy.model.shapes.CollectionShape
 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.join
-import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.smithy.makeRustBoxed
 import software.amazon.smithy.rust.codegen.core.smithy.module
+import software.amazon.smithy.rust.codegen.core.util.hasTrait
+import software.amazon.smithy.rust.codegen.core.util.letIf
 import software.amazon.smithy.rust.codegen.server.smithy.PubCrateConstraintViolationSymbolProvider
 import software.amazon.smithy.rust.codegen.server.smithy.ServerCodegenContext
 import software.amazon.smithy.rust.codegen.server.smithy.canReachConstrainedShape
+import software.amazon.smithy.rust.codegen.server.smithy.traits.ConstraintViolationRustBoxTrait
 import software.amazon.smithy.rust.codegen.server.smithy.traits.isReachableFromOperationInput
 
 class CollectionConstraintViolationGenerator(
@@ -38,16 +41,22 @@ class CollectionConstraintViolationGenerator(
     private val constraintsInfo: List<TraitInfo> = collectionConstraintsInfo.map { it.toTraitInfo() }
 
     fun render() {
-        val memberShape = model.expectShape(shape.member.target)
+        val targetShape = model.expectShape(shape.member.target)
         val constraintViolationSymbol = constraintViolationSymbolProvider.toSymbol(shape)
         val constraintViolationName = constraintViolationSymbol.name
-        val isMemberConstrained = memberShape.canReachConstrainedShape(model, symbolProvider)
+        val isMemberConstrained = targetShape.canReachConstrainedShape(model, symbolProvider)
         val constraintViolationVisibility = Visibility.publicIf(publicConstrainedTypes, Visibility.PUBCRATE)
 
         modelsModuleWriter.withInlineModule(constraintViolationSymbol.module()) {
             val constraintViolationVariants = constraintsInfo.map { it.constraintViolationVariant }.toMutableList()
             if (isMemberConstrained) {
                 constraintViolationVariants += {
+                    val memberConstraintViolationSymbol =
+                        constraintViolationSymbolProvider.toSymbol(targetShape).letIf(
+                            shape.member.hasTrait<ConstraintViolationRustBoxTrait>(),
+                        ) {
+                            it.makeRustBoxed()
+                        }
                     rustTemplate(
                         """
                         /// Constraint violation error when an element doesn't satisfy its own constraints.
@@ -56,7 +65,7 @@ class CollectionConstraintViolationGenerator(
                         ##[doc(hidden)]
                         Member(usize, #{MemberConstraintViolationSymbol})
                         """,
-                        "MemberConstraintViolationSymbol" to constraintViolationSymbolProvider.toSymbol(memberShape),
+                        "MemberConstraintViolationSymbol" to memberConstraintViolationSymbol,
                     )
                 }
             }

@@ -11,10 +11,13 @@ import software.amazon.smithy.model.traits.LengthTrait
 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.rustTemplate
+import software.amazon.smithy.rust.codegen.core.smithy.makeRustBoxed
 import software.amazon.smithy.rust.codegen.core.smithy.module
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
+import software.amazon.smithy.rust.codegen.core.util.letIf
 import software.amazon.smithy.rust.codegen.server.smithy.PubCrateConstraintViolationSymbolProvider
 import software.amazon.smithy.rust.codegen.server.smithy.ServerCodegenContext
+import software.amazon.smithy.rust.codegen.server.smithy.traits.ConstraintViolationRustBoxTrait
 import software.amazon.smithy.rust.codegen.server.smithy.traits.isReachableFromOperationInput
 
 class MapConstraintViolationGenerator(
@@ -47,7 +50,14 @@ class MapConstraintViolationGenerator(
             constraintViolationCodegenScopeMutableList.add("KeyConstraintViolationSymbol" to constraintViolationSymbolProvider.toSymbol(keyShape))
         }
         if (isValueConstrained(valueShape, model, symbolProvider)) {
-            constraintViolationCodegenScopeMutableList.add("ValueConstraintViolationSymbol" to constraintViolationSymbolProvider.toSymbol(valueShape))
+            constraintViolationCodegenScopeMutableList.add(
+                "ValueConstraintViolationSymbol" to
+                    constraintViolationSymbolProvider.toSymbol(valueShape).letIf(
+                        shape.value.hasTrait<ConstraintViolationRustBoxTrait>(),
+                    ) {
+                        it.makeRustBoxed()
+                    },
+            )
             constraintViolationCodegenScopeMutableList.add("KeySymbol" to constrainedShapeSymbolProvider.toSymbol(keyShape))
         }
         val constraintViolationCodegenScope = constraintViolationCodegenScopeMutableList.toTypedArray()