Update unions

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/CodegenVisitor.java

@@ -26,6 +26,7 @@
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.build.PluginContext;
 import software.amazon.smithy.codegen.core.Symbol;
@@ -48,6 +49,7 @@
 import software.amazon.smithy.model.shapes.UnionShape;
 import software.amazon.smithy.model.traits.EnumTrait;
 import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.utils.OptionalUtils;
 
 /**
  * Orchestrates Go client generation.
@@ -162,6 +164,17 @@ void execute() {
             shape.accept(this);
         }
 
+        // Generate a struct to handle unknown tags in unions
+        List<UnionShape> unions = serviceShapes.stream()
+                .map(Shape::asUnionShape)
+                .flatMap(OptionalUtils::stream)
+                .collect(Collectors.toList());
+        if (!unions.isEmpty()) {
+            writers.useShapeWriter(unions.get(0), writer -> {
+                UnionGenerator.generateUnknownUnion(writer, unions, symbolProvider);
+            });
+        }
+
         for (GoIntegration integration : integrations) {
             integration.writeAdditionalFiles(settings, model, symbolProvider, writers::useFileWriter);
             integration.writeAdditionalFiles(settings, model, symbolProvider, writers);

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/ShapeValueGenerator.java

@@ -17,6 +17,8 @@
 
 package software.amazon.smithy.go.codegen;
 
+import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Logger;
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.Symbol;
@@ -35,7 +37,9 @@
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeType;
+import software.amazon.smithy.model.shapes.SimpleShape;
 import software.amazon.smithy.model.shapes.StructureShape;
+import software.amazon.smithy.model.shapes.UnionShape;
 import software.amazon.smithy.model.traits.EnumTrait;
 import software.amazon.smithy.model.traits.StreamingTrait;
 
@@ -79,66 +83,90 @@ public void writeShapeValueInline(GoWriter writer, Shape shape, Node params) {
 
         switch (shape.getType()) {
             case STRUCTURE:
-                structDeclShapeValue(writer, shape.asStructureShape().get(), () -> {
-                    params.accept(new ShapeValueNodeVisitor(writer, this, shape));
-                });
+                structDeclShapeValue(writer, shape.asStructureShape().get(), params);
                 break;
 
             case SET:
             case LIST:
-                listDeclShapeValue(writer, (CollectionShape) shape, () -> {
-                    params.accept(new ShapeValueNodeVisitor(writer, this, shape));
-                });
+                listDeclShapeValue(writer, (CollectionShape) shape, params);
                 break;
 
             case MAP:
-                mapDeclShapeValue(writer, shape.asMapShape().get(), () -> {
-                    params.accept(new ShapeValueNodeVisitor(writer, this, shape));
-                });
+                mapDeclShapeValue(writer, shape.asMapShape().get(), params);
                 break;
 
             case UNION:
+                unionDeclShapeValue(writer, shape.asUnionShape().get(), params.expectObjectNode());
+                break;
+
             case DOCUMENT:
                 LOGGER.warning("Skipping " + shape.getType() + " shape type not supported, " + shape.getId());
                 writer.writeInline("nil");
                 break;
 
             default:
-
-                scalarWrapShapeValue(writer, shape, () -> {
-                    params.accept(new ShapeValueNodeVisitor(writer, this, shape));
-                });
+                writeScalarPointerInline(writer, shape, params);
         }
     }
 
     /**
      * Writes the declaration for a Go structure. Delegates to the runner for member fields within the structure.
      *
      * @param writer writer to write generated code with.
-     * @param shape  the structure shpae
-     * @param inner  lambda to populate struct member fields.
+     * @param shape  the structure shape
+     * @param params parameters to fill the generated shape declaration.
      */
