Keep a wrapped trailing line comment where the first run put it

open-java-format/src/main/java/com/palantir/javaformat/OpsBuilder.java

@@ -551,7 +551,7 @@ public OpsOutput build() {
                                             tokBefore.isSlashSlashComment() ? FillMode.FORCED : FillMode.UNIFIED,
                                             "",
                                             tokenOp.getPlusIndentCommentsBefore()));
-                            tokOps.putAll(j, makeComment(tokBefore));
+                            tokOps.putAll(j, makeComment(tokBefore, false));
                             space = tokBefore.isSlashStarComment();
                             newlines = 0;
                             lastWasComment = true;
@@ -621,7 +621,7 @@ public OpsOutput build() {
                                 tokOps.put(tokAfterPos, Comment.makeNonNlsMarkers(tokAfter, literals));
                                 tokOps.put(tokAfterPos, Break.makeForced());
                             } else {
-                                tokOps.putAll(tokAfterPos, makeComment(tokAfter));
+                                tokOps.putAll(tokAfterPos, makeComment(tokAfter, true));
                             }
                             if (breakAfter) {
                                 tokOps.put(tokAfterPos, Break.make(FillMode.FORCED, "", ZERO));
@@ -731,10 +731,9 @@ private static boolean isForcedBreak(Op op) {
         return op instanceof Break b && b.isForced();
     }
 
-    private static List<Op> makeComment(Input.Tok comment) {
-        return comment.isSlashStarComment()
-                ? ImmutableList.of(Comment.make(comment))
-                : ImmutableList.of(Comment.make(comment), Break.makeForced());
+    private static List<Op> makeComment(Input.Tok comment, boolean trailing) {
+        var doc = trailing ? Comment.makeTrailing(comment) : Comment.make(comment);
+        return comment.isSlashStarComment() ? List.of(doc) : List.of(doc, Break.makeForced());
     }
 
     @Override

open-java-format/src/main/java/com/palantir/javaformat/doc/Comment.java

@@ -35,9 +35,17 @@ public final class Comment extends Doc implements Op {
     /** For a comment of NON-NLS markers, the index of the string literal each marker marks, or -1; else empty. */
     private final List<Integer> nonNlsLiterals;
 
-    private Comment(Input.Tok tok, List<Integer> nonNlsLiterals) {
+    /** Whether the comment trails code on its line. */
+    private final boolean trailing;
+
+    /** Set by {@link #withFirstWordOnly}. */
+    private final boolean firstWordOnly;
+
+    private Comment(Input.Tok tok, List<Integer> nonNlsLiterals, boolean trailing, boolean firstWordOnly) {
         this.tok = tok;
         this.nonNlsLiterals = nonNlsLiterals;
+        this.trailing = trailing;
+        this.firstWordOnly = firstWordOnly;
     }
 
     /**
@@ -47,7 +55,12 @@ private Comment(Input.Tok tok, List<Integer> nonNlsLiterals) {
      * @return the new {@code Tok}
      */
     public static Comment make(Input.Tok tok) {
-        return new Comment(tok, List.of());
+        return new Comment(tok, List.of(), false, false);
+    }
+
+    /** A comment that trails code on its line. */
+    public static Comment makeTrailing(Input.Tok tok) {
+        return new Comment(tok, List.of(), true, false);
     }
 
     /**
@@ -58,7 +71,32 @@ public static Comment make(Input.Tok tok) {
      * @param literals for each marker, the index of the string literal it marks, or -1
      */
     public static Comment makeNonNlsMarkers(Input.Tok tok, List<Integer> literals) {
-        return new Comment(tok, List.copyOf(literals));
+        return new Comment(tok, List.copyOf(literals), false, false);
+    }
+
+    /**
+     * This comment, laid out by its slashes and its first word only, as a line comment that spans lines is. The next
+     * run reads a wrapped trailing comment back as one that spans lines, so laying it out that way from the start gives
+     * both runs the same layout, and the code before it is not broken for room that the comment does not get.
+     */
+    public Comment withFirstWordOnly() {
+        return new Comment(tok, nonNlsLiterals, trailing, true);
+    }
+
+    /**
+     * Whether this trailing line comment, counted in full by the layout, came out of {@code state} wrapped: it did not
+     * fit after the code before it, however that was broken.
+     */
+    public boolean wrappedIn(State state) {
+        return trailing
+                && tok.isSlashSlashComment()
+                && !countsFirstWordOnly()
+                && Newlines.containsBreaks(state.getTokText(this));
+    }
+
+    /** Whether the layout counts only the slashes and the first word of this line comment. */
+    boolean countsFirstWordOnly() {
+        return tok.isSlashSlashComment() && (firstWordOnly || Newlines.containsBreaks(tok.getOriginalText()));
     }
 
     @Override
@@ -79,6 +117,11 @@ protected float computeWidth() {
             }
             return width;
         }
+        if (countsFirstWordOnly()) {
+            // The rest of the comment goes where the wrap puts it, so the code before it is laid out as if it were not
+            // there; see withFirstWordOnly.
+            return JavaCommentsHelper.firstWordLength(getFlat());
+        }
         int idx = Newlines.firstBreak(tok.getOriginalText());
         // only count the first line of multi-line block comments
         if (tok.isComment()) {

open-java-format/src/main/java/com/palantir/javaformat/doc/Level.java

@@ -48,6 +48,7 @@
 import java.util.stream.Collector;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import javax.annotation.Nullable;
 
 /** A {@code Level} inside a {@link Doc}. */
 public final class Level extends Doc {
@@ -67,6 +68,10 @@ public final class Level extends Doc {
     @SuppressWarnings("Immutable") // Effectively immutable
     private final ImmutableSupplier<Integer> memoizedMaxDepth = Suppliers.memoize(() -> computeMaxDepth(docs))::get;
 
+    /** See {@link #hasCommentByFirstWord()}; computed when first asked for, once the docs are all in. */
+    @Nullable
+    private Boolean hasCommentByFirstWord;
+
     /**
      * Where {@link #tryBreakLastLevel} has already failed. An attempt can lay out all of the last inner level before it
      * fails, and each level above repeats it for every layout it tries, so without this deeply nested calls took time
@@ -124,15 +129,38 @@ protected Range<Integer> computeRange() {
 
     @Override
     public State computeBreaks(CommentsHelper commentsHelper, int maxWidth, State state, Obs.ExplorationNode observer) {
-        return tryToFitOnOneLine(maxWidth, state, docs)
-                .map(newWidth -> state.withColumn(newWidth).withLevelState(this, new State.LevelState(true)))
-                .orElseGet(() -> {
-                    Obs.LevelNode childLevel = observer.newChildNode(this, state);
-                    State newState =
-                            new BreakImpl(commentsHelper, maxWidth, state, childLevel).apply(getBreakBehaviour());
-
-                    return childLevel.finishLevel(state.updateAfterLevel(newState));
-                });
+        var oneLine = tryToFitOnOneLine(maxWidth, state, docs);
+        if (oneLine.isPresent() && !hasCommentByFirstWord()) {
+            return state.withColumn(oneLine.get()).withLevelState(this, new State.LevelState(true));
+        }
+
+        var childLevel = observer.newChildNode(this, state);
+        State newState;
+        if (oneLine.isPresent()) {
+            // A level that fits is written from its flat text, which would skip wrapping a comment that counted
+            // only its first word. So it is laid out with none of its breaks taken, and the comment is wrapped
+            // where it needs to.
+            newState = childLevel
+                    .explore(
+                            "one line, comments laid out",
+                            state,
+                            explorationNode -> tryToLayOutLevelOnOneLine(
+                                    commentsHelper, maxWidth, state, memoizedSplitsBreaks.get(), explorationNode))
+                    .markAccepted();
+        } else {
+            newState = new BreakImpl(commentsHelper, maxWidth, state, childLevel).apply(getBreakBehaviour());
+        }
+        return childLevel.finishLevel(state.updateAfterLevel(newState));
+    }
+
+    /** Whether this level or a level in it holds a comment that the layout counts by its first word only. */
+    private boolean hasCommentByFirstWord() {
+        if (hasCommentByFirstWord == null) {
+            hasCommentByFirstWord = docs.stream()
+                    .anyMatch(doc -> (doc instanceof Comment comment && comment.countsFirstWordOnly())
+                            || (doc instanceof Level level && level.hasCommentByFirstWord()));
+        }
+        return hasCommentByFirstWord;
     }
 
     /**

open-java-format/src/main/java/com/palantir/javaformat/java/Formatter.java

@@ -30,6 +30,7 @@
 import com.palantir.javaformat.OpsBuilder;
 import com.palantir.javaformat.OpsBuilder.OpsOutput;
 import com.palantir.javaformat.Utils;
+import com.palantir.javaformat.doc.Comment;
 import com.palantir.javaformat.doc.Doc;
 import com.palantir.javaformat.doc.DocBuilder;
 import com.palantir.javaformat.doc.Level;
@@ -48,6 +49,9 @@
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import javax.tools.Diagnostic;
 import javax.tools.DiagnosticCollector;
 import javax.tools.DiagnosticListener;
@@ -150,14 +154,27 @@ static JavaOutput format(
         opsBuilder.drain();
         OpsOutput opsOutput = opsBuilder.build();
 
-        Level doc = new DocBuilder().withOps(opsOutput.ops()).build();
-
-        // Don't even allocate all those JSON nodes if we're not going to write it out
-        Sink sink = debugMode ? new JsonSink() : new NoopSink();
-
-        Obs.ExplorationNode observationNode = Obs.createRoot(sink);
-        State finalState =
-                doc.computeBreaks(commentsHelper, options.maxLineLength(), State.startingState(), observationNode);
+        // A trailing line comment that came out wrapped did not fit however the code before it was broken. Laid
+        // out again by its first word only, as the next run lays out the wrapped comment it reads, it leaves that
+        // code as it would be without the comment, and the next run gives the same layout.
+        List<Op> ops = opsOutput.ops();
+        Level doc;
+        Sink sink;
+        State finalState;
+        while (true) {
+            doc = new DocBuilder().withOps(ops).build();
+
+            // Don't even allocate all those JSON nodes if we're not going to write it out
+            sink = debugMode ? new JsonSink() : new NoopSink();
+            finalState = doc.computeBreaks(
+                    commentsHelper, options.maxLineLength(), State.startingState(), Obs.createRoot(sink));
+
+            var wrapped = wrappedComments(ops, finalState);
+            if (wrapped.isEmpty()) {
+                break;
+            }
+            ops = withFirstWordOnly(ops, wrapped);
+        }
 
         JavaOutput javaOutput = new JavaOutput(javaInput, opsOutput.inputMetadata());
         doc.write(finalState, javaOutput);
@@ -169,6 +186,25 @@ static JavaOutput format(
         return javaOutput;
     }
 
+    /** The comments of {@code ops} that came out of {@code state} wrapped, see {@link Comment#wrappedIn}. */
+    private static Set<Comment> wrappedComments(List<Op> ops, State state) {
+        var wrapped = new HashSet<Comment>();
+        for (var op : ops) {
+            if (op instanceof Comment comment && comment.wrappedIn(state)) {
+                wrapped.add(comment);
+            }
+        }
+        return wrapped;
+    }
+
+    /** {@code ops} with each of the {@code wrapped} comments laid out by its first word only. */
+    private static List<Op> withFirstWordOnly(List<Op> ops, Set<Comment> wrapped) {
+        return ops.stream()
+                .map(op ->
+                        op instanceof Comment comment && wrapped.contains(comment) ? comment.withFirstWordOnly() : op)
+                .toList();
+    }
+
     static JCCompilationUnit parseJcCompilationUnit(Context context, String sourceText) throws FormatterException {
         DiagnosticCollector<JavaFileObject> diagnostics = new DiagnosticCollector<>();
         context.put(DiagnosticListener.class, diagnostics);

open-java-format/src/main/java/com/palantir/javaformat/java/JavaCommentsHelper.java

@@ -189,14 +189,28 @@ public static String addMissingSpace(String line) {
         return line;
     }
 
-    private static String lineCommentPrefix(String line) {
+    static String lineCommentPrefix(String line) {
         int prefixLength = 0;
         while (prefixLength < line.length() && line.charAt(prefixLength) == '/') {
             prefixLength++;
         }
         return "/".repeat(prefixLength);
     }
 
+    /**
+     * The length of the slashes and the first word of a line comment, which stay on its line however it is wrapped:
+     * {@link #wrapLineComments} breaks a line only at whitespace after them.
+     */
+    public static int firstWordLength(String line) {
+        var prefixLength = lineCommentPrefix(line).length();
+        for (var i = prefixLength + 1; i < line.length(); i++) {
+            if (CharMatcher.breakingWhitespace().matches(line.charAt(i))) {
+                return i;
+            }
+        }
+        return line.length();
+    }
+
     private List<String> wrapLineComments(List<String> lines, int column0) {
         List<String> result = new ArrayList<>();
         for (String line : lines) {

open-java-format/src/main/java/com/palantir/javaformat/java/JavaInput.java

@@ -445,7 +445,10 @@ public CharSequence getCharContent(boolean ignoreEncodingErrors) throws IOExcept
                 }
             }
             if (strings.size() == 1) {
-                toks.add(new Tok(isNumbered ? kN++ : -1, originalTokText, tokText, charI, columnI, isToken, t.kind()));
+                if (!tokText.startsWith("//") || !continueTrailingComment(toks, originalTokText, text)) {
+                    toks.add(new Tok(
+                            isNumbered ? kN++ : -1, originalTokText, tokText, charI, columnI, isToken, t.kind()));
+                }
                 charI += originalTokText.length();
                 columnI = updateColumn(columnI, originalTokText);
 
@@ -497,6 +500,56 @@ private static Collection<JCDiagnostic> getDiagnostics(DeferredDiagnosticHandler
         }
     }
 
+    /**
+     * Adds {@code comment}, a line comment that starts a line of its own, to the line comment that trails code on the
+     * line above when it starts in that comment's column, the way the formatter writes a trailing comment it wraps. The
+     * two become one tok, with the line break and the indentation between them, as a run of {@code ///} lines from JDK
+     * 23 on is. Read as a comment of its own, the continuation went to the indentation of the code on the next run.
+     *
+     * @param toks the toks so far, which end with the trailing comment, a newline and the indentation of
+     *     {@code comment} when it continues one
+     * @return whether {@code comment} continues the trailing comment and is now part of it
+     */
+    private static boolean continueTrailingComment(List<Tok> toks, String comment, String text) {
+        var n = toks.size();
+        if (n < 3) {
+            return false;
+        }
+
+        var trailing = toks.get(n - 3);
+        var newline = toks.get(n - 2);
+        var indentation = toks.get(n - 1);
+        if (!trailing.isSlashSlashComment() || !newline.isNewline() || !isSpaces(indentation)) {
+            return false;
+        }
+
+        // What comes before the trailing comment on its line is read in the text, since Tok.getColumn() counts an
+        // operator such as == once for each of its characters. Code there makes the comment a trailing one, and its
+        // length is the comment's column.
+        var position = trailing.getPosition();
+        var lineStart = Math.max(text.lastIndexOf('\n', position), text.lastIndexOf('\r', position)) + 1;
+        var beforeComment = text.substring(lineStart, position);
+        if (beforeComment.isBlank()
+                || beforeComment.length() != indentation.length()
+                || !JavaCommentsHelper.lineCommentPrefix(trailing.getOriginalText())
+                        .equals(JavaCommentsHelper.lineCommentPrefix(comment))
+                || NonNlsMarkers.isMarkersOnly(trailing.getOriginalText())
+                || NonNlsMarkers.isMarkersOnly(comment)) {
+            return false;
+        }
+
+        var continued =
+                trailing.getOriginalText() + newline.getOriginalText() + indentation.getOriginalText() + comment;
+        toks.subList(n - 3, n).clear();
+        toks.add(new Tok(
+                trailing.getIndex(), continued, continued, position, trailing.getColumn(), false, trailing.kind()));
+        return true;
+    }
+
+    private static boolean isSpaces(Tok tok) {
+        return !tok.isToken() && !tok.isNewline() && !tok.isComment();
+    }
+
     @SuppressWarnings("for-rollout:NullAway")
     private static int updateColumn(int columnI, String originalTokText) {
         Integer last = Iterators.getLast(Newlines.lineOffsetIterator(originalTokText));

open-java-format/src/test/java/com/palantir/javaformat/java/FormatterTest.java

@@ -212,26 +212,6 @@ void trailingCommentWithoutTerminalNewline() throws Exception {
         assertThat(FORMATTER.formatSource("/*\n * my comment */")).isEqualTo("/*\n * my comment */\n");
     }
 
-    @Test
-    void wrapLineComment() throws Exception {
-        // Not a case in test-cases yet: a second run moves the wrapped part of the comment to the indent of the block,
-        // and the cases check that a second run changes nothing.
-        assertThat(FORMATTER.formatSource("class T {\n"
-                        + "  public static void main(String[] args) { // one long incredibly"
-                        + " unbroken sentence moving from topic to topic so that no-one had a"
-                        + " chance to interrupt;\n"
-                        + "  }\n"
-                        + "}\n"))
-                .isEqualTo("class T {\n"
-                        + "    public static void main(\n"
-                        + "            String[]\n"
-                        + "                    args) { // one long incredibly unbroken sentence moving"
-                        + " from topic to topic so that no-one had a\n"
-                        + "                            // chance to interrupt;\n"
-                        + "    }\n"
-                        + "}\n");
-    }
-
     @Test
     void throwsFormatterException() throws Exception {
         assertThatThrownBy(() -> FORMATTER.formatSourceAndFixImports("package foo; public class {"))