Indent text blocks in the main formatting pass

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

@@ -75,7 +75,6 @@
 import com.google.googlejavaformat.FormattingError;
 import com.google.googlejavaformat.Indent;
 import com.google.googlejavaformat.Input;
-import com.google.googlejavaformat.Newlines;
 import com.google.googlejavaformat.Op;
 import com.google.googlejavaformat.OpenOp;
 import com.google.googlejavaformat.OpsBuilder.BlankLineWanted;
@@ -523,7 +522,7 @@
         if (afterFirstToken) {
           builder.forcedBreak();
         }
         builder.open(row.iterator().next().getKind() == NEW_ARRAY || cols == 1 ? ZERO : plusFour);
         boolean firstInRow = true;
         for (ExpressionTree item : row) {
           if (!firstInRow) {
@@ -763,7 +762,7 @@
         CollapseEmptyOrNot.YES,
         AllowLeadingBlankLine.YES,
         AllowTrailingBlankLine.YES);
     if (node.getStatement().getKind() == BLOCK) {
       builder.space();
     } else {
       builder.breakOp(" ");
@@ -1079,11 +1078,11 @@
     builder.open(plusFour);
     builder.open(
         node.getInitializer().size() > 1
                 && node.getInitializer().get(0).getKind() == Tree.Kind.EXPRESSION_STATEMENT
             ? plusFour
             : ZERO);
     if (!node.getInitializer().isEmpty()) {
       if (node.getInitializer().get(0).getKind() == VARIABLE) {
         PeekingIterator<StatementTree> it =
             Iterators.peekingIterator(node.getInitializer().iterator());
         visitVariables(
@@ -1147,7 +1146,7 @@
     while (true) {
       expressions.add(node.getCondition());
       statements.add(node.getThenStatement());
       if (node.getElseStatement() != null && node.getElseStatement().getKind() == IF) {
         node = (IfTree) node.getElseStatement();
       } else {
         break;
@@ -1756,21 +1755,8 @@
   public Void visitLiteral(LiteralTree node, Void unused) {
     sync(node);
     String sourceForNode = getSourceForNode(node, getCurrentPath());
-    if (sourceForNode.startsWith("\"\"\"")) {
-      String separator = Newlines.guessLineSeparator(sourceForNode);
-      ImmutableList<String> initialLines = sourceForNode.lines().collect(toImmutableList());
-      String stripped = initialLines.stream().skip(1).collect(joining(separator)).stripIndent();
-      // Use the last line of the text block to determine if it is deindented to column 0, by
-      // comparing the length of the line in the input source with the length after processing
-      // the text block contents with stripIndent().
-      boolean deindent =
-          getLast(initialLines).stripTrailing().length()
-              == Streams.findLast(stripped.lines()).orElseThrow().stripTrailing().length();
-      if (deindent) {
-        Indent indent = Indent.Const.make(Integer.MIN_VALUE / indentMultiplier, indentMultiplier);
-        builder.breakOp(indent);
-      }
-      token(sourceForNode);
+    if (sourceForNode.startsWith(TEXT_BLOCK_DELIMITER)) {
+      visitTextBlock(sourceForNode);
       return null;
     }
     if (isUnaryMinusLiteral(sourceForNode)) {
@@ -1781,6 +1767,39 @@
     return null;
   }
 
+  private static final String TEXT_BLOCK_DELIMITER = "\"\"\"";
+
+  private void visitTextBlock(String sourceForNode) {
+    // 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 = sourceForNode.lines().collect(toImmutableList());
+    String stripped = initialLines.stream().skip(1).collect(joining("\n")).stripIndent();
+    ImmutableList<String> strippedLines = stripped.lines().collect(toImmutableList());
+    String lastLine = getLast(strippedLines);
+    // Use the last line of the text block to determine if it is deindented to column 0, by
+    // comparing the length of the line in the input source with the length after processing
+    // the text block contents with stripIndent().
+    boolean deindent = getLast(initialLines).stripTrailing().length() == lastLine.length();
+    if (deindent) {
+      Indent indent = Indent.Const.make(Integer.MIN_VALUE / indentMultiplier, indentMultiplier);
+      builder.breakOp(indent);
+    }
+    ImmutableList.Builder<String> lines = ImmutableList.builder();
+    lines.add(initialLines.get(0).stripTrailing());
+    lines.addAll(strippedLines.subList(0, strippedLines.size() - 1));
+    String withoutDelimiter =
+        lastLine.substring(0, lastLine.length() - TEXT_BLOCK_DELIMITER.length()).stripTrailing();
+    if (!withoutDelimiter.isEmpty()) {
+      lines.add(withoutDelimiter + '\\');
+    }
+    // If the trailing line is just """, indenting it 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'.
+    lines.add(TEXT_BLOCK_DELIMITER);
+    builder.multilineToken(sourceForNode, lines.build());
+  }
+
   // A negative numeric literal -n is usually represented as unary minus on n,
   // but that doesn't work for integer or long MIN_VALUE. The parser works
   // around that by representing it directly as a signed literal (with no
@@ -2987,7 +3006,7 @@
     builder.close();
   }
 
   /** Collapse chains of {@code .} operators, across multiple {@link ASTNode} types. */
 
   /// Output a "." node.
   ///

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

@@ -14,11 +14,9 @@
 
 package com.google.googlejavaformat.java;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Iterables.getLast;
 import static com.google.googlejavaformat.java.Trees.getEndPosition;
 import static com.google.googlejavaformat.java.Trees.getStartPosition;
-import static java.lang.Math.min;
 import static java.util.stream.Collectors.joining;
 
 import com.google.common.base.CharMatcher;
@@ -130,34 +128,26 @@ private int visualColumn(int position) {
     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);
+      new LongStringsScanner(longStringLiterals).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 class LongStringsScanner extends TreePathScanner<Void, Void> {
 
       private final List<TreePath> longStringLiterals;
-      private final List<Tree> textBlocks;
 
-      LongStringsAndTextBlockScanner(List<TreePath> longStringLiterals, List<Tree> textBlocks) {
+      LongStringsScanner(List<TreePath> longStringLiterals) {
         this.longStringLiterals = longStringLiterals;
-        this.textBlocks = textBlocks;
       }
 
       @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(TEXT_BLOCK_DELIMITER)) {
-          textBlocks.add(literalTree);
+        if (input.startsWith(TEXT_BLOCK_DELIMITER, getStartPosition(literalTree))) {
           return null;
         }
         Tree parent = getCurrentPath().getParentPath().getLeaf();
@@ -178,56 +168,6 @@ public Void visitLiteral(LiteralTree literalTree, Void aVoid) {
       }
     }
 
-    private void indentTextBlocks(
-        TreeRangeMap<Integer, String> replacements, List<Tree> textBlocks) {
-      for (Tree tree : textBlocks) {
-        int startPosition = lineMap.getStartPosition(lineMap.getLineNumber(getStartPosition(tree)));
-        int endPosition = getEndPosition(tree, unit);
-        String text = input.substring(startPosition, endPosition);
-        int leadingWhitespace = CharMatcher.whitespace().negate().indexIn(text);
-
-        // 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 = initialLines.stream().skip(1).collect(joining(separator)).stripIndent();
-        ImmutableList<String> lines = stripped.lines().collect(toImmutableList());
-        boolean deindent =
-            getLast(initialLines).stripTrailing().length()
-                == getLast(lines).stripTrailing().length();
-
-        String prefix =
-            deindent ? "" : options.indentString(visualColumn(startPosition + leadingWhitespace));
-
-        StringBuilder output = new StringBuilder(prefix).append(initialLines.get(0).stripLeading());
-        for (int i = 0; i < lines.size(); i++) {
-          String line = lines.get(i);
-          String trimmed = line.stripTrailing();
-          output.append(separator);
-          if (!trimmed.isEmpty()) {
-            // Don't add incidental leading whitespace to empty lines
-            output.append(prefix);
-          }
-          if (i == lines.size() - 1) {
-            String withoutDelimiter =
-                trimmed
-                    .substring(0, trimmed.length() - TEXT_BLOCK_DELIMITER.length())
-                    .stripTrailing();
-            if (!withoutDelimiter.stripLeading().isEmpty()) {
-              output.append(withoutDelimiter).append('\\').append(separator).append(prefix);
-            }
-            // If the trailing line is just """, indenting it 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(TEXT_BLOCK_DELIMITER);
-          } else {
-            output.append(line);
-          }
-        }
-        replacements.put(Range.closedOpen(startPosition, endPosition), output.toString());
-      }
-    }
-
     private void wrapLongStrings(
         TreeRangeMap<Integer, String> replacements, List<TreePath> longStringLiterals) {
       for (TreePath path : longStringLiterals) {
@@ -450,14 +390,13 @@ private static boolean noComments(
   private static final CharMatcher STRING_CONCAT_DELIMITER =
       CharMatcher.whitespace().or(CharMatcher.anyOf("\"+"));
 
-  /// Returns true if any lines in the given Java source exceed the column limit, or contain a `"""`
-  /// that could indicate a text block.
+  /// Returns true if any lines in the given Java source exceed the column limit.
   private static boolean needWrapping(int columnLimit, String input, JavaFormatterOptions options) {
     // TODO(cushon): consider adding Newlines.lineIterable?
     Iterator<String> it = Newlines.lineIterator(input);
     while (it.hasNext()) {
       String line = it.next();
-      if (options.visualLength(line) > columnLimit || line.contains(TEXT_BLOCK_DELIMITER)) {
+      if (options.visualLength(line) > columnLimit) {
         return true;
       }
     }

core/src/main/java/com/google/googlejavaformat/layout/Doc.java

@@ -56,7 +56,8 @@
 ///
 /// A [Choice] offers a preferred alternative layout for a group, used when the group doesn't fit
 /// flat but the preferred layout's first line does.
-public sealed interface Doc permits Doc.Text, Doc.Tok, Doc.Space, Doc.Break, Doc.Group, Doc.Choice {
+public sealed interface Doc
+    permits Doc.Text, Doc.Lines, Doc.Tok, Doc.Space, Doc.Break, Doc.Group, Doc.Choice {
 
   /// Width used for documents that cannot be laid out on a single line.
   int MAX_LINE_WIDTH = 1000;
@@ -117,7 +118,7 @@ public Text(String text) {
 
     @Override
     public int width() {
-      return Newlines.firstBreak(text) >= 0 ? MAX_LINE_WIDTH : text.length();
+      return text.length();
     }
 
     @Override
@@ -126,6 +127,20 @@ public void appendFlat(StringBuilder sb) {
     }
   }
 
+  /// A multi-line input token whose non-empty lines after the first are indented to the column
+  /// where the token starts, for example text blocks.
+  record Lines(ImmutableList<String> lines, Range<Integer> range) implements Doc {
+    @Override
+    public int width() {
+      return MAX_LINE_WIDTH;
+    }
+
+    @Override
+    public void appendFlat(StringBuilder sb) {
+      throw new UnsupportedOperationException();
+    }
+  }
+
   /// Input text that isn't a token (the equivalent of the `Doc.Tok` leaf in
   /// [com.google.googlejavaformat.Doc]). This is usually a comment. It can also be whitespace or
   /// newlines around an *imaginary* token ([RealOrImaginary#IMAGINARY]): a token the AST visitor

core/src/main/java/com/google/googlejavaformat/layout/LayoutBuilder.java

@@ -119,7 +119,8 @@ public void drain() {
                 token,
                 RealOrImaginary.IMAGINARY,
                 ZERO,
-                /* breakAndIndentTrailingComment= */ Optional.empty()));
+                /* breakAndIndentTrailingComment= */ Optional.empty(),
+                /* linesIndentedToTokenStart= */ Optional.empty()));
       }
     }
     this.inputPosition = endPosition;
@@ -257,14 +258,29 @@ public void token(
       RealOrImaginary realOrImaginary,
       Indent plusIndentCommentsBefore,
       Optional<Indent> breakAndIndentTrailingComment) {
+    token(
+        token,
+        realOrImaginary,
+        plusIndentCommentsBefore,
+        breakAndIndentTrailingComment,
+        Optional.empty());
+  }
+
+  private void token(
+      String token,
+      RealOrImaginary realOrImaginary,
+      Indent plusIndentCommentsBefore,
+      Optional<Indent> breakAndIndentTrailingComment,
+      Optional<ImmutableList<String>> linesIndentedToTokenStart) {
     ImmutableList<? extends Input.Token> tokens = input.getTokens();
     if (peekToken().filter(token::equals).isPresent()) {
       events.add(
           new TokenOp(
               tokens.get(tokenI++),
               RealOrImaginary.REAL,
               plusIndentCommentsBefore,
-              breakAndIndentTrailingComment));
+              breakAndIndentTrailingComment,
+              linesIndentedToTokenStart));
     } else if (realOrImaginary.isReal()) {
       // Generated a "bad" token, which doesn't exist on the input. Complain, unless (for example)
       // we're guessing at an optional token, in which case it's dropped.
@@ -275,6 +291,12 @@ public void token(
     }
   }
 
+  /// Emits a real multi-line token whose non-empty lines after the first are indented to the column
+  /// where the token starts, for example a text block.
+  public void multilineToken(String token, ImmutableList<String> lines) {
+    token(token, RealOrImaginary.REAL, ZERO, Optional.empty(), Optional.of(lines));
+  }
+
   public void op(String op) {
     int opN = op.length();
     for (int i = 0; i < opN; i++) {
@@ -592,9 +614,11 @@ private static Doc buildTree(List<Event> events) {
         }
         case TokenOp tokenOp -> {
           Input.Tok tok = tokenOp.token().getTok();
+          Range<Integer> range = Range.singleton(tok.getIndex()).canonical(INTEGERS);
           appendLevel.add(
-              new Doc.Text(
-                  tok.getOriginalText(), Range.singleton(tok.getIndex()).canonical(INTEGERS)));
+              tokenOp.linesIndentedToTokenStart().isPresent()
+                  ? new Doc.Lines(tokenOp.linesIndentedToTokenStart().get(), range)
+                  : new Doc.Text(tok.getOriginalText(), range));
         }
         case TokOp tokOp -> appendLevel.add(new Doc.Tok(tokOp.tok()));
         case HugPointOp hugPointOp -> appendLevel.add(hugPointOp.hugPoint().marker);
@@ -751,7 +775,8 @@ private record TokenOp(
       Input.Token token,
       RealOrImaginary realOrImaginary,
       Indent plusIndentCommentsBefore,
-      Optional<Indent> breakAndIndentTrailingComment)
+      Optional<Indent> breakAndIndentTrailingComment,
+      Optional<ImmutableList<String>> linesIndentedToTokenStart)
       implements Event {}
 
   private record TokOp(Input.Tok tok) implements Event {}

core/src/main/java/com/google/googlejavaformat/layout/LayoutSolver.java

@@ -51,7 +51,7 @@ public static String format(Doc doc, int maxWidth) {
     for (Action action : solve(doc, maxWidth, /* commentsHelper= */ null)) {
       switch (action) {
         case Append a -> sb.append(a.text());
-        case Newline n -> sb.append('\n');
+        case Newline n -> sb.repeat("\n", n.count());
         case IndentTo i -> sb.repeat(" ", i.spaces());
       }
     }
@@ -63,7 +63,7 @@ public static void format(Doc doc, int maxWidth, Output output, CommentsHelper c
     for (Action action : solve(doc, maxWidth, commentsHelper)) {
       switch (action) {
         case Append a -> output.append(a.text(), a.range());
-        case Newline n -> output.append("\n", Doc.EMPTY_RANGE);
+        case Newline n -> output.append("\n".repeat(n.count()), Doc.EMPTY_RANGE);
         case IndentTo i -> output.indent(i.spaces());
       }
     }
@@ -73,7 +73,12 @@ private sealed interface Action {}
 
   private record Append(String text, Range<Integer> range) implements Action {}
 
-  private record Newline() implements Action {}
+  /// Emits `count` line breaks; `count` is more than 1 when a [Doc.Lines] contains blank lines.
+  private record Newline(int count) implements Action {
+    Newline() {
+      this(1);
+    }
+  }
 
   private record IndentTo(int spaces) implements Action {}
 
@@ -232,13 +237,35 @@ private State leaf(Doc doc, State state) {
         case Doc.Group g -> group(g, state);
         case Doc.Choice c -> choice(c, state);
         case Doc.Text t -> {
-          int newline = Newlines.firstBreak(t.text());
-          if (newline >= 0) {
-            endLine(state.column() + newline);
-          }
           actions.add(new Append(t.text(), t.range()));
           yield state.withColumn(state.column() + t.width());
         }
+        case Doc.Lines l -> {
+          ImmutableList<String> lines = l.lines();
+          endLine(state.column() + lines.get(0).length());
+          if (firstLineEnded) {
+            yield state;
+          }
+          int indent = lineIndent;
+          actions.add(new Append(lines.get(0), l.range()));
+          int newlines = 0;
+          for (String line : lines.subList(1, lines.size())) {
+            newlines++;
+            if (line.isEmpty()) {
+              continue;
+            }
+            actions.add(new Newline(newlines));
+            newlines = 0;
+            if (indent > 0) {
+              actions.add(new IndentTo(indent));
+            }
+            actions.add(new Append(line, l.range()));
+          }
+          if (newlines > 0) {
+            actions.add(new Newline(newlines));
+          }
+          yield state.withColumn(state.column() + l.width());
+        }
         case Doc.Space s -> {
           actions.add(new Append(" ", Doc.EMPTY_RANGE));
           yield state.withColumn(state.column() + 1);