-    protected void structDeclShapeValue(GoWriter writer, StructureShape shape, Runnable inner) {
+    protected void structDeclShapeValue(GoWriter writer, StructureShape shape, Node params) {
         Symbol symbol = symbolProvider.toSymbol(shape);
 
         writer.write("&$T{", symbol);
-        inner.run();
+        params.accept(new ShapeValueNodeVisitor(writer, this, shape));
         writer.writeInline("}");
     }
 
+    /**
+     * Writes the declaration for a Go union.
+     *
+     * @param writer writer to write generated code with.
+     * @param shape  the union shape.
+     * @param params the params.
+     */
+    protected void unionDeclShapeValue(GoWriter writer, UnionShape shape, ObjectNode params) {
+        Symbol symbol = symbolProvider.toSymbol(shape);
+        for (Map.Entry<StringNode, Node> entry : params.getMembers().entrySet()) {
+            Optional<MemberShape> member = shape.getMember(entry.getKey().toString());
+            if (member.isPresent()) {
+                Shape target = model.expectShape(member.get().getTarget());
+                Symbol memberSymbol = SymbolUtils.createValueSymbolBuilder(
+                        symbolProvider.toMemberName(member.get()),
+                        symbol.getNamespace()
+                ).build();
+
+                writer.writeInline("&$T{Value: ", memberSymbol);
+                if (target instanceof SimpleShape) {
+                    writeScalarValueInline(writer, target, entry.getValue());
+                } else {
+                    writeShapeValueInline(writer, target, entry.getValue());
+                }
+                writer.writeInline("}");
+            }
+            return;
+        }
+    }
+
     /**
      * Writes the declaration for a Go slice. Delegates to the runner for fields within the slice.
      *
      * @param writer writer to write generated code with.
      * @param shape  the collection shape
-     * @param inner  lambda to populate collection member fields.
+     * @param params parameters to fill the generated shape declaration.
      */
-    protected void listDeclShapeValue(GoWriter writer, CollectionShape shape, Runnable inner) {
+    protected void listDeclShapeValue(GoWriter writer, CollectionShape shape, Node params) {
         Shape memberShape = model.expectShape(shape.getMember().getTarget());
         Symbol memberSymbol = symbolProvider.toSymbol(memberShape);
 
         writer.write("[]$P{", memberSymbol);
-        inner.run();
+        params.accept(new ShapeValueNodeVisitor(writer, this, shape));
         writer.writeInline("}");
     }
 
@@ -147,17 +175,17 @@ protected void listDeclShapeValue(GoWriter writer, CollectionShape shape, Runnab
      *
      * @param writer writer to write generated code with.
      * @param shape  the map shape.
-     * @param inner  lambda to populate map key/value fields.
+     * @param params parameters to fill the generated shape declaration.
      */
-    protected void mapDeclShapeValue(GoWriter writer, MapShape shape, Runnable inner) {
+    protected void mapDeclShapeValue(GoWriter writer, MapShape shape, Node params) {
         Shape valueShape = model.expectShape(shape.getValue().getTarget());
         Shape keyShape = model.expectShape(shape.getKey().getTarget());
 
         Symbol valueSymbol = symbolProvider.toSymbol(valueShape);
         Symbol keySymbol = symbolProvider.toSymbol(keyShape);
 
         writer.write("map[$T]$P{", keySymbol, valueSymbol);
-        inner.run();
+        params.accept(new ShapeValueNodeVisitor(writer, this, shape));
         writer.writeInline("}");
     }
 
@@ -166,9 +194,9 @@ protected void mapDeclShapeValue(GoWriter writer, MapShape shape, Runnable inner
      *
      * @param writer writer to write generated code with.
      * @param shape  scalar shape.
-     * @param inner  lambda to populate the actual shape value that will be wrapped.
+     * @param params parameters to fill the generated shape declaration.
      */
-    protected void scalarWrapShapeValue(GoWriter writer, Shape shape, Runnable inner) {
+    protected void writeScalarPointerInline(GoWriter writer, Shape shape, Node params) {
         boolean withPtrImport = true;
         String closing = ")";
 
@@ -178,30 +206,15 @@ protected void scalarWrapShapeValue(GoWriter writer, Shape shape, Runnable inner
                 break;
 
             case BLOB:
-                if (shape.hasTrait(StreamingTrait.class)) {
-                    writer.addUseImports(SmithyGoDependency.SMITHY_IO);
-                    writer.addUseImports(SmithyGoDependency.BYTES);
-                    writer.writeInline("smithyio.ReadSeekNopCloser{ReadSeeker: bytes.NewReader([]byte(");
-                    closing += ")}";
-                } else {
-                    writer.writeInline("[]byte(");
-                }
+                closing = "";
                 withPtrImport = false;
                 break;
 
             case STRING:
                 // Enum are not pointers, but string alias values
-                if (shape.hasTrait(StreamingTrait.class)) {
-                    writer.addUseImports(SmithyGoDependency.SMITHY_IO);
-                    writer.addUseImports(SmithyGoDependency.STRINGS);
-                    writer.writeInline("smithyio.ReadSeekNopCloser{ReadSeeker: strings.NewReader(");
-                    closing += "}";
-
-                } else if (shape.hasTrait(EnumTrait.class)) {
-                    Symbol enumSymbol = symbolProvider.toSymbol(shape);
-                    writer.writeInline("$T(", enumSymbol);
+                if (shape.hasTrait(StreamingTrait.class) || shape.hasTrait(EnumTrait.class)) {
+                    closing = "";
                     withPtrImport = false;
-
                 } else {
                     writer.writeInline("ptr.String(");
                 }
@@ -238,7 +251,7 @@ protected void scalarWrapShapeValue(GoWriter writer, Shape shape, Runnable inner
 
             case BIG_INTEGER:
             case BIG_DECIMAL:
-                inner.run();
+                writeScalarValueInline(writer, shape, params);
                 return;
 
             default:
@@ -249,7 +262,44 @@ protected void scalarWrapShapeValue(GoWriter writer, Shape shape, Runnable inner
             writer.addUseImports(SmithyGoDependency.SMITHY_PTR);
         }
 
-        inner.run();
+        writeScalarValueInline(writer, shape, params);
+        writer.writeInline(closing);
+    }
+
+    protected void writeScalarValueInline(GoWriter writer, Shape shape, Node params) {
+        String closing = "";
+        switch (shape.getType()) {
+            case BLOB:
+                if (shape.hasTrait(StreamingTrait.class)) {
+                    writer.addUseImports(SmithyGoDependency.SMITHY_IO);
+                    writer.addUseImports(SmithyGoDependency.BYTES);
+                    writer.writeInline("smithyio.ReadSeekNopCloser{ReadSeeker: bytes.NewReader([]byte(");
+                    closing = "))}";
+                } else {
+                    writer.writeInline("[]byte(");
+                    closing = ")";
+                }
+                break;
+
+            case STRING:
+                // Enum are not pointers, but string alias values
+                if (shape.hasTrait(StreamingTrait.class)) {
+                    writer.addUseImports(SmithyGoDependency.SMITHY_IO);
+                    writer.addUseImports(SmithyGoDependency.STRINGS);
+                    writer.writeInline("smithyio.ReadSeekNopCloser{ReadSeeker: strings.NewReader(");
+                    closing = ")}";
+
+                } else if (shape.hasTrait(EnumTrait.class)) {
+                    Symbol enumSymbol = symbolProvider.toSymbol(shape);
+                    writer.writeInline("$T(", enumSymbol);
+                    closing = ")";
+                }
+                break;
+
+            default:
+                break;
+        }
+        params.accept(new ShapeValueNodeVisitor(writer, this, shape));
         writer.writeInline(closing);
     }
 

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/SymbolVisitor.java

@@ -84,7 +84,13 @@ final class SymbolVisitor implements SymbolProvider, ShapeVisitor<Symbol> {
         this.rootModuleName = rootModuleName;
         this.typesPackageName = rootModuleName + "/types";
 
-        ReservedWords reservedWords = new ReservedWordsBuilder()
+        // Reserve the generated names for union members, including the unknown case.
+        ReservedWordsBuilder reservedNames = new ReservedWordsBuilder()
+                .put(UnionGenerator.UNKNOWN_MEMBER_NAME,
+                        escapeWithTrailingUnderscore(UnionGenerator.UNKNOWN_MEMBER_NAME));
+        reserveUnionMemberNames(model, reservedNames);
+
+        ReservedWords reservedMembers = new ReservedWordsBuilder()
                 // Since Go only exports names if the first character is upper case and all
                 // the go reserved words are lower case, it's functionally impossible to conflict,
                 // so we only need to protect against common names. As of now there's only one.
@@ -96,7 +102,8 @@ final class SymbolVisitor implements SymbolProvider, ShapeVisitor<Symbol> {
                 .forEach(this::reserveInterfaceMemberAccessors);
 
         escaper = ReservedWordSymbolProvider.builder()
-                .memberReservedWords(reservedWords)
+                .nameReservedWords(reservedNames.build())
+                .memberReservedWords(reservedMembers)
                 // Only escape words when the symbol has a definition file to
                 // prevent escaping intentional references to built-in types.
                 .escapePredicate((shape, symbol) -> !StringUtils.isEmpty(symbol.getDefinitionFile()))
@@ -111,11 +118,28 @@ final class SymbolVisitor implements SymbolProvider, ShapeVisitor<Symbol> {
                 .build();
 
         errorMemberEscaper = ReservedWordSymbolProvider.builder()
-                .memberReservedWords(ReservedWords.compose(reservedWords, reservedErrorMembers))
+                .memberReservedWords(ReservedWords.compose(reservedMembers, reservedErrorMembers))
                 .escapePredicate((shape, symbol) -> !StringUtils.isEmpty(symbol.getDefinitionFile()))
                 .buildEscaper();
     }
 
+    /**
+     * Reserves generated member names for unions.
+     *
+     * <p>These have the format {UnionName}Member{MemberName}.
+     *
+     * @param model The model whose unions should be reserved.
+     * @param builder A reserved words builder to add on to.
+     */
+    private void reserveUnionMemberNames(Model model, ReservedWordsBuilder builder) {
+        model.shapes(UnionShape.class).forEach(union -> {
+            for (MemberShape member : union.getAllMembers().values()) {
+                String memberName = formatUnionMemberName(union, member);
+                builder.put(memberName, escapeWithTrailingUnderscore(memberName));
+            }
+        });
+    }
+
     private boolean supportsInheritance(Shape shape) {
         return shape.isStructureShape() && shape.hasTrait(ErrorTrait.class);
     }
@@ -174,7 +198,12 @@ private Symbol linkArchetypeShape(Shape shape, Symbol symbol) {
 
     @Override
     public String toMemberName(MemberShape shape) {
-        String memberName = escaper.escapeMemberName(getDefaultMemberName(shape));
+        String memberName = getDefaultMemberName(shape);
+        Shape container = model.expectShape(shape.getContainer());
+        if (container.isUnionShape()) {
+            memberName = formatUnionMemberName(container.asUnionShape().get(), shape);
+        }
+        memberName = escaper.escapeMemberName(memberName);
 
         // Escape words reserved for the specific container.
         if (structureSpecificMemberEscapers.containsKey(shape.getContainer())) {
@@ -188,6 +217,10 @@ public String toMemberName(MemberShape shape) {
         return memberName;
     }
 
+    private String formatUnionMemberName(UnionShape union, MemberShape member) {
+        return String.format("%sMember%s", getDefaultShapeName(union), getDefaultMemberName(member));
+    }
+
     private String getDefaultMemberName(MemberShape shape) {
         return StringUtils.capitalize(shape.getMemberName());
     }
@@ -385,7 +418,7 @@ private Optional<String> getNameOfBoundOperation(StructureShape shape) {
     @Override
     public Symbol unionShape(UnionShape shape) {
         String name = getDefaultShapeName(shape);
-        return SymbolUtils.createPointableSymbolBuilder(shape, name, typesPackageName)
+        return SymbolUtils.createValueSymbolBuilder(shape, name, typesPackageName)
                 .definitionFile("./types/types.go")
                 .build();
     }