Remove optionality from documents that are in unions

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/SymbolVisitor.kt

@@ -168,7 +168,7 @@ class SymbolVisitor(
         // an Input shape, then the field is _not optional_.
         val httpLabeledInput =
             container.hasTrait<SyntheticInputTrait>() && member.hasTrait<HttpLabelTrait>()
-        return if (nullableIndex.isNullable(member) && !httpLabeledInput || model.expectShape(member.target).isDocumentShape) {
+        return if (nullableIndex.isNullable(member) && !httpLabeledInput) {
             symbol.makeOptional()
         } else symbol
     }

codegen/src/test/kotlin/software/amazon/smithy/rust/codegen/generators/StructureGeneratorTest.kt

@@ -58,8 +58,13 @@ class StructureGeneratorTest {
             // test that sensitive can be applied directly to a member or to the shape
             secretKey: SecretKey
         }
+
+        structure StructWithDoc {
+            doc: Document
+        }
         """.asSmithyModel()
         val struct = model.lookup<StructureShape>("com.test#MyStruct")
+        val structWithDoc = model.lookup<StructureShape>("com.test#StructWithDoc")
         val inner = model.lookup<StructureShape>("com.test#Inner")
         val credentials = model.lookup<StructureShape>("com.test#Credentials")
         val error = model.lookup<StructureShape>("com.test#MyError")
@@ -178,4 +183,20 @@ class StructureGeneratorTest {
 
         writer.compileAndTest()
     }
+
+    @Test
+    fun `documents are optional in structs`() {
+        val provider = testSymbolProvider(model)
+        val writer = RustWriter.forModule("lib")
+        StructureGenerator(model, provider, writer, structWithDoc).render()
+
+        writer.compileAndTest(
+            """
+            let _struct = StructWithDoc {
+                // This will only compile if the document is optional
+                doc: None
+            };
+            """
+        )
+    }
 }

codegen/src/test/kotlin/software/amazon/smithy/rust/codegen/generators/UnionGeneratorTest.kt

@@ -18,42 +18,37 @@ import software.amazon.smithy.rust.codegen.util.lookup
 class UnionGeneratorTest {
     @Test
     fun `generate basic unions`() {
-        val model = """
-        namespace test
-        union MyUnion {
-            stringConfig: String,
-            @documentation("This *is* documentation about the member")
-            intConfig: PrimitiveInteger
-        }
-        """.asSmithyModel()
-        val provider: SymbolProvider = testSymbolProvider(model)
-        val writer = RustWriter.forModule("model")
-        val generator = UnionGenerator(model, provider, writer, model.lookup("test#MyUnion"))
-        generator.render()
+        val writer = generateUnion(
+            """
+            union MyUnion {
+                stringConfig: String,
+                @documentation("This *is* documentation about the member")
+                intConfig: PrimitiveInteger
+            }
+            """
+        )
+
         writer.compileAndTest(
             """
-        let var_a = MyUnion::StringConfig("abc".to_string());
-        let var_b = MyUnion::IntConfig(10);
-        assert_ne!(var_a, var_b);
-        assert_eq!(var_a, var_a);
-        """
+            let var_a = MyUnion::StringConfig("abc".to_string());
+            let var_b = MyUnion::IntConfig(10);
+            assert_ne!(var_a, var_b);
+            assert_eq!(var_a, var_a);
+            """
         )
         writer.toString() shouldContain "#[non_exhaustive]"
     }
 
     @Test
     fun `generate conversion helper methods`() {
-        val model = """
-        namespace test
-        union MyUnion {
-            stringValue: String,
-            intValue: PrimitiveInteger
-        }
-        """.asSmithyModel()
-        val provider: SymbolProvider = testSymbolProvider(model)
-        val writer = RustWriter.forModule("model")
-        val generator = UnionGenerator(model, provider, writer, model.lookup("test#MyUnion"))
-        generator.render()
+        val writer = generateUnion(
+            """
+            union MyUnion {
+                stringValue: String,
+                intValue: PrimitiveInteger
+            }
+            """
+        )
 
         writer.compileAndTest(
             """
@@ -67,7 +62,26 @@ class UnionGeneratorTest {
             assert_eq!(bar.is_int_value(), true);
             assert_eq!(bar.as_string_value(), None);
             assert_eq!(bar.as_int_value(), Some(&10));
-        """
+            """
+        )
+    }
+
+    @Test
+    fun `documents are not optional in unions`() {
+        val writer = generateUnion("union MyUnion { doc: Document, other: String }")
+        writer.compileAndTest(
+            """
+            // If the document isn't optional, this will compile
+            MyUnion::Doc(smithy_types::Document::Null);
+            """
         )
     }
+
+    private fun generateUnion(modelSmithy: String, unionName: String = "MyUnion"): RustWriter {
+        val model = "namespace test\n$modelSmithy".asSmithyModel()
+        val provider: SymbolProvider = testSymbolProvider(model)
+        val writer = RustWriter.forModule("model")
+        UnionGenerator(model, provider, writer, model.lookup("test#$unionName")).render()
+        return writer
+    }
 }