Stop deeply nested calls from taking exponential time to format

Since palantir#1535 (2.89.0), the levels above the
innermost eight first try to break only their last inner level and fall
back to breaking normally. That attempt can lay out all of the inner
level before it fails, and in deeply nested code it nearly always fails
once the closing parentheses no longer fit in the line, so every level
above repeats it for each layout it tries. Nine nested
Map.ofEntries(Map.entry(...)) took 69.5 s to format, and the file from
palantir#1632 took 34 s with the jar and 60 s with
the native binary (#35).

A failed attempt leaves nothing behind but the failure, so a level now
remembers the states it failed from and answers at once when asked
again. The key holds what the attempt reads: its arguments and the
state's column, indents, pending break and taken break tags.

The new test formats the nested calls under a 10 s limit, which the old
code exceeds. They now take 0.45 s, and that file 1.05 s with the jar
and 0.73 s with the native binary. The layout does not change: the
output for both inputs and for the 15,747 files of the JDK 21 sources is
byte for byte the same, and formatting those sources takes no longer.

Author: abashev

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

@@ -41,9 +41,11 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.OptionalInt;
+import java.util.Set;
 import java.util.stream.Collector;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -67,6 +69,13 @@ public final class Level extends Doc {
     @SuppressWarnings("Immutable") // Effectively immutable
     private final ImmutableSupplier<Integer> memoizedMaxDepth = Suppliers.memoize(() -> computeMaxDepth(docs))::get;
 
+    /**
+     * 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
+     * exponential in their depth to format.
+     */
+    private final Set<LastLevelAttempt> failedLastLevelAttempts = new HashSet<>();
+
     /** The immutable characteristics of this level determined before the level contents are available. */
     private final OpenOp openOp;
 
@@ -364,9 +373,21 @@ private Optional<State> tryBreakLastLevel(
         }
         Level innerLevel = ((Level) getLast(docs));
 
-        return tryBreakInnerLevel(commentsHelper, maxWidth, state, explorationNode, innerLevel, isSimpleInliningSoFar);
+        LastLevelAttempt attempt = new LastLevelAttempt(maxWidth, isSimpleInliningSoFar, state.layoutInputs());
+        if (failedLastLevelAttempts.contains(attempt)) {
+            return Optional.empty();
+        }
+        Optional<State> result =
+                tryBreakInnerLevel(commentsHelper, maxWidth, state, explorationNode, innerLevel, isSimpleInliningSoFar);
+        if (result.isEmpty()) {
+            failedLastLevelAttempts.add(attempt);
+        }
+        return result;
     }
 
+    /** A call of {@link #tryBreakLastLevel}, by everything besides this level that decides its outcome. */
+    private record LastLevelAttempt(int maxWidth, boolean isSimpleInliningSoFar, State.LayoutInputs inputs) {}
+
     @SuppressWarnings("for-rollout:NullAway")
     private Optional<State> tryInlineSuffix(
             CommentsHelper commentsHelper,

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

@@ -26,6 +26,7 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
+import java.util.Objects;
 import org.immutables.value.Value;
 import org.immutables.value.Value.Parameter;
 
@@ -221,6 +222,27 @@ State withTokState(Comment comment, TokState tokState) {
                 .build();
     }
 
+    /**
+     * The part of this state that laying out a level from here reads. It leaves out the line count, which is only ever
+     * compared with another count from the same starting point; the branching coefficient, which nothing reads; and the
+     * states of breaks, levels and comments, since a layout only reads those of the docs it has laid out itself.
+     */
+    LayoutInputs layoutInputs() {
+        return new LayoutInputs(lastIndent(), indent(), column(), mustBreak(), breakTagsTaken());
+    }
+
+    /** See {@link #layoutInputs()}. */
+    record LayoutInputs(int lastIndent, int indent, int column, boolean mustBreak, Set<BreakTag> breakTagsTaken) {
+        /**
+         * Leaves out the taken break tags: hashing them on every lookup costs more than comparing them on the lookups
+         * that match everything else. Equality still compares them.
+         */
+        @Override
+        public int hashCode() {
+            return Objects.hash(lastIndent, indent, column, mustBreak);
+        }
+    }
+
     public static class Builder extends ImmutableState.Builder {}
 
     public static Builder builder() {

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

@@ -19,6 +19,7 @@
 import static com.palantir.javaformat.java.JavaFormatterOptions.Style;
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
 
 import com.google.common.base.Joiner;
 import com.google.common.collect.Range;
@@ -32,6 +33,7 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.List;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -473,4 +475,20 @@ void producesFormattingChangesOnAlreadyFormattedFiles() throws FormatterExceptio
                         formattedClass, List.of(Range.closedOpen(0, formattedClass.length()))))
                 .isNotEmpty();
     }
+
+    @Test
+    void formatsDeeplyNestedCallsQuickly() {
+        // Nine nested Map.ofEntries(Map.entry(...)) are one more than fit in 120 columns with each argument kept on its
+        // call's line. Every such attempt then failed, but only after laying out everything inside it, and each level
+        // above repeated it for every layout it tried: formatting took time exponential in the depth.
+        String call = "Map.ofEntries(Map.entry(\"key0\", \"value0\"), Map.entry(\"key1\", \"value1\"))";
+        for (int depth = 0; depth < 9; depth++) {
+            call = "Map.ofEntries(Map.entry(\"key" + depth + "\", " + call + "))";
+        }
+        String input = "class DeepNesting {\n    Object foo() {\n        return " + call + ";\n    }\n}\n";
+        Formatter formatter = Formatter.createFormatter(
+                JavaFormatterOptions.builder().style(Style.OJF).build());
+
+        assertTimeoutPreemptively(Duration.ofSeconds(10), () -> formatter.formatSource(input));
+    }
 }