Fix trait-codegen failure when string literals contain dollar signs

Add defer() and writeLazy() to AbstractCodeWriter, enabling lazy
resolution of format expressions at toString() time without re-parsing
the entire output. Update TraitCodegenWriter to use defer() for type
name placeholders instead of a second format() pass that would
misinterpret literal $ characters in generated code.

Author: sugmanue

smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java

@@ -484,6 +484,9 @@ static Stream<Arguments> loadsModelTests() {
                 Arguments.of("defaults/defaults.smithy",
                         StructDefaultsTrait.class,
                         MapUtils.of("getDefaultString", "default")),
+                Arguments.of("defaults/defaults.smithy",
+                        StructDefaultsTrait.class,
+                        MapUtils.of("getDefaultDollarString", "def$ault")),
                 Arguments.of("defaults/defaults.smithy",
                         StructDefaultsTrait.class,
                         MapUtils.of("getDefaultByte", (byte) 1)),

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/integrations/javadoc/DocumentationTraitInterceptor.java

@@ -19,7 +19,7 @@ final class DocumentationTraitInterceptor implements CodeInterceptor<JavaDocSect
     @Override
     public void write(TraitCodegenWriter writer, String previousText, JavaDocSection section) {
         String docs = section.shape().expectTrait(DocumentationTrait.class).getValue();
-        writer.write(escapeAt(docs));
+        writer.writeWithNoFormatting(escapeAt(docs));
 
         if (!previousText.isEmpty()) {
             // Add spacing if tags have been added to the javadoc

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

@@ -135,9 +135,9 @@ public String toString() {
         builder.append(getPackageHeader()).append(getNewline());
         builder.append(getImportContainer().toString()).append(getNewline());
 
-        // Handle duplicates that may need to use full name
+        // Resolve type name conflicts before toString() resolves deferred values
         resolveNameContext();
-        builder.append(format(super.toString()));
+        builder.append(super.toString());
 
         return builder.toString();
     }
@@ -242,8 +242,9 @@ private String getPlaceholder(Symbol 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}", normalizedSymbol.getFullName());
+            // Return a deferred value that resolves the type name at toString() time
+            String fullName = normalizedSymbol.getFullName();
+            return defer(() -> format("${" + fullName + ":L}"));
         }
 
         // Recursively process the symbols using Java Type.

smithy-trait-codegen/src/test/resources/META-INF/smithy/defaults/defaults.smithy

@@ -16,6 +16,10 @@ structure StructDefaults {
     @default("default")
     defaultString: String
 
+    /// Doc comment that states that this member has a default value of `def$ault`.
+    @default("def$ault")
+    defaultDollarString: String
+
     @default(1)
     defaultByte: Byte
 

smithy-utils/src/main/java/software/amazon/smithy/utils/AbstractCodeWriter.java

@@ -19,6 +19,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.regex.Pattern;
 
 /**
@@ -620,6 +621,8 @@ public abstract class AbstractCodeWriter<T extends AbstractCodeWriter<T>> {
     private boolean trailingNewline = true;
     private int trimBlankLines = -1;
     private boolean enableStackTraceComments;
+    private final Map<String, Supplier<String>> deferredValues = new HashMap<>();
+    private int deferredCounter = 0;
 
     /**
      * Creates a new SimpleCodeWriter that uses "\n" for a newline, four spaces
@@ -784,6 +787,13 @@ public char getExpressionStart() {
     public String toString() {
         String result = currentState.toString();
 
+        // Resolve any deferred values.
+        if (!deferredValues.isEmpty()) {
+            for (Map.Entry<String, Supplier<String>> entry : deferredValues.entrySet()) {
+                result = result.replace(entry.getKey(), entry.getValue().get());
+            }
+        }
+
         // Trim excessive blank lines.
         if (trimBlankLines > -1) {
             StringBuilder builder = new StringBuilder(result.length());
@@ -1756,6 +1766,45 @@ public final String format(Object content, Object... args) {
         return result.toString();
     }
 
+    /**
+     * Registers a deferred value that will be resolved when {@link #toString()} is called.
+     *
+     * <p>This method returns a sentinel string that can be embedded in the writer's output
+     * (for example, as the return value of a custom formatter). When the writer's content is
+     * materialized via {@code toString()}, all sentinels are replaced with the result of
+     * invoking their corresponding supplier.
+     *
+     * <p>This mechanism allows code generators to defer decisions (such as whether to use
+     * a short or fully-qualified type name) until all content has been written, without
+     * requiring a second formatting pass that would re-interpret expression characters
+     * in the output.
+     *
+     * @param supplier A supplier that produces the final string value at resolution time.
+     * @return A sentinel string to embed in the output.
+     */
+    protected final String defer(Supplier<String> supplier) {
+        String id = "\u0000\u0000" + (deferredCounter++) + "\u0000\u0000";
+        deferredValues.put(id, supplier);
+        return id;
+    }
+
+    /**
+     * Writes a lazily-evaluated value to the writer.
+     *
+     * <p>The supplier is not invoked immediately. Instead, a sentinel is written to the
+     * output, and the supplier is invoked when {@link #toString()} is called. This is
+     * useful for writing content that depends on information not yet available at write
+     * time.
+     *
+     * @param supplier A supplier that produces the string value at resolution time.
+     * @return Returns self.
+     */
+    @SuppressWarnings("unchecked")
+    public T writeLazy(Supplier<String> supplier) {
+        writeInlineWithNoFormatting(defer(supplier));
+        return (T) this;
+    }
+
     /**
      * A simple helper method that makes it easier to invoke the built-in {@code C}
      * (call) formatter using a {@link Consumer} where {@code T} is the specific type

smithy-utils/src/test/java/software/amazon/smithy/utils/SimpleCodeWriterTest.java

@@ -1204,4 +1204,92 @@ public void ensuresNewlineIsPresent() {
 
         assertThat(writer.toString(), equalTo("Foo\nBar\nBaz\nBam\n"));
     }
+
+    @Test
+    public void writeLazyResolvesAtToStringTime() {
+        SimpleCodeWriter writer = new SimpleCodeWriter();
+        String[] value = {"initial"};
+        writer.writeInline("Hello ");
+        writer.writeLazy(() -> value[0]);
+        value[0] = "world";
+
+        assertThat(writer.toString(), equalTo("Hello world\n"));
+    }
+
+    @Test
+    public void writeLazyWorksWithSurroundingContent() {
+        SimpleCodeWriter writer = new SimpleCodeWriter();
+        writer.write("before");
+        writer.writeLazy(() -> "lazy");
+        writer.write("after");
+
+        assertThat(writer.toString(), equalTo("before\nlazyafter\n"));
+    }
+
+    @Test
+    public void writeLazyDoesNotReinterpretDollarSigns() {
+        SimpleCodeWriter writer = new SimpleCodeWriter();
+        writer.writeWithNoFormatting("String s = \"$100\";");
+        writer.writeLazy(() -> "resolved");
+
+        assertThat(writer.toString(), equalTo("String s = \"$100\";\nresolved\n"));
+    }
+
+    @Test
+    public void deferCanBeUsedInCustomFormatter() {
+        MyWriter writer = new MyWriter();
+        writer.putFormatter('T', (value, indent) -> writer.defer(() -> "Resolved<" + value + ">"));
+        writer.write("Type: $T", "Foo");
+
+        assertThat(writer.toString(), equalTo("Type: Resolved<Foo>\n"));
+    }
+
+    @Test
+    public void deferredValuesResolveBeforeBlankLineTrimming() {
+        SimpleCodeWriter writer = new SimpleCodeWriter();
+        writer.trimBlankLines();
+        writer.write("before");
+        writer.writeLazy(() -> "lazy");
+        writer.write("after");
+
+        assertThat(writer.toString(), equalTo("before\nlazyafter\n"));
+    }
+
+    @Test
+    public void multipleDeferredValuesResolveIndependently() {
+        SimpleCodeWriter writer = new SimpleCodeWriter();
+        writer.writeInline("A=");
+        writer.writeLazy(() -> "1");
+        writer.writeInline(" B=");
+        writer.writeLazy(() -> "2");
+
+        assertThat(writer.toString(), equalTo("A=1 B=2\n"));
+    }
+
+    @Test
+    public void deferredFormatDoesNotReinterpretDollarSignsInOutput() {
+        // This test reproduces the TraitCodegenWriter problem: a custom $T formatter
+        // uses defer() to lazily resolve type names via context variables. The output
+        // also contains literal $ characters (e.g., from string constants in generated code).
+        // Without defer(), the old approach would re-format the entire output, causing
+        // the literal $ to be misinterpreted as an expression start.
+        MyWriter writer = new MyWriter();
+
+        // Register a formatter that defers type name resolution (like TraitCodegenWriter's $T)
+        writer.putFormatter('T', (value, indent) -> {
+            String fullName = value.toString();
+            return writer.defer(() -> writer.format("${" + fullName + ":L}"));
+        });
+
+        // Set up context that maps full name to short name (like resolveNameContext does)
+        writer.putContext("com.example.MyType", "MyType");
+
+        // Write code that includes both a $T reference AND a literal dollar sign
+        writer.write("public $T getValue() {", "com.example.MyType");
+        writer.writeWithNoFormatting("    return \"costs $100\";");
+        writer.write("}");
+
+        String result = writer.toString();
+        assertThat(result, equalTo("public MyType getValue() {\n    return \"costs $100\";\n}\n"));
+    }
 }