Use sealed interfaces and records for Unions

Author: adwsingh

codegen/codegen-core/src/it/java/software/amazon/smithy/java/codegen/test/RecursionTests.java

@@ -106,7 +106,7 @@ void multiplyRecursiveUnionWorks() {
         var output = builder.build();
         assertEquals(recursive.hashCode(), output.hashCode());
         assertEquals(recursive, output);
-        assertNotEquals(AttributeValue.Type.$UNKNOWN, output.type());
+        assertNotEquals(AttributeValue.$Unknown.class, output.getClass());
     }
 
     @Test

codegen/codegen-core/src/it/java/software/amazon/smithy/java/codegen/test/UnionTest.java

@@ -6,6 +6,8 @@
 package software.amazon.smithy.java.codegen.test;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.math.BigDecimal;
@@ -25,14 +27,13 @@
 import software.amazon.smithy.java.codegen.test.model.NestedStruct;
 import software.amazon.smithy.java.codegen.test.model.NestedUnion;
 import software.amazon.smithy.java.codegen.test.model.UnionType;
-import software.amazon.smithy.java.core.schema.SerializableShape;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.java.core.serde.document.Document;
 import software.amazon.smithy.model.shapes.ShapeType;
 
 public class UnionTest {
 
-    static Stream<SerializableShape> unionTypes() {
+    static Stream<UnionType> unionTypes() {
         return Stream.of(
                 new UnionType.BooleanValueMember(true),
                 new UnionType.ListValueMember(List.of("a", "b")),
@@ -103,7 +104,11 @@ void unknownUnionDeser() {
         document.deserializeInto(builder);
         var output = builder.build();
 
-        assertEquals(UnionType.Type.$UNKNOWN, output.type());
+        if (output instanceof UnionType.$Unknown(String memberName)) {
+            assertEquals("UNKNOWN!!!", memberName);
+        } else {
+            throw new AssertionError("Expected $Unknown variant");
+        }
         assertEquals("UNKNOWN!!!", output.getValue());
     }
 
@@ -112,4 +117,19 @@ void unknownUnionSerFails() {
         var union = UnionType.builder().$unknownMember("foo").build();
         assertThrows(UnsupportedOperationException.class, () -> Document.of(union));
     }
+
+    @Test
+    void showExhaustiveMatchRequired() {
+        var union = getNestedUnion();
+        switch (union) {
+            case NestedUnion.AMember(var a) -> assertNotNull(a);
+            case NestedUnion.BMember(var b) -> assertNotEquals(0, b);
+            case NestedUnion.$Unknown(var memberName) -> assertNotNull(memberName);
+            default -> throw new IllegalArgumentException("Unexpected member: " + union);
+        }
+    }
+
+    private static NestedUnion getNestedUnion() {
+        return NestedUnion.builder().a("a").build();
+    }
 }

codegen/codegen-core/src/main/java/software/amazon/smithy/java/codegen/generators/BuilderGenerator.java

@@ -13,6 +13,7 @@
 import software.amazon.smithy.java.core.schema.SchemaUtils;
 import software.amazon.smithy.java.core.schema.ShapeBuilder;
 import software.amazon.smithy.model.Model;
+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.ShapeType;
@@ -21,7 +22,7 @@
  * Generates a static nested {@code Builder} class for a Java class.
  */
 abstract class BuilderGenerator implements Runnable {
-    private final JavaWriter writer;
+    protected final JavaWriter writer;
     protected final Shape shape;
     protected final SymbolProvider symbolProvider;
     protected final Model model;
@@ -53,14 +54,14 @@ public void run() {
                         /**
                          * @return returns a new Builder.
                          */
-                        public static Builder builder() {
+                        ${^inInterface}public ${/inInterface}static Builder builder() {
                             return new Builder();
                         }
 
                         /**
                          * Builder for {@link ${shape:T}}.
                          */
-                        public static final class Builder implements ${sdkShapeBuilder:T}<${shape:T}>${?isStaged}, ${#stages}${value:L}${^key.last}, ${/key.last}${/stages}${/isStaged} {
+                        ${^inInterface}public static ${/inInterface}final class Builder implements ${sdkShapeBuilder:T}<${shape:T}>${?isStaged}, ${#stages}${value:L}${^key.last}, ${/key.last}${/stages}${/isStaged} {
                             ${builderProperties:C|}
 
                             ${builderConstructor:C|}
@@ -91,6 +92,7 @@ public static final class Builder implements ${sdkShapeBuilder:T}<${shape:T}>${?
         writer.putContext("setMemberValue", writer.consumer(this::generateSetMemberValue));
         boolean isStaged = !this.stageInterfaces().isEmpty();
         writer.putContext("isStaged", isStaged);
+        writer.putContext("inInterface", inInterface());
         if (isStaged) {
             writer.putContext("stages", this.stageInterfaces());
             writer.putContext("stageGen", writer.consumer(this::generateStages));
@@ -130,6 +132,14 @@ protected List<String> stageInterfaces() {
         return List.of();
     }
 
+    protected boolean inInterface() {
+        return false;
+    }
+
+    protected String getMemberSchemaName(MemberShape member) {
+        return CodegenUtils.toMemberSchemaName(symbolProvider.toMemberName(member));
+    }
+
     protected void generateSetMemberValue(JavaWriter writer) {
         // Don't override the default implementation that throws if there are no members.
         if (shape.members().isEmpty() || (shape.getType() == ShapeType.ENUM || shape.getType() == ShapeType.INT_ENUM)) {
@@ -157,7 +167,7 @@ protected void generateMemberValueSetters(JavaWriter writer) {
             writer.pushState();
             writer.putContext("memberName", symbolProvider.toMemberName(member));
             writer.putContext("type", symbolProvider.toSymbol(member));
-            writer.putContext("memberSchema", CodegenUtils.toMemberSchemaName(symbolProvider.toMemberName(member)));
+            writer.putContext("memberSchema", getMemberSchemaName(member));
             writer.putContext("schemaUtilsClass", SchemaUtils.class);
             writer.putContext("isNullable", CodegenUtils.isNullableMember(model, member));
             writer.write(

codegen/codegen-core/src/main/java/software/amazon/smithy/java/codegen/generators/IdStringGenerator.java

@@ -16,13 +16,18 @@
  * @param shape Shape to write ID for
  */
 @SmithyInternalApi
-public record IdStringGenerator(JavaWriter writer, Shape shape) implements Runnable {
+public record IdStringGenerator(JavaWriter writer, Shape shape, boolean isInterface) implements Runnable {
+
+    public IdStringGenerator(JavaWriter writer, Shape shape) {
+        this(writer, shape, false);
+    }
 
     @Override
     public void run() {
         writer.pushState();
         writer.putContext("shapeId", ShapeId.class);
-        writer.write("public static final ${shapeId:T} $$ID = $$SCHEMA.id();");
+        var qualifiers = isInterface ? "" : "public static final ";
+        writer.write(qualifiers + "${shapeId:T} $$ID = $$SCHEMA.id();");
         writer.popState();
     }
 }

codegen/codegen-core/src/main/java/software/amazon/smithy/java/codegen/generators/SchemaFieldGenerator.java

@@ -90,11 +90,9 @@ public Void structureShape(StructureShape shape) {
     public Void unionShape(UnionShape shape) {
         writer.pushState();
         writer.putContext("name", schemaFieldOrder.getSchemaFieldName(shape, writer));
-        writer.write("public static final ${schemaClass:T} $$SCHEMA = ${name:L};");
-
-        for (var member : shape.members()) {
-            writeMemberProperty(member);
-        }
+        // Unions are sealed interfaces, so $SCHEMA is implicitly public static final
+        // Member schemas are generated inside each record variant, not at the interface level
+        writer.write("${schemaClass:T} $$SCHEMA = ${name:L};");
 
         writer.popState();
         return null;