Fix crash when breaking up long comment after simple lambda body (#203)

Make lambda/assignment logic more resilient so it doesn't crash when encountering long comments.

changelog/@unreleased/pr-203.v2.yml

@@ -0,0 +1,6 @@
+type: fix
+fix:
+  description: Make lambda/assignment logic more resilient so it doesn't crash when
+    encountering long comments.
+  links:
+  - https://github.com/palantir/palantir-java-format/pull/203

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

@@ -260,7 +260,12 @@ private Optional<State> handle_breakOnlyIfInnerLevelsThenFitOnOneLine(
             boolean keepIndent,
             Obs.ExplorationNode explorationNode) {
 
-        boolean anyLevelWasBroken = brokenState.numLines() != state.numLines() + 1;
+        // Note: we are not checking if the brokenState produced one extra line compared to state, as this can be
+        // misleading if there is no level but a single comment that got reflowed onto multiple lines (see palantir-11).
+        boolean anyLevelWasBroken = docs.stream()
+                .filter(doc -> doc instanceof Level)
+                .map(doc -> ((Level) doc))
+                .anyMatch(level -> !brokenState.isOneLine(level));
 
         if (!anyLevelWasBroken) {
             return Optional.of(brokenState);

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

@@ -18,7 +18,6 @@
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
-import com.google.common.base.Preconditions;
 import com.google.errorprone.annotations.Immutable;
 import com.palantir.javaformat.Indent;
 import fj.data.Set;
@@ -100,9 +99,10 @@ boolean isOneLine(Level level) {
     }
 
     String getTokText(Comment comment) {
-        return Preconditions.checkNotNull(
-                        tokStates().get(comment).toNull(), "Expected Tok state to exist for: %s", comment)
-                .text();
+        // A TokState will only be present if computeBreaks was called.
+        // That won't always happen, for example when the level containing this comment was one-lined.
+        // Note: if the parent level was inlined, this method itself also won't get called, unless we're in debug mode.
+        return tokStates().get(comment).map(TokState::text).orSome(comment::getFlat);
     }
 
     /** Record whether break was taken. */

palantir-java-format/src/test/resources/com/palantir/javaformat/java/testdata/palantir-11.input

@@ -0,0 +1,13 @@
+public class Palantir11 {
+    private void foo() {
+        boolean answer = strategy.accept(Strategies.visitor(
+                greaterThan -> true, // we don't need to validate greaterThan because it'll roll up to a good version
+                exact -> !coordinates.contains(MavenCoordinate.of(productId, exact.getVersion())),
+                remove -> true,
+                stay -> true,
+                stayWithExceptions ->
+                        !coordinates.contains(MavenCoordinate.of(
+                                productId,
+                                stayWithExceptions.getDeploymentsExceptionVersion().getVersion()))));
+    }
+}

palantir-java-format/src/test/resources/com/palantir/javaformat/java/testdata/palantir-11.output

@@ -0,0 +1,12 @@
+public class Palantir11 {
+    private void foo() {
+        boolean answer = strategy.accept(Strategies.visitor(
+                greaterThan -> true, // we don't need to validate greaterThan because it'll roll up to a good version
+                exact -> !coordinates.contains(MavenCoordinate.of(productId, exact.getVersion())),
+                remove -> true,
+                stay -> true,
+                stayWithExceptions -> !coordinates.contains(MavenCoordinate.of(
+                        productId,
+                        stayWithExceptions.getDeploymentsExceptionVersion().getVersion()))));
+    }
+}