Fix codegen for unions with the @httpPayload trait

CHANGELOG.next.toml

@@ -84,3 +84,9 @@ message = "Generate a region setter when a model uses SigV4."
 references = ["smithy-rs#2960"]
 meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "client" }
 author = "jdisanti"
+
+[[smithy-rs]]
+message = "Fix code generation for union members with the `@httpPayload` trait."
+references = ["smithy-rs#2969", "smithy-rs#1896"]
+meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "all" }
+author = "jdisanti"

codegen-client-test/model/rest-xml-extras.smithy

@@ -21,6 +21,9 @@ service RestXmlExtras {
         StringHeader,
         CreateFoo,
         RequiredMember,
+        // TODO(https://github.com/awslabs/smithy-rs/issues/2968): Remove the following once these tests are included in Smithy
+        // They're being added in https://github.com/smithy-lang/smithy/pull/1908
+        HttpPayloadWithUnion,
     ]
 }
 
@@ -257,3 +260,96 @@ structure RequiredMemberInputOutput {
     @required
     requiredString: String
 }
+
+// TODO(https://github.com/awslabs/smithy-rs/issues/2968): Delete the HttpPayloadWithUnion tests below once Smithy vends them
+// They're being added in https://github.com/smithy-lang/smithy/pull/1908
+
+/// This example serializes a union in the payload.
+@idempotent
+@http(uri: "/HttpPayloadWithUnion", method: "PUT")
+operation HttpPayloadWithUnion {
+    input: HttpPayloadWithUnionInputOutput,
+    output: HttpPayloadWithUnionInputOutput
+}
+
+apply HttpPayloadWithUnion @httpRequestTests([
+    {
+        id: "RestXmlHttpPayloadWithUnion",
+        documentation: "Serializes a union in the payload.",
+        protocol: restXml,
+        method: "PUT",
+        uri: "/HttpPayloadWithUnion",
+        body: """
+              <UnionPayload>
+                  <greeting>hello</greeting>
+              </UnionPayload>""",
+        bodyMediaType: "application/xml",
+        headers: {
+            "Content-Type": "application/xml",
+        },
+        requireHeaders: [
+            "Content-Length"
+        ],
+        params: {
+            nested: {
+                greeting: "hello"
+            }
+        }
+    },
+    {
+        id: "RestXmlHttpPayloadWithUnsetUnion",
+        documentation: "No payload is sent if the union has no value.",
+        protocol: restXml,
+        method: "PUT",
+        uri: "/HttpPayloadWithUnion",
+        body: "",
+        headers: {
+            "Content-Type": "application/xml",
+            "Content-Length": "0"
+        },
+        params: {}
+    }
+])
+
+apply HttpPayloadWithUnion @httpResponseTests([
+    {
+        id: "RestXmlHttpPayloadWithUnion",
+        documentation: "Serializes a union in the payload.",
+        protocol: restXml,
+        code: 200,
+        body: """
+              <UnionPayload>
+                  <greeting>hello</greeting>
+              </UnionPayload>""",
+        bodyMediaType: "application/xml",
+        headers: {
+            "Content-Type": "application/xml",
+        },
+        params: {
+            nested: {
+                greeting: "hello"
+            }
+        }
+    },
+    {
+        id: "RestXmlHttpPayloadWithUnsetUnion",
+        documentation: "No payload is sent if the union has no value.",
+        protocol: restXml,
+        code: 200,
+        body: "",
+        headers: {
+            "Content-Type": "application/xml",
+            "Content-Length": "0"
+        },
+        params: {}
+    }
+])
+
+structure HttpPayloadWithUnionInputOutput {
+    @httpPayload
+    nested: UnionPayload,
+}
+
+union UnionPayload {
+    greeting: String
+}

codegen-core/common-test-models/rest-json-extras.smithy

@@ -65,6 +65,9 @@ service RestJsonExtras {
         NullInNonSparse,
         CaseInsensitiveErrorOperation,
         EmptyStructWithContentOnWireOp,
+        // TODO(https://github.com/awslabs/smithy-rs/issues/2968): Remove the following once these tests are included in Smithy
+        // They're being added in https://github.com/smithy-lang/smithy/pull/1908
+        HttpPayloadWithUnion,
     ],
     errors: [ExtraError]
 }
