Ensure trailing newlines in javadocs and method bodies (#732)

* Add RecordingAppendable in LineWrapper for tracking last emitted char

* Check lastChar in javadoc emission to emit newline if necessary

Resolves #731

* Move trailing newline check to emit() overload for reuse

Allows using from anywhere emitting a CodeBlock

* Ensure trailing newlines in method bodies

Resolves #722

* Add dedicated trailing newline in javadoc test

* Fix modifier ordering

* Fix rebase conflict

Co-authored-by: Egor Andreevich <andreevich.egor@gmail.com>

5 files changed, 100 insertions, 5 deletions

diff --git a/src/main/java/com/squareup/javapoet/CodeWriter.java b/src/main/java/com/squareup/javapoet/CodeWriter.java
index b2b088b..b2bc3d3 100644
--- a/src/main/java/com/squareup/javapoet/CodeWriter.java
+++ b/src/main/java/com/squareup/javapoet/CodeWriter.java
@@ -149,7 +149,7 @@ final class CodeWriter {
     emit("/**\n");
     javadoc = true;
     try {
-      emit(javadocCodeBlock);
+      emit(javadocCodeBlock, true);
     } finally {
       javadoc = false;
     }
@@ -219,6 +219,10 @@ final class CodeWriter {
   }
 
   public CodeWriter emit(CodeBlock codeBlock) throws IOException {
+    return emit(codeBlock, false);
+  }
+
+  public CodeWriter emit(CodeBlock codeBlock, boolean ensureTrailingNewline) throws IOException {
     int a = 0;
     ClassName deferredTypeName = null; // used by "import static" logic
     ListIterator<String> partIterator = codeBlock.formatParts.listIterator();
@@ -307,6 +311,9 @@ final class CodeWriter {
           break;
       }
     }
+    if (ensureTrailingNewline && out.lastChar() != '\n') {
+      emit("\n");
+    }
     return this;
   }
 
diff --git a/src/main/java/com/squareup/javapoet/LineWrapper.java b/src/main/java/com/squareup/javapoet/LineWrapper.java
index 6aa3131..928d9f4 100644
--- a/src/main/java/com/squareup/javapoet/LineWrapper.java
+++ b/src/main/java/com/squareup/javapoet/LineWrapper.java
@@ -24,7 +24,7 @@ import static com.squareup.javapoet.Util.checkNotNull;
  * or soft-wrapping spaces using {@link #wrappingSpace}.
  */
 final class LineWrapper {
-  private final Appendable out;
+  private final RecordingAppendable out;
   private final String indent;
   private final int columnLimit;
   private boolean closed;
@@ -47,11 +47,16 @@ final class LineWrapper {
 
   LineWrapper(Appendable out, String indent, int columnLimit) {
     checkNotNull(out, "out == null");
-    this.out = out;
+    this.out = new RecordingAppendable(out);
     this.indent = indent;
     this.columnLimit = columnLimit;
   }
 
+  /** @return the last emitted char or {@link Character#MIN_VALUE} if nothing emitted yet. */
+  char lastChar() {
+    return out.lastChar;
+  }
+
   /** Emit {@code s}. This may be buffered to permit line wraps to be inserted. */
   void append(String s) throws IOException {
     if (closed) throw new IllegalStateException("closed");
@@ -134,4 +139,33 @@ final class LineWrapper {
   private enum FlushType {
     WRAP, SPACE, EMPTY;
   }
+
+  /** A delegating {@link Appendable} that records info about the chars passing through it. */
+  static final class RecordingAppendable implements Appendable {
+    private final Appendable delegate;
+
+    char lastChar = Character.MIN_VALUE;
+
+    RecordingAppendable(Appendable delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override public Appendable append(CharSequence csq) throws IOException {
+      int length = csq.length();
+      if (length != 0) {
+        lastChar = csq.charAt(length - 1);
+      }
+      return delegate.append(csq);
+    }
+
+    @Override public Appendable append(CharSequence csq, int start, int end) throws IOException {
+      CharSequence sub = csq.subSequence(start, end);
+      return append(sub);
+    }
+
+    @Override public Appendable append(char c) throws IOException {
+      lastChar = c;
+      return delegate.append(c);
+    }
+  }
 }
diff --git a/src/main/java/com/squareup/javapoet/MethodSpec.java b/src/main/java/com/squareup/javapoet/MethodSpec.java
index b06290f..2f9be0e 100644
--- a/src/main/java/com/squareup/javapoet/MethodSpec.java
+++ b/src/main/java/com/squareup/javapoet/MethodSpec.java
@@ -132,7 +132,7 @@ public final class MethodSpec {
       codeWriter.emit(" {\n");
 
       codeWriter.indent();
-      codeWriter.emit(code);
+      codeWriter.emit(code, true);
       codeWriter.unindent();
 
       codeWriter.emit("}\n");
diff --git a/src/test/java/com/squareup/javapoet/MethodSpecTest.java b/src/test/java/com/squareup/javapoet/MethodSpecTest.java
index 789661d..8704ec1 100644
--- a/src/test/java/com/squareup/javapoet/MethodSpecTest.java
+++ b/src/test/java/com/squareup/javapoet/MethodSpecTest.java
@@ -393,4 +393,27 @@ public final class MethodSpecTest {
     builder.typeVariables.remove(1);
     assertThat(builder.build().typeVariables).containsExactly(t);
   }
+
+  @Test public void ensureTrailingNewline() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("method")
+        .addCode("codeWithNoNewline();")
+        .build();
+
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "void method() {\n"
+        + "  codeWithNoNewline();\n"
+        + "}\n");
+  }
+
+  /** Ensures that we don't add a duplicate newline if one is already present. */
+  @Test public void ensureTrailingNewlineWithExistingNewline() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("method")
+        .addCode("codeWithNoNewline();\n") // Have a newline already, so ensure we're not adding one
+        .build();
+
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "void method() {\n"
+        + "  codeWithNoNewline();\n"
+        + "}\n");
+  }
 }
diff --git a/src/test/java/com/squareup/javapoet/TypeSpecTest.java b/src/test/java/com/squareup/javapoet/TypeSpecTest.java
index 6175607..f76e6f0 100644
--- a/src/test/java/com/squareup/javapoet/TypeSpecTest.java
+++ b/src/test/java/com/squareup/javapoet/TypeSpecTest.java
@@ -1852,7 +1852,8 @@ public final class TypeSpecTest {
         + "  }\n"
         + "\n"
         + "  /**\n"
-        + "   * chosen by fair dice roll ;) */\n"
+        + "   * chosen by fair dice roll ;)\n"
+        + "   */\n"
         + "  public int getRandomQuantity() {\n"
         + "    return 4;\n"
         + "  }\n"
@@ -2516,4 +2517,34 @@ public final class TypeSpecTest {
     builder.originatingElements.clear();
     assertThat(builder.build().originatingElements).isEmpty();
   }
+    
+  @Test public void javadocWithTrailingLineDoesNotAddAnother() {
+    TypeSpec spec = TypeSpec.classBuilder("Taco")
+        .addJavadoc("Some doc with a newline\n")
+        .build();
+
+    assertThat(toString(spec)).isEqualTo(""
+        + "package com.squareup.tacos;\n"
+        + "\n"
+        + "/**\n"
+        + " * Some doc with a newline\n"
+        + " */\n"
+        + "class Taco {\n"
+        + "}\n");
+  }
+
+  @Test public void javadocEnsuresTrailingLine() {
+    TypeSpec spec = TypeSpec.classBuilder("Taco")
+        .addJavadoc("Some doc with a newline")
+        .build();
+
+    assertThat(toString(spec)).isEqualTo(""
+        + "package com.squareup.tacos;\n"
+        + "\n"
+        + "/**\n"
+        + " * Some doc with a newline\n"
+        + " */\n"
+        + "class Taco {\n"
+        + "}\n");
+  }
 }