From a575f221a9f238a0c2a094cf9531b577ae706b0b Mon Sep 17 00:00:00 2001
From: Michael Dowling <mtdowling@gmail.com>
Date: Thu, 10 Mar 2022 12:53:29 -0800
Subject: [PATCH] Ensure code writer always has a StringBuilder

There was an edge case where contents written to nested sections of
captured sections would be swallowed due to how code writer lazily
created StringBuilders. This change ensures that the root writer and
every interceptable writer always define a StringBuilder to ensure no
text is omitted.
---
 .../smithy/utils/AbstractCodeWriter.java       |  9 +++------
 .../amazon/smithy/utils/CodeWriterTest.java    | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/smithy-utils/src/main/java/software/amazon/smithy/utils/AbstractCodeWriter.java b/smithy-utils/src/main/java/software/amazon/smithy/utils/AbstractCodeWriter.java
index e9348b3c670..dadd173501e 100644
--- a/smithy-utils/src/main/java/software/amazon/smithy/utils/AbstractCodeWriter.java
+++ b/smithy-utils/src/main/java/software/amazon/smithy/utils/AbstractCodeWriter.java
@@ -1766,7 +1766,6 @@ private final class State {
         private CopyOnWriteRef<CodeWriterFormatterContainer> formatters;
         private CopyOnWriteRef<CodeInterceptorContainer<T>> interceptors;
 
-        /** This StringBuilder, if null, will only be created lazily when needed. */
         private StringBuilder builder;
 
         /**
@@ -1777,6 +1776,7 @@ private final class State {
         private boolean isInline;
 
         State() {
+            builder = new StringBuilder();
             isRoot = true;
             CodeWriterFormatterContainer formatterContainer = new CodeWriterFormatterContainer();
             DEFAULT_FORMATTERS.forEach(formatterContainer::putFormatter);
@@ -1810,13 +1810,10 @@ private void copyStateFrom(State copy) {
 
         @Override
         public String toString() {
-            return builder == null ? "" : builder.toString();
+            return getBuilder().toString();
         }
 
         StringBuilder getBuilder() {
-            if (builder == null) {
-                builder = new StringBuilder();
-            }
             return builder;
         }
 
@@ -1921,7 +1918,7 @@ private void makeInterceptableCodeSection(CodeSection section) {
             // level is reset back to the root, and the newline prefix is removed.
             // Indentation and prefixes are added automatically if/when the
             // captured text is written into the parent state.
-            currentState.builder = null;
+            currentState.builder = new StringBuilder();
             currentState.newlinePrefix = "";
             dedent(-1);
         }
diff --git a/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java b/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java
index 707e3d59707..f0793df8c54 100644
--- a/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java
+++ b/smithy-utils/src/test/java/software/amazon/smithy/utils/CodeWriterTest.java
@@ -1030,4 +1030,22 @@ public void ensuresNewlineIsPresent() {
 
         assertThat(writer.toString(), equalTo("Foo\nBar\nBaz\nBam\n"));
     }
+
+    // Dynamically creating builders for a section doesn't work when pushing
+    // sections into captured sections that haven't written anything yet.
+    // It ends up swallowing the text written to sub-sections because the
+    // top-level section's builder wasn't written yet. This test makes sure
+    // that doesn't happen.
+    @Test
+    public void alwaysWritesToParentBuilders() {
+        SimpleCodeWriter writer = new SimpleCodeWriter();
+
+        writer.pushState("Hi");
+        writer.pushState();
+        writer.write("Hello");
+        writer.popState();
+        writer.popState();
+
+        assertThat(writer.toString(), equalTo("Hello\n"));
+    }
 }