Re-indent text blocks

PiperOrigin-RevId: 609526784

Author: cushon

core/src/main/java/com/google/googlejavaformat/java/JavaInputAstVisitor.java

@@ -1667,6 +1667,9 @@ public Void visitMemberSelect(MemberSelectTree node, Void unused) {
   public Void visitLiteral(LiteralTree node, Void unused) {
     sync(node);
     String sourceForNode = getSourceForNode(node, getCurrentPath());
+    if (sourceForNode.startsWith("\"\"\"")) {
+      builder.forcedBreak();
+    }
     if (isUnaryMinusLiteral(sourceForNode)) {
       token("-");
       sourceForNode = sourceForNode.substring(1).trim();

core/src/main/java/com/google/googlejavaformat/java/StringWrapper.java

@@ -14,6 +14,7 @@
 
 package com.google.googlejavaformat.java;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Iterables.getLast;
 import static java.lang.Math.min;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -44,6 +45,7 @@
 import com.sun.tools.javac.util.Position;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.lang.reflect.Method;
 import java.net.URI;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
@@ -59,6 +61,7 @@
 import javax.tools.JavaFileObject;
 import javax.tools.SimpleJavaFileObject;
 import javax.tools.StandardLocation;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Wraps string literals that exceed the column limit. */
 public final class StringWrapper {
@@ -72,7 +75,7 @@ public static String wrap(String input, Formatter formatter) throws FormatterExc
    */
   static String wrap(final int columnLimit, String input, Formatter formatter)
       throws FormatterException {
-    if (!longLines(columnLimit, input)) {
+    if (!needWrapping(columnLimit, input)) {
       // fast path
       return input;
     }
@@ -111,20 +114,56 @@ static String wrap(final int columnLimit, String input, Formatter formatter)
 
   private static TreeRangeMap<Integer, String> getReflowReplacements(
       int columnLimit, final String input) throws FormatterException {
-    JCTree.JCCompilationUnit unit = parse(input, /* allowStringFolding= */ false);
-    String separator = Newlines.guessLineSeparator(input);
+    return new Reflower(columnLimit, input).getReflowReplacements();
+  }
+
+  private static class Reflower {
+
+    private final String input;
+    private final int columnLimit;
+    private final String separator;
+    private final JCTree.JCCompilationUnit unit;
+    private final Position.LineMap lineMap;
+
+    Reflower(int columnLimit, String input) throws FormatterException {
+      this.columnLimit = columnLimit;
+      this.input = input;
+      this.separator = Newlines.guessLineSeparator(input);
+      this.unit = parse(input, /* allowStringFolding= */ false);
+      this.lineMap = unit.getLineMap();
+    }
+
+    TreeRangeMap<Integer, String> getReflowReplacements() {
+      // Paths to string literals that extend past the column limit.
+      List<TreePath> longStringLiterals = new ArrayList<>();
+      // Paths to text blocks to be re-indented.
+      List<Tree> textBlocks = new ArrayList<>();
+      new LongStringsAndTextBlockScanner(longStringLiterals, textBlocks)
+          .scan(new TreePath(unit), null);
+      TreeRangeMap<Integer, String> replacements = TreeRangeMap.create();
+      indentTextBlocks(replacements, textBlocks);
+      wrapLongStrings(replacements, longStringLiterals);
+      return replacements;
+    }
+
+    private class LongStringsAndTextBlockScanner extends TreePathScanner<Void, Void> {
+
+      private final List<TreePath> longStringLiterals;
+      private final List<Tree> textBlocks;
+
+      LongStringsAndTextBlockScanner(List<TreePath> longStringLiterals, List<Tree> textBlocks) {
+        this.longStringLiterals = longStringLiterals;
+        this.textBlocks = textBlocks;
+      }
 
-    // Paths to string literals that extend past the column limit.
-    List<TreePath> toFix = new ArrayList<>();
-    final Position.LineMap lineMap = unit.getLineMap();
-    new TreePathScanner<Void, Void>() {
       @Override
       public Void visitLiteral(LiteralTree literalTree, Void aVoid) {
         if (literalTree.getKind() != Kind.STRING_LITERAL) {
           return null;
         }
         int pos = getStartPosition(literalTree);
         if (input.substring(pos, min(input.length(), pos + 3)).equals("\"\"\"")) {
+          textBlocks.add(literalTree);
           return null;
         }
         Tree parent = getCurrentPath().getParentPath().getLeaf();
@@ -140,44 +179,114 @@ public Void visitLiteral(LiteralTree literalTree, Void aVoid) {
         if (lineMap.getColumnNumber(lineEnd) - 1 <= columnLimit) {
           return null;
         }
-        toFix.add(getCurrentPath());
+        longStringLiterals.add(getCurrentPath());
         return null;
       }
-    }.scan(new TreePath(unit), null);
-
-    TreeRangeMap<Integer, String> replacements = TreeRangeMap.create();
-    for (TreePath path : toFix) {
-      // Find the outermost contiguous enclosing concatenation expression
-      TreePath enclosing = path;
-      while (enclosing.getParentPath().getLeaf().getKind() == Tree.Kind.PLUS) {
-        enclosing = enclosing.getParentPath();
+    }
+
+    private void indentTextBlocks(
+        TreeRangeMap<Integer, String> replacements, List<Tree> textBlocks) {
+      for (Tree tree : textBlocks) {
+        int startPosition = getStartPosition(tree);
+        int endPosition = getEndPosition(unit, tree);
+        String text = input.substring(startPosition, endPosition);
+
+        // Find the source code of the text block with incidental whitespace removed.
+        // The first line of the text block is always """, and it does not affect incidental
+        // whitespace.
+        ImmutableList<String> initialLines = text.lines().collect(toImmutableList());
+        String stripped = stripIndent(initialLines.stream().skip(1).collect(joining(separator)));
+        ImmutableList<String> lines = stripped.lines().collect(toImmutableList());
+        int deindent =
+            initialLines.get(1).stripTrailing().length() - lines.get(0).stripTrailing().length();
+
+        int startColumn = lineMap.getColumnNumber(startPosition);
+        String prefix =
+            (deindent == 0 || lines.stream().anyMatch(x -> x.length() + startColumn > columnLimit))
+                ? ""
+                : " ".repeat(startColumn - 1);
+
+        StringBuilder output = new StringBuilder("\"\"\"");
+        for (int i = 0; i < lines.size(); i++) {
+          String line = lines.get(i);
+          String trimmed = line.stripLeading().stripTrailing();
+          output.append(separator);
+          if (!trimmed.isEmpty()) {
+            // Don't add incidental leading whitespace to empty lines
+            output.append(prefix);
+          }
+          if (i == lines.size() - 1 && trimmed.equals("\"\"\"")) {
+            // If the trailing line is just """, indenting is more than the prefix of incidental
+            // whitespace has no effect, and results in a javac text-blocks warning that 'trailing
+            // white space will be removed'.
+            output.append("\"\"\"");
+          } else {
+            output.append(line);
+          }
+        }
+        replacements.put(Range.closedOpen(startPosition, endPosition), output.toString());
       }
-      // Is the literal being wrapped the first in a chain of concatenation expressions?
-      // i.e. `ONE + TWO + THREE`
-      // We need this information to handle continuation indents.
-      AtomicBoolean first = new AtomicBoolean(false);
-      // Finds the set of string literals in the concat expression that includes the one that needs
-      // to be wrapped.
-      List<Tree> flat = flatten(input, unit, path, enclosing, first);
-      // Zero-indexed start column
-      int startColumn = lineMap.getColumnNumber(getStartPosition(flat.get(0))) - 1;
-
-      // Handling leaving trailing non-string tokens at the end of the literal,
-      // e.g. the trailing `);` in `foo("...");`.
-      int end = getEndPosition(unit, getLast(flat));
-      int lineEnd = end;
-      while (Newlines.hasNewlineAt(input, lineEnd) == -1) {
-        lineEnd++;
+    }
+
+    private void wrapLongStrings(
+        TreeRangeMap<Integer, String> replacements, List<TreePath> longStringLiterals) {
+      for (TreePath path : longStringLiterals) {
+        // Find the outermost contiguous enclosing concatenation expression
+        TreePath enclosing = path;
+        while (enclosing.getParentPath().getLeaf().getKind() == Kind.PLUS) {
+          enclosing = enclosing.getParentPath();
+        }
+        // Is the literal being wrapped the first in a chain of concatenation expressions?
+        // i.e. `ONE + TWO + THREE`
+        // We need this information to handle continuation indents.
+        AtomicBoolean first = new AtomicBoolean(false);
+        // Finds the set of string literals in the concat expression that includes the one that
+        // needs
+        // to be wrapped.
+        List<Tree> flat = flatten(input, unit, path, enclosing, first);
+        // Zero-indexed start column
+        int startColumn = lineMap.getColumnNumber(getStartPosition(flat.get(0))) - 1;
+
+        // Handling leaving trailing non-string tokens at the end of the literal,
+        // e.g. the trailing `);` in `foo("...");`.
+        int end = getEndPosition(unit, getLast(flat));
+        int lineEnd = end;
+        while (Newlines.hasNewlineAt(input, lineEnd) == -1) {
+          lineEnd++;
+        }
+        int trailing = lineEnd - end;
+
+        // Get the original source text of the string literals, excluding `"` and `+`.
+        ImmutableList<String> components = stringComponents(input, unit, flat);
+        replacements.put(
+            Range.closedOpen(getStartPosition(flat.get(0)), getEndPosition(unit, getLast(flat))),
+            reflow(separator, columnLimit, startColumn, trailing, components, first.get()));
       }
-      int trailing = lineEnd - end;
+    }
+  }
+
+  private static final Method STRIP_INDENT = getStripIndent();
+
+  private static @Nullable Method getStripIndent() {
+    if (Runtime.version().feature() < 15) {
+      return null;
+    }
+    try {
+      return String.class.getMethod("stripIndent");
+    } catch (NoSuchMethodException e) {
+      throw new LinkageError(e.getMessage(), e);
+    }
+  }
 
-      // Get the original source text of the string literals, excluding `"` and `+`.
-      ImmutableList<String> components = stringComponents(input, unit, flat);
-      replacements.put(
-          Range.closedOpen(getStartPosition(flat.get(0)), getEndPosition(unit, getLast(flat))),
-          reflow(separator, columnLimit, startColumn, trailing, components, first.get()));
+  private static String stripIndent(String input) {
+    if (STRIP_INDENT == null) {
+      return input;
+    }
+    try {
+      return (String) STRIP_INDENT.invoke(input);
+    } catch (ReflectiveOperationException e) {
+      throw new LinkageError(e.getMessage(), e);
     }
-    return replacements;
   }
 
   /**
@@ -364,13 +473,16 @@ private static int getStartPosition(Tree tree) {
     return ((JCTree) tree).getStartPosition();
   }
 
-  /** Returns true if any lines in the given Java source exceed the column limit. */
-  private static boolean longLines(int columnLimit, String input) {
+  /**
+   * Returns true if any lines in the given Java source exceed the column limit, or contain a {@code
+   * """} that could indicate a text block.
+   */
+  private static boolean needWrapping(int columnLimit, String input) {
     // TODO(cushon): consider adding Newlines.lineIterable?
     Iterator<String> it = Newlines.lineIterator(input);
     while (it.hasNext()) {
       String line = it.next();
-      if (line.length() > columnLimit) {
+      if (line.length() > columnLimit || line.contains("\"\"\"")) {
         return true;
       }
     }
@@ -385,7 +497,6 @@ private static JCTree.JCCompilationUnit parse(String source, boolean allowString
     context.put(DiagnosticListener.class, diagnostics);
     Options.instance(context).put("--enable-preview", "true");
     Options.instance(context).put("allowStringFolding", Boolean.toString(allowStringFolding));
-    JCTree.JCCompilationUnit unit;
     JavacFileManager fileManager = new JavacFileManager(context, true, UTF_8);
     try {
       fileManager.setLocation(StandardLocation.PLATFORM_CLASS_PATH, ImmutableList.of());
@@ -404,7 +515,7 @@ public CharSequence getCharContent(boolean ignoreEncodingErrors) {
     JavacParser parser =
         parserFactory.newParser(
             source, /* keepDocComments= */ true, /* keepEndPos= */ true, /* keepLineMap= */ true);
-    unit = parser.parseCompilationUnit();
+    JCTree.JCCompilationUnit unit = parser.parseCompilationUnit();
     unit.sourcefile = sjfo;
     Iterable<Diagnostic<? extends JavaFileObject>> errorDiagnostics =
         Iterables.filter(diagnostics.getDiagnostics(), Formatter::errorDiagnostic);

core/src/test/java/com/google/googlejavaformat/java/StringWrapperTest.java

@@ -60,19 +60,90 @@ public void textBlock() throws Exception {
         lines(
             "package com.mypackage;",
             "public class ReproBug {",
-            "    private String myString;",
-            "    private ReproBug() {",
-            "        String str =",
-            "                \"\"\"",
-            "               "
-                + " {\"sourceEndpoint\":\"ri.something.1-1.object-internal.1\",\"targetEndpoint\":\"ri.some"
-                + "thing.1-1.object-internal.2\",\"typeId\":\"typeId\"}\"\"\";",
-            "        myString = str;",
-            "    }",
+            "  private String myString;",
+            "  private ReproBug() {",
+            "    String str =",
+            "        \"\"\"",
+            "{\"sourceEndpoint\":\"ri.something.1-1.object-internal.1\",\"targetEndpoint"
+                + "\":\"ri.something.1-1.object-internal.2\",\"typeId\":\"typeId\"}\"\"\";",
+            "    myString = str;",
+            "  }",
             "}");
     assertThat(StringWrapper.wrap(100, input, new Formatter())).isEqualTo(input);
   }
 
+  // Test that whitespace handling on text block lines only removes spaces, not other control
+  // characters.
+  @Test
+  public void textBlockControlCharacter() throws Exception {
+    assumeTrue(Runtime.version().feature() >= 15);
+    // We want an actual control character in the Java source being formatted, not a unicode escape,
+    // i.e. the escape below doesn't need to be double-escaped.
+    String input =
+        lines(
+            "package p;",
+            "public class T {",
+            "  String s =",
+            "      \"\"\"",
+            "      \u0007lorem",
+            "      \u0007",
+            "      ipsum",
+            "      \"\"\";",
+            "}");
+    String actual = StringWrapper.wrap(100, input, new Formatter());
+    assertThat(actual).isEqualTo(input);
+  }
+
+  @Test
+  public void textBlockTrailingWhitespace() throws Exception {
+    assumeTrue(Runtime.version().feature() >= 15);
+    String input =
+        lines(
+            "public class T {",
+            "  String s =",
+            "      \"\"\"",
+            "      lorem   ",
+            "      ipsum",
+            "      \"\"\";",
+            "}");
+    String expected =
+        lines(
+            "public class T {",
+            "  String s =",
+            "      \"\"\"",
+            "      lorem",
+            "      ipsum",
+            "      \"\"\";",
+            "}");
+    String actual = StringWrapper.wrap(100, input, new Formatter());
+    assertThat(actual).isEqualTo(expected);
+  }
+
+  @Test
+  public void textBlockSpaceTabMix() throws Exception {
+    assumeTrue(Runtime.version().feature() >= 15);
+    String input =
+        lines(
+            "public class T {",
+            "  String s =",
+            "      \"\"\"",
+            "      lorem",
+            "     \tipsum",
+            "      \"\"\";",
+            "}");
+    String expected =
+        lines(
+            "public class T {",
+            "  String s =",
+            "      \"\"\"",
+            "      lorem",
+            "      ipsum",
+            "      \"\"\";",
+            "}");
+    String actual = StringWrapper.wrap(100, input, new Formatter());
+    assertThat(actual).isEqualTo(expected);
+  }
+
   private static String lines(String... line) {
     return Joiner.on('\n').join(line) + '\n';
   }