@@ -348,3 +351,96 @@ structure EmptyStructWithContentOnWireOpOutput {
 operation EmptyStructWithContentOnWireOp {
     output: EmptyStructWithContentOnWireOpOutput,
 }
+
+// TODO(https://github.com/awslabs/smithy-rs/issues/2968): Delete the HttpPayloadWithUnion tests below once Smithy vends them
+// They're being added in https://github.com/smithy-lang/smithy/pull/1908
+
+/// This examples serializes a union in the payload.
+@idempotent
+@http(uri: "/HttpPayloadWithUnion", method: "PUT")
+operation HttpPayloadWithUnion {
+    input: HttpPayloadWithUnionInputOutput,
+    output: HttpPayloadWithUnionInputOutput
+}
+
+structure HttpPayloadWithUnionInputOutput {
+    @httpPayload
+    nested: UnionPayload,
+}
+
+union UnionPayload {
+    greeting: String
+}
+
+apply HttpPayloadWithUnion @httpRequestTests([
+    {
+        id: "RestJsonHttpPayloadWithUnion",
+        documentation: "Serializes a union in the payload.",
+        protocol: restJson1,
+        method: "PUT",
+        uri: "/HttpPayloadWithUnion",
+        body: """
+              {
+                  "greeting": "hello"
+              }""",
+        bodyMediaType: "application/json",
+        headers: {
+            "Content-Type": "application/json"
+        },
+        requireHeaders: [
+            "Content-Length"
+        ],
+        params: {
+            nested: {
+                greeting: "hello"
+            }
+        }
+    },
+    {
+        id: "RestJsonHttpPayloadWithUnsetUnion",
+        documentation: "No payload is sent if the union has no value.",
+        protocol: restJson1,
+        method: "PUT",
+        uri: "/HttpPayloadWithUnion",
+        body: "",
+        headers: {
+            "Content-Type": "application/json",
+            "Content-Length": "0"
+        },
+        params: {}
+    }
+])
+
+apply HttpPayloadWithUnion @httpResponseTests([
+    {
+        id: "RestJsonHttpPayloadWithUnion",
+        documentation: "Serializes a union in the payload.",
+        protocol: restJson1,
+        code: 200,
+        body: """
+              {
+                  "greeting": "hello"
+              }""",
+        bodyMediaType: "application/json",
+        headers: {
+            "Content-Type": "application/json"
+        },
+        params: {
+            nested: {
+                greeting: "hello"
+            }
+        }
+    },
+    {
+        id: "RestJsonHttpPayloadWithUnsetUnion",
+        documentation: "No payload is sent if the union has no value.",
+        protocol: restJson1,
+        code: 200,
+        body: "",
+        headers: {
+            "Content-Type": "application/json",
+            "Content-Length": "0"
+        },
+        params: {}
+    }
+])

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/HttpBoundProtocolPayloadGenerator.kt

@@ -269,7 +269,7 @@ class HttpBoundProtocolPayloadGenerator(
                                 """,
                             )
                             is StructureShape -> rust("#T()", serializerGenerator.unsetStructure(targetShape))
-                            is UnionShape -> throw CodegenException("Currently unsupported. Tracking issue: https://github.com/awslabs/smithy-rs/issues/1896")
+                            is UnionShape -> rust("#T()", serializerGenerator.unsetUnion(targetShape))
                             else -> throw CodegenException("`httpPayload` on member shapes targeting shapes of type ${targetShape.type} is unsupported")
                         }
                     }

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/JsonSerializerGenerator.kt

@@ -36,6 +36,7 @@ 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.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.customize.NamedCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.customize.Section
@@ -170,7 +171,7 @@ class JsonSerializerGenerator(
     private val runtimeConfig = codegenContext.runtimeConfig
     private val protocolFunctions = ProtocolFunctions(codegenContext)
     private val codegenScope = arrayOf(
-        "String" to RuntimeType.String,
+        *preludeScope,
         "Error" to runtimeConfig.serializationError(),
         "SdkBody" to RuntimeType.sdkBody(runtimeConfig),
         "JsonObjectWriter" to RuntimeType.smithyJson(runtimeConfig).resolve("serialize::JsonObjectWriter"),
@@ -232,7 +233,7 @@ class JsonSerializerGenerator(
     }
 
     override fun unsetStructure(structure: StructureShape): RuntimeType =
-        ProtocolFunctions.crossOperationFn("rest_json_unsetpayload") { fnName ->
+        ProtocolFunctions.crossOperationFn("rest_json_unset_struct_payload") { fnName ->
             rustTemplate(
                 """
                 pub fn $fnName() -> #{ByteSlab} {
@@ -243,6 +244,14 @@ class JsonSerializerGenerator(
             )
         }
 
+    override fun unsetUnion(union: UnionShape): RuntimeType =
+        ProtocolFunctions.crossOperationFn("rest_json_unset_union_payload") { fnName ->
+            rustTemplate(
+                "pub fn $fnName() -> #{ByteSlab} { #{Vec}::new() }",
+                *codegenScope,
+            )
+        }
+
     override fun operationInputSerializer(operationShape: OperationShape): RuntimeType? {
         // Don't generate an operation JSON serializer if there is no JSON body.
         val httpDocumentMembers = httpBindingResolver.requestMembers(operationShape, HttpLocation.DOCUMENT)

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/QuerySerializerGenerator.kt

@@ -121,6 +121,10 @@ abstract class QuerySerializerGenerator(codegenContext: CodegenContext) : Struct
         TODO("AwsQuery doesn't support payload serialization")
     }
 
+    override fun unsetUnion(union: UnionShape): RuntimeType {
+        TODO("AwsQuery doesn't support payload serialization")
+    }
+
     override fun operationInputSerializer(operationShape: OperationShape): RuntimeType? {
         val inputShape = operationShape.inputShape(model)
         return protocolFunctions.serializeFn(inputShape, fnNameSuffix = "input") { fnName ->

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/StructuredDataSerializerGenerator.kt

@@ -9,6 +9,7 @@ import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.shapes.StructureShape
+import software.amazon.smithy.model.shapes.UnionShape
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 
 interface StructuredDataSerializerGenerator {
@@ -27,12 +28,22 @@ interface StructuredDataSerializerGenerator {
      * Generate the correct data when attempting to serialize a structure that is unset
      *
      * ```rust
-     * fn rest_json_unsetpayload() -> Vec<u8> {
+     * fn rest_json_unset_struct_payload() -> Vec<u8> {
      *     ...
      * }
      */
     fun unsetStructure(structure: StructureShape): RuntimeType
 
+    /**
+     * Generate the correct data when attempting to serialize a union that is unset
+     *
+     * ```rust
+     * fn rest_json_unset_union_payload() -> Vec<u8> {
+     *     ...
+     * }
+     */
+    fun unsetUnion(union: UnionShape): RuntimeType
+
     /**
      * Generate a serializer for an operation input structure.
      * This serializer is only used by clients.

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/XmlBindingTraitSerializerGenerator.kt

@@ -35,6 +35,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.stripOuter
 import software.amazon.smithy.rust.codegen.core.rustlang.withBlock
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.smithy.generators.UnionGenerator
 import software.amazon.smithy.rust.codegen.core.smithy.generators.renderUnknownVariant
 import software.amazon.smithy.rust.codegen.core.smithy.generators.serializationError
@@ -176,8 +177,8 @@ class XmlBindingTraitSerializerGenerator(
         }
     }
 
-    override fun unsetStructure(structure: StructureShape): RuntimeType {
-        return ProtocolFunctions.crossOperationFn("rest_xml_unset_payload") { fnName ->
+    override fun unsetStructure(structure: StructureShape): RuntimeType =
+        ProtocolFunctions.crossOperationFn("rest_xml_unset_struct_payload") { fnName ->
             rustTemplate(
                 """
                 pub fn $fnName() -> #{ByteSlab} {
@@ -187,7 +188,15 @@ class XmlBindingTraitSerializerGenerator(
                 "ByteSlab" to RuntimeType.ByteSlab,
             )
         }
-    }
+
+    override fun unsetUnion(union: UnionShape): RuntimeType =
+        ProtocolFunctions.crossOperationFn("rest_xml_unset_union_payload") { fnName ->
+            rustTemplate(
+                "pub fn $fnName() -> #{ByteSlab} { #{Vec}::new() }",
+                *preludeScope,
+                "ByteSlab" to RuntimeType.ByteSlab,
+            )
+        }
 
     override fun operationOutputSerializer(operationShape: OperationShape): RuntimeType? {
         val outputShape = operationShape.outputShape(model)