Constraint member types are refactored as standalone shapes.

Author: Fahad Zubair

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/protocols/eventstream/ClientEventStreamBaseRequirements.kt

@@ -23,6 +23,7 @@ import software.amazon.smithy.rust.codegen.client.testutil.testSymbolProvider
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.implBlock
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenTarget
+import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.generators.BuilderGenerator
 import software.amazon.smithy.rust.codegen.core.testutil.EventStreamTestModels
@@ -51,6 +52,7 @@ abstract class ClientEventStreamBaseRequirements : EventStreamTestRequirements<C
     )
 
     override fun renderBuilderForShape(
+        rustCrate: RustCrate,
         writer: RustWriter,
         codegenContext: ClientCodegenContext,
         shape: StructureShape,

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

@@ -18,6 +18,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenTarget
 import software.amazon.smithy.rust.codegen.core.smithy.DirectedWalker
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.generators.UnionGenerator
 import software.amazon.smithy.rust.codegen.core.smithy.generators.renderUnknownVariant
@@ -62,6 +63,7 @@ interface EventStreamTestRequirements<C : CodegenContext> {
 
     /** Render a builder for the given shape */
     fun renderBuilderForShape(
+        rustCrate: RustCrate,
         writer: RustWriter,
         codegenContext: C,
         shape: StructureShape,

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ConstrainedShapeSymbolProvider.kt

@@ -19,7 +19,10 @@ import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShortShape
 import software.amazon.smithy.model.shapes.StringShape
+import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.traits.LengthTrait
+import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
+import software.amazon.smithy.rust.codegen.core.rustlang.RustReservedWords
 import software.amazon.smithy.rust.codegen.core.rustlang.RustType
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.WrappingSymbolProvider
@@ -32,6 +35,8 @@ import software.amazon.smithy.rust.codegen.core.smithy.symbolBuilder
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
 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 software.amazon.smithy.rust.codegen.server.smithy.generators.serverBuilderModule
 
 /**
  * The [ConstrainedShapeSymbolProvider] returns, for a given _directly_
@@ -56,14 +61,19 @@ class ConstrainedShapeSymbolProvider(
     private val base: RustSymbolProvider,
     private val model: Model,
     private val serviceShape: ServiceShape,
+    private val publicConstrainedTypes : Boolean = true
 ) : WrappingSymbolProvider(base) {
     private val nullableIndex = NullableIndex.of(model)
 
     private fun publicConstrainedSymbolForMapOrCollectionShape(shape: Shape): Symbol {
         check(shape is MapShape || shape is CollectionShape)
+        // FZ rebase
+        // val rustType = RustType.Opaque(shape.contextName(serviceShape).toPascalCase())
+        // return symbolBuilder(shape, rustType).locatedIn(ServerRustModule.Model).build()
 
-        val rustType = RustType.Opaque(shape.contextName(serviceShape).toPascalCase())
-        return symbolBuilder(shape, rustType).locatedIn(ServerRustModule.Model).build()
+        val (name, module) = getMemberNameAndModule(shape, serviceShape, ServerRustModule.Model, !publicConstrainedTypes)
+        val rustType = RustType.Opaque(name)
+        return symbolBuilder(shape, rustType).locatedIn(module).build()
     }
 
     override fun toSymbol(shape: Shape): Symbol {
@@ -74,8 +84,14 @@ class ConstrainedShapeSymbolProvider(
                 val target = model.expectShape(shape.target)
                 val targetSymbol = this.toSymbol(target)
                 // Handle boxing first, so we end up with `Option<Box<_>>`, not `Box<Option<_>>`.
-                handleOptionality(handleRustBoxing(targetSymbol, shape), shape, nullableIndex, base.config().nullabilityCheckMode)
+                handleOptionality(
+                    handleRustBoxing(targetSymbol, shape),
+                    shape,
+                    nullableIndex,
+                    base.config().nullabilityCheckMode,
+                )
             }
+
             is MapShape -> {
                 if (shape.isDirectlyConstrained(base)) {
                     check(shape.hasTrait<LengthTrait>()) {
@@ -91,6 +107,7 @@ class ConstrainedShapeSymbolProvider(
                         .build()
                 }
             }
+
             is CollectionShape -> {
                 if (shape.isDirectlyConstrained(base)) {
                     check(constrainedCollectionCheck(shape)) {
@@ -105,8 +122,15 @@ class ConstrainedShapeSymbolProvider(
 
             is StringShape, is IntegerShape, is ShortShape, is LongShape, is ByteShape, is BlobShape -> {
                 if (shape.isDirectlyConstrained(base)) {
-                    val rustType = RustType.Opaque(shape.contextName(serviceShape).toPascalCase())
-                    symbolBuilder(shape, rustType).locatedIn(ServerRustModule.Model).build()
+                    // FZ rebase
+                    //val rustType = RustType.Opaque(shape.contextName(serviceShape).toPascalCase())
+                    //symbolBuilder(shape, rustType).locatedIn(ServerRustModule.Model).build()
+
+                    // A standalone constrained shape goes into `ModelsModule`, but one
+                    // arising from a constrained member shape goes into a module for the container.
+                    val (name, module) = getMemberNameAndModule(shape, serviceShape, ServerRustModule.Model, !publicConstrainedTypes)
+                    val rustType = RustType.Opaque(name)
+                    symbolBuilder(shape, rustType).locatedIn(module).build()
                 } else {
                     base.toSymbol(shape)
                 }
@@ -122,9 +146,51 @@ class ConstrainedShapeSymbolProvider(
      *  - That it has no unsupported constraints applied.
      */
     private fun constrainedCollectionCheck(shape: CollectionShape): Boolean {
-        val supportedConstraintTraits = supportedCollectionConstraintTraits.mapNotNull { shape.getTrait(it).orNull() }.toSet()
+        val supportedConstraintTraits =
+            supportedCollectionConstraintTraits.mapNotNull { shape.getTrait(it).orNull() }.toSet()
         val allConstraintTraits = allConstraintTraits.mapNotNull { shape.getTrait(it).orNull() }.toSet()
 
-        return supportedConstraintTraits.isNotEmpty() && allConstraintTraits.subtract(supportedConstraintTraits).isEmpty()
+        return supportedConstraintTraits.isNotEmpty() && allConstraintTraits.subtract(supportedConstraintTraits)
+            .isEmpty()
+    }
+
+
+    /**
+     * Returns the pair (Rust Symbol Name, Inline Module) for the shape. At the time of model transformation all
+     * constrained member shapes are extracted and are given a model-wide unique name. However, the generated code
+     * for the new shapes is in a module that is named after the containing shape (structure, list, map or union).
+     * The new shape's Rust Symbol is renamed from `{structureName}{memberName}` to  `{structure_name}::{member_name}`
+     */
+    private fun getMemberNameAndModule(
+        shape: Shape,
+        serviceShape: ServiceShape,
+        defaultModule: RustModule.LeafModule,
+        pubCrateServerBuilder: Boolean,
+    ): Pair<String, RustModule.LeafModule> {
+        val (container, member) =
+            shape.overriddenConstrainedMemberInfo() ?: return Pair(shape.contextName(serviceShape), defaultModule)
+
+        return if (container is StructureShape) {
+            val builderModule = container.serverBuilderModule(base, pubCrateServerBuilder)
+            val renameTo = member.memberName ?: member.id.name
+            Pair(renameTo.toPascalCase(), builderModule)
+        } else {
+            // For List, Union and Map, the new shape defined for a constrained member shape
+            // need to be placed into an inline module named `pub {container_name_in_snake_case}`
+            val innerModuleName = RustReservedWords.escapeIfNeeded(container.id.name.toSnakeCase()) + if (pubCrateServerBuilder) {
+                    "_internal"
+            } else {
+                    ""
+            }
+
+            val innerModule = RustModule.new(
+                innerModuleName,
+                visibility = Visibility.publicIf(!pubCrateServerBuilder, Visibility.PUBCRATE),
+                parent = defaultModule,
+                inline = true,
+            )
+            val renameTo = member.memberName ?: member.id.name
+            Pair(renameTo.toPascalCase(), innerModule)
+        }
     }
 }

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ConstraintViolationSymbolProvider.kt

@@ -29,8 +29,10 @@ import software.amazon.smithy.rust.codegen.core.smithy.contextName
 import software.amazon.smithy.rust.codegen.core.smithy.locatedIn
 import software.amazon.smithy.rust.codegen.core.smithy.module
 import software.amazon.smithy.rust.codegen.core.smithy.rustType
+import software.amazon.smithy.rust.codegen.core.util.getTrait
 import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 import software.amazon.smithy.rust.codegen.server.smithy.generators.serverBuilderSymbol
+import software.amazon.smithy.rust.codegen.server.smithy.traits.SyntheticStructureFromConstrainedMemberTrait
 
 /**
  * The [ConstraintViolationSymbolProvider] returns, for a given constrained
@@ -79,14 +81,34 @@ class ConstraintViolationSymbolProvider(
 
     private fun Shape.shapeModule(): RustModule.LeafModule {
         val documentation = if (publicConstrainedTypes && this.isDirectlyConstrained(base)) {
-            "See [`${this.contextName(serviceShape)}`]."
+            if (this.hasTrait(SyntheticStructureFromConstrainedMemberTrait.ID)) {
+                val symbol = base.toSymbol(this)
+                "See [`${this.contextName(serviceShape)}`]($symbol)."
+            } else {
+                "See [`${this.contextName(serviceShape)}`]."
+            }
         } else {
             null
         }
-        return RustModule.new(
+
+        val syntheticTrait = getTrait<SyntheticStructureFromConstrainedMemberTrait>()
+
+        val (module, name) = if (syntheticTrait != null) {
+            // For constrained member shapes, the ConstraintViolation code needs to go in an inline rust module
+            // that is a descendant of the module that contains the extracted shape itself.
+            val overriddenMemberModule = this.getParentAndInlineModuleForConstrainedMember(base, !publicConstrainedTypes)!!
+            val name = syntheticTrait.member.memberName
+            Pair(overriddenMemberModule.second, RustReservedWords.escapeIfNeeded(name).toSnakeCase())
+        } else {
             // Need to use the context name so we get the correct name for maps.
-            name = RustReservedWords.escapeIfNeeded(this.contextName(serviceShape)).toSnakeCase(),
+            Pair(ModelsModule, RustReservedWords.escapeIfNeeded(this.contextName(serviceShape)).toSnakeCase())
+        }
+
+        return RustModule.new(
+            name = name,
             visibility = visibility,
+            // FZ rebase
+            //parent = ServerRustModule.Model,
             parent = ServerRustModule.Model,
             inline = true,
             documentation = documentation,

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/Constraints.kt

@@ -14,6 +14,7 @@ import software.amazon.smithy.model.shapes.IntegerShape
 import software.amazon.smithy.model.shapes.LongShape
 import software.amazon.smithy.model.shapes.MapShape
 import software.amazon.smithy.model.shapes.MemberShape
+import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShortShape
 import software.amazon.smithy.model.shapes.SimpleShape
@@ -26,11 +27,21 @@ import software.amazon.smithy.model.traits.PatternTrait
 import software.amazon.smithy.model.traits.RangeTrait
 import software.amazon.smithy.model.traits.RequiredTrait
 import software.amazon.smithy.model.traits.UniqueItemsTrait
+import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
+import software.amazon.smithy.rust.codegen.core.rustlang.RustReservedWords
+import software.amazon.smithy.rust.codegen.core.rustlang.Visibility
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.DirectedWalker
+import software.amazon.smithy.rust.codegen.core.smithy.ModelsModule
+import software.amazon.smithy.rust.codegen.core.smithy.contextName
 import software.amazon.smithy.rust.codegen.core.smithy.isOptional
+import software.amazon.smithy.rust.codegen.core.smithy.module
 import software.amazon.smithy.rust.codegen.core.util.UNREACHABLE
+import software.amazon.smithy.rust.codegen.core.util.getTrait
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
+import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
+import software.amazon.smithy.rust.codegen.server.smithy.generators.serverBuilderModule
+import software.amazon.smithy.rust.codegen.server.smithy.traits.SyntheticStructureFromConstrainedMemberTrait
 
 /**
  * This file contains utilities to work with constrained shapes.
@@ -160,3 +171,52 @@ fun Shape.typeNameContainsNonPublicType(
     is StructureShape, is UnionShape -> false
     else -> UNREACHABLE("the above arms should be exhaustive, but we received shape: $this")
 }
+
+
+/**
+ * For synthetic shapes that are added to the model because of member constrained shapes, it returns
+ * the "container" and "the member shape" that originally had the constraint trait. For all other
+ * shapes, it returns null.
+ */
+fun Shape.overriddenConstrainedMemberInfo(): Pair<Shape, MemberShape>? {
+    val trait = getTrait<SyntheticStructureFromConstrainedMemberTrait>() ?: return null
+    return Pair(trait.container, trait.member)
+}
+
+
+/**
+ * Returns the parent and the inline module that this particular shape should go in.
+ */
+fun Shape.getParentAndInlineModuleForConstrainedMember(symbolProvider: SymbolProvider, pubCrateServerBuilder: Boolean): Pair<RustModule.LeafModule, RustModule.LeafModule>? {
+    //val container = overriddenConstrainedMemberInfo()?.first ?: return null
+    val overriddenTrait = getTrait<SyntheticStructureFromConstrainedMemberTrait>() ?: return null
+    return if (overriddenTrait.container is StructureShape) {
+        val structureModule = symbolProvider.toSymbol(overriddenTrait.container).module()
+        val builderModule = overriddenTrait.container.serverBuilderModule(symbolProvider, pubCrateServerBuilder)
+        Pair(structureModule, builderModule)
+    }
+    else {
+        // For constrained member shapes, the ConstraintViolation code needs to go in an inline rust module
+        // that is a descendant of the module that contains the extracted shape itself.
+        return if (!pubCrateServerBuilder) {
+            // List, union and map types need to go into their own module
+            val shapeSymbol = symbolProvider.toSymbol(this)
+            val shapeModule = shapeSymbol.module()
+            check(!shapeModule.parent.isInline()) {
+                "parent module of $id should not be an inline module"
+            }
+            Pair(shapeModule.parent as RustModule.LeafModule, shapeModule)
+        }
+        else {
+            val name = RustReservedWords.escapeIfNeeded(overriddenTrait.container.id.name).toSnakeCase() + "_internal"
+            val innerModule = RustModule.new(
+                    name = name,
+                    visibility = Visibility.PUBCRATE,
+                    parent = ModelsModule,
+                    inline = true,
+                )
+
+            Pair(ModelsModule, innerModule)
+        }
+    }
+}