Use deferred comment serialization to fix prominent issues

gson/src/main/java/com/google/gson/stream/JsonWriter.java

@@ -202,6 +202,8 @@ public class JsonWriter implements Closeable, Flushable {
 
   private boolean serializeNulls = true;
 
+  private List<String> deferredComment = new LinkedList<>();
+
   /**
    * Creates a new instance that writes a JSON-encoded stream to {@code out}.
    * For best performance, ensure {@link Writer} is buffered; wrapping in
@@ -302,26 +304,6 @@ public class JsonWriter implements Closeable, Flushable {
     return omitQuotes;
   }
 
-  /**
-   * Insert a comment at the current location.
-   * May create a new line.
-   * This writer MUST be lenient to use this
-   */
-  public JsonWriter comment(String comment) throws IOException {
-    if (!lenient) throw new MalformedJsonException("Cannot write comment in non-lenient JsonWriter.");
-    if (comment == null || comment.isBlank()) return this;
-    String[] parts = comment.split("\n");
-    if (indent == null) {
-      out.append("/* ").append(String.join(" / ", parts)).append(" */");
-    } else {
-      for (String s : parts) {
-        newline();
-        out.append("// ").append(s);
-      }
-    }
-    return this;
-  }
-
   /**
    * Begins encoding a new array. Each call to this method must be paired with
    * a call to {@link #endArray}.
@@ -387,6 +369,12 @@ public class JsonWriter implements Closeable, Flushable {
       throw new IllegalStateException("Dangling name: " + deferredName);
     }
 
+    if (!deferredComment.isEmpty()) {
+      newline();
+      writeDeferredComment();
+      context = nonempty;
+    }
+
     stackSize--;
     if (context == nonempty) {
       newline();
@@ -419,6 +407,39 @@ public class JsonWriter implements Closeable, Flushable {
     stack[stackSize - 1] = topOfStack;
   }
 
+  /**
+   * Insert a comment at the current location.
+   * May create a new line.
+   * This writer MUST be lenient to use this
+   */
+  public JsonWriter comment(String comment) throws IOException {
+    if (!lenient) throw new MalformedJsonException("Cannot write comment in non-lenient JsonWriter.");
+    if (comment == null || comment.isBlank()) return this;
+    String[] parts = comment.split("\n");
+    Collections.addAll(deferredComment, parts);
+    if (peek() == NONEMPTY_DOCUMENT) {
+      newline();
+      writeDeferredComment();
+    }
+    return this;
+  }
+
+  private void writeDeferredComment() throws IOException {
+    if (!deferredComment.isEmpty()) {
+      if (indent == null) {
+        out.append("/* ").append(String.join(" / ", deferredComment)).append(" */");
+      } else {
+        boolean first = true;
+        for (String s : deferredComment) {
+          if (!first) newline();
+          first = false;
+          out.append("// ").append(s);
+        }
+      }
+      deferredComment.clear();
+    }
+  }
+
   /**
    * Encodes the property name.
    *
@@ -689,6 +710,10 @@ public class JsonWriter implements Closeable, Flushable {
       throw new IllegalStateException("Nesting problem.");
     }
     newline();
+    if (!deferredComment.isEmpty()) {
+      writeDeferredComment();
+      newline();
+    }
     replaceTop(DANGLING_NAME);
   }
 
@@ -702,26 +727,42 @@ public class JsonWriter implements Closeable, Flushable {
     switch (peek()) {
     case NONEMPTY_DOCUMENT:
       if (!lenient) {
-        throw new IllegalStateException(
-            "JSON must have only one top-level value.");
+        throw new IllegalStateException("JSON must have only one top-level value.");
       }
       // fall-through
     case EMPTY_DOCUMENT: // first in document
       replaceTop(NONEMPTY_DOCUMENT);
+      if (!deferredComment.isEmpty()) {
+        writeDeferredComment();
+        newline();
+      }
       break;
 
     case EMPTY_ARRAY: // first in array
       replaceTop(NONEMPTY_ARRAY);
       newline();
+      if (!deferredComment.isEmpty()) {
+        writeDeferredComment();
+        newline();
+      }
       break;
 
     case NONEMPTY_ARRAY: // another in array
       out.append(',');
       newline();
+      if (!deferredComment.isEmpty()) {
+        writeDeferredComment();
+        newline();
+      }
       break;
 
     case DANGLING_NAME: // value for name
       out.append(separator);
+      if (!deferredComment.isEmpty()) {
+        newline();
+        writeDeferredComment();
+        newline();
+      }
       replaceTop(NONEMPTY_OBJECT);
       break;
 

gson/src/test/java/com/google/gson/jf/CommentsTest.java

@@ -17,9 +17,14 @@
 package com.google.gson.jf;
 
 import com.google.gson.*;
+import com.google.gson.common.*;
 import com.google.gson.reflect.TypeToken;
+
+import java.io.*;
 import java.util.Arrays;
 import java.util.List;
+
+import com.google.gson.stream.*;
 import junit.framework.TestCase;
 
 /**
@@ -43,4 +48,45 @@ public final class CommentsTest extends TestCase {
     List<String> abc = new GsonBuilder().setLenient().create().fromJson(json, new TypeToken<List<String>>() {}.getType());
     assertEquals(Arrays.asList("a", "b", "c"), abc);
   }
+
+  public void testWriteComments() throws IOException {
+    String expectedJson = "// comment at file head\n"
+        + "[\n"
+        + "  // comment directly after context\n"
+        + "  \"a\",\n"
+        + "  // comment before context\n"
+        + "  {\n"
+        + "    // comment directly after object context\n"
+        + "    \"b\": \"c\",\n"
+        + "    // comment before name, after value\n"
+        + "    \"d\": \"e\"\n"
+        + "  }\n"
+        + "  // comment before context end\n"
+        + "]\n"
+        + "// comment behind the object";
+
+    StringWriter sw = new StringWriter();
+    JsonWriter jw = new JsonWriter(sw);
+    jw.setLenient(true);
+    jw.setIndent("  ");
+
+    jw.comment("comment at file head")
+      .beginArray()
+        .comment("comment directly after context")
+        .value("a")
+        .comment("comment before context")
+        .beginObject()
+          .comment("comment directly after object context")
+          .name("b").value("c")
+          .comment("comment before name, after value")
+          .name("d").value("e")
+        .endObject()
+        .comment("comment before context end")
+      .endArray()
+      .comment("comment behind the object");
+
+    jw.close();
+    assertEquals(expectedJson, sw.toString());
+    sw.close();
+  }
 }