Fix crash when formatting method references (#1630)

 Fix crash when formatting method references

Author: crogoz

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

@@ -67,15 +67,11 @@ public Float visitLevel(Level level) {
         if (found.isPresent()) {
             return visit(level.getDocs().get(found.getAsInt()));
         }
-        // Otherwise, assert that we encountered a break and move on.
-        if (StartsWithBreakVisitor.INSTANCE.visit(level) != Result.YES) {
-            // Avoid computing expensive level.representation() if we aren't throwing it in an exception.
-            throw new IllegalStateException(String.format(
-                    "Didn't find expected break at the beginning of level.\n%s",
-                    level.representation(State.startingState())));
-        }
-
-        return 0f;
+        // Level starts with tokens directly (e.g. method reference qualifier) — sum doc widths up to the first break.
+        return level.getDocs().stream()
+                .takeWhile(doc -> !(doc instanceof Break))
+                .map(Doc::getWidth)
+                .reduce(0f, Float::sum);
     }
 
     /**

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

@@ -80,5 +80,18 @@ class A {
         2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2
             + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2
             + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2;
+        return StreamEx.of(itemsToCreate)
+                .append(itemsToEdit)
+                .append(itemsToDelete)
+                .distinct()
+                .mapToEntry(item -> getArgsGenerator()
+                        .getNewArgs(
+                                extendedItems.getOrDefault(item, Set.of()),
+                                getOwningResource(
+                                        item,
+                                        getModificationContext()::getRestrictionStatusOrDefault,
+                                        RestrictionStatus::getPackageId,
+                                        getModificationContext())))
+                .toImmutableMap();
   }
 }

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

@@ -70,5 +70,18 @@ class A {
         something = 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2
                 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2
                 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2 + 2;
+        return StreamEx.of(itemsToCreate)
+                .append(itemsToEdit)
+                .append(itemsToDelete)
+                .distinct()
+                .mapToEntry(item -> getArgsGenerator()
+                        .getNewArgs(
+                                extendedItems.getOrDefault(item, Set.of()),
+                                getOwningResource(
+                                        item,
+                                        getModificationContext()::getRestrictionStatusOrDefault,
+                                        RestrictionStatus::getPackageId,
+                                        getModificationContext())))
+                .toImmutableMap();
     }
 }