Update TraitCodegenWriter (#2255)

Update TraitCodegenWriter to address a few minor issues and improve ergonomics. 
Changes include:
1. updates to writer to use a static factory class
2. adding a CapitalizingFormatter formatter for $U
3. Normalizing symbol references before adding to duplicate map to avoid false positive duplicates 
4. Putting resolved placeholder values directly into context map rather than instantiating a second hashmap.

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/TraitCodegenContext.java

@@ -40,7 +40,7 @@ public final class TraitCodegenContext implements CodegenContext<TraitCodegenSet
         this.fileManifest = fileManifest;
         this.integrations = integrations;
         this.writerDelegator = new WriterDelegator<>(fileManifest, symbolProvider,
-                (filename, namespace) -> new TraitCodegenWriter(filename, namespace, settings));
+                new TraitCodegenWriter.Factory(settings));
     }
 
     @Override

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

@@ -26,7 +26,6 @@
 import software.amazon.smithy.traitcodegen.writer.TraitCodegenWriter;
 import software.amazon.smithy.utils.BuilderRef;
 import software.amazon.smithy.utils.SmithyBuilder;
-import software.amazon.smithy.utils.StringUtils;
 
 /**
  * Generates a static builder for a Java class.
@@ -243,29 +242,27 @@ protected Void getDefault(Shape shape) {
         public Void listShape(ListShape shape) {
             writer.openBlock("public Builder $1L($2B $1L) {", "}",
                     memberName, symbolProvider.toSymbol(shape), () -> {
-                        writer.write("clear$L();", StringUtils.capitalize(memberName));
+                        writer.write("clear$U();", memberName);
                         writer.write("this.$1L.get().addAll($1L);", memberName);
                         writer.writeWithNoFormatting("return this;");
                     }).newLine();
 
             // Clear all
-            writer.openBlock("public Builder clear$L() {", "}",
-                    StringUtils.capitalize(memberName), () -> {
+            writer.openBlock("public Builder clear$U() {", "}", memberName, () -> {
                         writer.write("$L.get().clear();", memberName);
                         writer.writeWithNoFormatting("return this;");
                     }).newLine();
 
             // Set one
-            writer.openBlock("public Builder add$L($T value) {", "}",
-                    StringUtils.capitalize(memberName), symbolProvider.toSymbol(shape.getMember()),
-                    () -> {
+            writer.openBlock("public Builder add$U($T value) {", "}",
+                    memberName, symbolProvider.toSymbol(shape.getMember()), () -> {
                         writer.write("$L.get().add(value);", memberName);
                         writer.write("return this;");
                     }).newLine();
 
             // Remove one
-            writer.openBlock("public Builder remove$L($T value) {", "}",
-                    StringUtils.capitalize(memberName), symbolProvider.toSymbol(shape.getMember()),
+            writer.openBlock("public Builder remove$U($T value) {", "}",
+                    memberName, symbolProvider.toSymbol(shape.getMember()),
                     () -> {
                         writer.write("$L.get().remove(value);", memberName);
                         writer.write("return this;");
@@ -278,31 +275,30 @@ public Void mapShape(MapShape shape) {
             // Set all
             writer.openBlock("public Builder $1L($2B $1L) {", "}",
                     memberName, symbolProvider.toSymbol(shape), () -> {
-                        writer.write("clear$L();", StringUtils.capitalize(memberName));
+                        writer.write("clear$U();", memberName);
                         writer.write("this.$1L.get().putAll($1L);", memberName);
                         writer.write("return this;");
                     });
             writer.newLine();
 
             // Clear all
-            writer.openBlock("public Builder clear$L() {", "}", StringUtils.capitalize(memberName), () -> {
+            writer.openBlock("public Builder clear$U() {", "}", memberName, () -> {
                 writer.write("this.$L.get().clear();", memberName);
                 writer.write("return this;");
             }).newLine();
 
             // Set one
             MemberShape keyShape = shape.getKey();
             MemberShape valueShape = shape.getValue();
-            writer.openBlock("public Builder put$L($T key, $T value) {", "}",
-                    StringUtils.capitalize(memberName), symbolProvider.toSymbol(keyShape),
-                    symbolProvider.toSymbol(valueShape), () -> {
+            writer.openBlock("public Builder put$U($T key, $T value) {", "}",
+                    memberName, symbolProvider.toSymbol(keyShape), symbolProvider.toSymbol(valueShape), () -> {
                         writer.write("this.$L.get().put(key, value);", memberName);
                         writer.write("return this;");
                     }).newLine();
 
             // Remove one
-            writer.openBlock("public Builder remove$L($T $L) {", "}",
-                    StringUtils.capitalize(memberName), symbolProvider.toSymbol(keyShape), memberName, () -> {
+            writer.openBlock("public Builder remove$U($T $L) {", "}",
+                    memberName, symbolProvider.toSymbol(keyShape), memberName, () -> {
                         writer.write("this.$1L.get().remove($1L);", memberName);
                         writer.write("return this;");
                     }).newLine();

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/GetterGenerator.java

@@ -24,7 +24,6 @@
 import software.amazon.smithy.traitcodegen.TraitCodegenUtils;
 import software.amazon.smithy.traitcodegen.sections.GetterSection;
 import software.amazon.smithy.traitcodegen.writer.TraitCodegenWriter;
-import software.amazon.smithy.utils.StringUtils;
 
 /**
  * Generates getter methods for each shape member or the value type held by the trait.
@@ -118,16 +117,15 @@ public Void structureShape(StructureShape shape) {
                 // then do not wrap return in an Optional
                 writer.pushState(new GetterSection(member));
                 if (member.isRequired()) {
-                    writer.openBlock("public $T get$L() {", "}",
+                    writer.openBlock("public $T get$U() {", "}",
                             symbolProvider.toSymbol(member),
-                            StringUtils.capitalize(symbolProvider.toMemberName(member)),
+                            symbolProvider.toMemberName(member),
                             () -> writer.write("return $L;", symbolProvider.toMemberName(member)));
                     writer.popState();
                     writer.newLine();
                 } else {
-                    writer.openBlock("public $T<$T> get$L() {", "}",
-                            Optional.class, symbolProvider.toSymbol(member),
-                            StringUtils.capitalize(symbolProvider.toMemberName(member)),
+                    writer.openBlock("public $T<$T> get$U() {", "}",
+                            Optional.class, symbolProvider.toSymbol(member), symbolProvider.toMemberName(member),
                             () -> writer.write("return $T.ofNullable($L);",
                                     Optional.class, symbolProvider.toMemberName(member)));
                     writer.popState();
@@ -137,9 +135,9 @@ public Void structureShape(StructureShape shape) {
                     // getter as a convenience method as well.
                     Shape target = model.expectShape(member.getTarget());
                     if (target.isListShape() || target.isMapShape()) {
-                        writer.openBlock("public $T get$LOrEmpty() {", "}",
+                        writer.openBlock("public $T get$UOrEmpty() {", "}",
                                 symbolProvider.toSymbol(member),
-                                StringUtils.capitalize(symbolProvider.toMemberName(member)),
+                                symbolProvider.toMemberName(member),
                                 () -> writer.write("return $L;", symbolProvider.toMemberName(member)));
                     }
                 }

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/ToNodeGenerator.java

@@ -34,7 +34,6 @@
 import software.amazon.smithy.model.traits.TraitDefinition;
 import software.amazon.smithy.traitcodegen.TraitCodegenUtils;
 import software.amazon.smithy.traitcodegen.writer.TraitCodegenWriter;
-import software.amazon.smithy.utils.StringUtils;
 
 /**
  * Generates methods to serialize a Java class to a smithy {@code Node}.
@@ -145,8 +144,8 @@ public Void structureShape(StructureShape shape) {
                             mem.getMemberName(),
                             (Runnable) () -> mem.accept(new ToNodeMapperVisitor(symbolProvider.toMemberName(mem))));
                 } else {
-                    writer.write(".withOptionalMember($S, get$L().map(m -> $C))",
-                            mem.getMemberName(), StringUtils.capitalize(symbolProvider.toMemberName(mem)),
+                    writer.write(".withOptionalMember($S, get$U().map(m -> $C))",
+                            mem.getMemberName(), symbolProvider.toMemberName(mem),
                             (Runnable) () -> mem.accept(new ToNodeMapperVisitor("m")));
                 }
             }

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/writer/TraitCodegenWriter.java

@@ -20,24 +20,28 @@
 import software.amazon.smithy.traitcodegen.SymbolProperties;
 import software.amazon.smithy.traitcodegen.TraitCodegenSettings;
 import software.amazon.smithy.traitcodegen.TraitCodegenUtils;
+import software.amazon.smithy.utils.ListUtils;
 import software.amazon.smithy.utils.StringUtils;
 
 /**
  * Writes Java code for trait definitions.
  *
  * <p>This writer supports two custom formatters, a Java type formatter '$T' and
  * a Base type formatter '$B'.
- * <ul>
- *     <li>{@link JavaTypeFormatter}|{@code 'T'}: This formatter handles the formatting of
+ * <dl>
+ *     <dt>{@link JavaTypeFormatter}|{@code 'T'}</dt>
+ *     <dd>This formatter handles the formatting of
  *     Java types and also ensures that parameterized types (such as {@code List<String>} are
- *     written correctly.
- *
- *     <li>{@link BaseTypeFormatter}|{@code 'B'}: This formatter allows you to use the base type
+ *     written correctly.</dd>
+ *     <dt>{@link BaseTypeFormatter}|{@code 'B'}</dt>
+ *     <dd>This formatter allows you to use the base type
  *     for a trait. For example a String Trait may have a base type of {@code ShapeId}. To write
  *     this base type, use the {@code $B} formatter and provide the trait symbol. Note that
  *     if no base type is found (i.e. type is not a trait) then this formatter behaves exactly the
- *     same as the {@link JavaTypeFormatter}.
- * </ul>
+ *     same as the {@link JavaTypeFormatter}.</dd>
+ *     <dt>{@link CapitalizingFormatter}|{@code 'U'}</dt>
+ *     <dd>This formatter will capitalize the first letter of any string literal it is used to format.</dd>
+ * </dl>
  */
 public class TraitCodegenWriter extends SymbolWriter<TraitCodegenWriter, TraitCodegenImportContainer> {
     private static final int MAX_LINE_LENGTH = 120;
@@ -62,6 +66,7 @@ public TraitCodegenWriter(String fileName,
 
         putFormatter('T', new JavaTypeFormatter());
         putFormatter('B', new BaseTypeFormatter());
+        putFormatter('U', new CapitalizingFormatter());
     }
 
 
@@ -118,14 +123,13 @@ public String toString() {
         builder.append(getImportContainer().toString()).append(getNewline());
 
         // Handle duplicates that may need to use full name
-        putContext(resolveNameContext());
+        resolveNameContext();
         builder.append(format(super.toString()));
 
         return builder.toString();
     }
 
-    private Map<String, Object> resolveNameContext() {
-        Map<String, Object> contextMap = new HashMap<>();
+    private void resolveNameContext() {
         for (Map.Entry<String, Set<Symbol>> entry : symbolNames.entrySet()) {
             Set<Symbol> duplicates = entry.getValue();
             // If the duplicates list has more than one entry
@@ -135,18 +139,16 @@ private Map<String, Object> resolveNameContext() {
                     // If we are in the namespace of a Symbol, use its
                     // short name, otherwise use the full name
                     if (dupe.getNamespace().equals(namespace)) {
-                        contextMap.put(dupe.getFullName(), dupe.getName());
+                        putContext(dupe.getFullName(), dupe.getName());
                     } else {
-                        contextMap.put(dupe.getFullName(), dupe.getFullName());
+                        putContext(dupe.getFullName(), dupe.getFullName());
                     }
                 });
             } else {
                 Symbol symbol = duplicates.iterator().next();
-                contextMap.put(symbol.getFullName(), symbol.getName());
+                putContext(symbol.getFullName(), symbol.getName());
             }
         }
-
-        return contextMap;
     }
 
     public String getPackageHeader() {
@@ -170,6 +172,26 @@ public void override() {
         writeWithNoFormatting("@Override");
     }
 
+    /**
+     * A factory class to create {@link TraitCodegenWriter}s.
+     */
+    public static final class Factory implements SymbolWriter.Factory<TraitCodegenWriter> {
+
+        private final TraitCodegenSettings settings;
+
+        /**
+         * @param settings The Trait codegen plugin settings.
+         */
+        public Factory(TraitCodegenSettings settings) {
+            this.settings = settings;
+        }
+
+        @Override
+        public TraitCodegenWriter apply(String filename, String namespace) {
+            return new TraitCodegenWriter(filename, namespace, settings);
+        }
+    }
+
     /**
      * Implements a formatter for {@code $T} that formats Java types.
      */
@@ -207,15 +229,17 @@ public String apply(Object type, String indent) {
         }
 
         private String getPlaceholder(Symbol symbol) {
+            Symbol normalizedSymbol = symbol.toBuilder().references(ListUtils.of()).build();
+
             // Add symbol to import container
-            addImport(symbol);
+            addImport(normalizedSymbol);
 
             // Add symbol to symbol map, so we can handle potential type name conflicts
-            Set<Symbol> nameSet = symbolNames.computeIfAbsent(symbol.getName(), n -> new HashSet<>());
-            nameSet.add(symbol);
+            Set<Symbol> nameSet = symbolNames.computeIfAbsent(normalizedSymbol.getName(), n -> new HashSet<>());
+            nameSet.add(normalizedSymbol);
 
             // Return a placeholder value that will be filled when toString is called
-            return format("$${$L:L}", symbol.getFullName());
+            return format("$${$L:L}", normalizedSymbol.getFullName());
         }
     }
 
@@ -239,4 +263,21 @@ public String apply(Object type, String indent) {
             return javaTypeFormatter.apply(symbol, indent);
         }
     }
+
+    /**
+     * Implements a formatter for {@code $U} that capitalizes the first letter of a string literal.
+     */
+    private static final class CapitalizingFormatter implements BiFunction<Object, String, String> {
+        @Override
+        public String apply(Object type, String indent) {
+            if (type instanceof String) {
+                return StringUtils.capitalize((String) type);
+            }
+            throw new IllegalArgumentException(
+                    "Invalid type provided for $U. Expected a String but found: `"
+                            + type + "`."
+            );
+        }
+    }
+
 }