Better lambda/assignment inlining logic (#169)

Reduce situations where an inlined lambda can get a newline after the arrow, and allow inlining new Class() calls in more contexts.

This is a two pronged change:
* **When part of a simple inlining**: prefer partially inlining (instead of putting entire lambda body expression on the 2nd line) only if this results in fewer lines used, *or* it's a long method chain, like a builder.
* **When part of a complex inlining**: partially inline iff it fits and the expression being inlined is complex, otherwise abort the inlining chain

-- simple / complex refer to the complexity of the levels partially inlined so far in the given inlining chain, where if any of the levels is complex, the inlining is considered complex.

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

@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: Reduce situations where an inlined lambda can get a newline after the
+    arrow, and allow inlining new Class() calls in more contexts.
+  links:
+  - https://github.com/palantir/palantir-java-format/pull/169

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

@@ -254,18 +254,43 @@ private Optional<State> handle_breakOnlyIfInnerLevelsThenFitOnOneLine(
             State brokenState,
             boolean keepIndent,
             Obs.ExplorationNode explorationNode) {
-        List<Level> innerLevels = this.docs.stream()
-                .filter(doc -> doc instanceof Level)
-                .map(doc -> ((Level) doc))
-                .collect(Collectors.toList());
 
-        boolean anyLevelWasBroken = innerLevels.stream().anyMatch(level -> !brokenState.isOneLine(level));
+        boolean anyLevelWasBroken = brokenState.numLines() != state.numLines() + 1;
 
         if (!anyLevelWasBroken) {
             return Optional.of(brokenState);
         }
+
+        Optional<State> partiallyInlinedStateOpt =
+                tryInlinePrefixOntoCurrentLine(commentsHelper, maxWidth, state, keepIndent, explorationNode);
+        if (!partiallyInlinedStateOpt.isPresent()) {
+            return Optional.empty();
+        }
+        State partiallyInlinedState = partiallyInlinedStateOpt.get();
+
+        boolean bodyIsComplex = this.docs.stream()
+                .filter(doc -> doc instanceof Level)
+                .map(doc -> ((Level) doc))
+                .anyMatch(il -> il.openOp.complexity() == Complexity.COMPLEX);
+
+        if (bodyIsComplex || partiallyInlinedState.numLines() < brokenState.numLines()) {
+            return Optional.of(partiallyInlinedState);
+        }
+        return Optional.empty();
+    }
+
+    private Optional<State> tryInlinePrefixOntoCurrentLine(
+            CommentsHelper commentsHelper,
+            int maxWidth,
+            State state,
+            boolean keepIndent,
+            Obs.ExplorationNode explorationNode) {
         // Find the last level, skipping empty levels (that contain nothing, or are made up
         // entirely of other empty levels).
+        List<Level> innerLevels = this.docs.stream()
+                .filter(doc -> doc instanceof Level)
+                .map(doc -> ((Level) doc))
+                .collect(Collectors.toList());
 
         // Last level because there might be other in-between levels after the initial break like `new
         // int[]
@@ -451,14 +476,11 @@ private static Optional<State> tryBreakInnerLevel_checkInner(
                                             commentsHelper, maxWidth, state1, exp, isSimpleInlining))
                             .map(Exploration::markAccepted);
                 })
-                .inlineSuffix(() -> {
-                    State state1 = state.withIndentIncrementedBy(innerLevel.getPlusIndent());
-                    return explorationNode
-                            .newChildNode(innerLevel, state1)
-                            .maybeExplore("recurse into inner tryInlineSuffix", state1, exp ->
-                                    innerLevel.tryInlineSuffix(commentsHelper, maxWidth, state1, exp, isSimpleInlining))
-                            .map(Exploration::markAccepted);
-                })
+                .inlineSuffix(() -> explorationNode
+                        .newChildNode(innerLevel, state)
+                        .maybeExplore("recurse into inner tryInlineSuffix", state, exp ->
+                                innerLevel.tryInlineSuffix(commentsHelper, maxWidth, state, exp, isSimpleInlining))
+                        .map(Exploration::markAccepted))
                 .breakOnlyIfInnerLevelsThenFitOnOneLine(keepIndentWhenInlined -> {
                     // This case currently only matches lambda _expressions_ (without curlies)
                     State state1 =
@@ -485,21 +507,30 @@ private static Optional<State> tryBreakInnerLevel_checkInner(
                                     case CHECK_INNER:
                                         return Optional.empty();
                                     case ACCEPT_INLINE_CHAIN:
+                                        Exploration broken =
+                                                innerLevel.breakNormally(state, levelNode, commentsHelper, maxWidth);
+                                        return innerLevel.handle_breakOnlyIfInnerLevelsThenFitOnOneLine(
+                                                commentsHelper,
+                                                maxWidth,
+                                                state1,
+                                                broken.state(),
+                                                keepIndentWhenInlined,
+                                                explorationNode);
                                     case ACCEPT_INLINE_CHAIN_IF_SIMPLE_OTHERWISE_CHECK_INNER:
-                                        // we only want to allow inlining in these cases
+                                        // We will allow inlining this kind of level but only if the level itself is
+                                        // not simple.
+                                        if (lastLevel2.openOp.complexity() == Complexity.SIMPLE) {
+                                            return Optional.empty();
+                                        }
+                                        return innerLevel.tryInlinePrefixOntoCurrentLine(
+                                                commentsHelper,
+                                                maxWidth,
+                                                state1,
+                                                keepIndentWhenInlined,
+                                                explorationNode);
+                                    default:
+                                        throw new RuntimeException("Unknown breakabilityIfLastLevel: " + lastLevel2);
                                 }
-
-                                // Note: intentionally using 'state' and not 'state1' as Level.breakNormally will always
-                                // add the level's plusIndent.
-                                Exploration broken =
-                                        innerLevel.breakNormally(state, levelNode, commentsHelper, maxWidth);
-                                return innerLevel.handle_breakOnlyIfInnerLevelsThenFitOnOneLine(
-                                        commentsHelper,
-                                        maxWidth,
-                                        state1,
-                                        broken.state(),
-                                        keepIndentWhenInlined,
-                                        explorationNode);
                             })
                             .map(Exploration::markAccepted);
                 })

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

@@ -42,6 +42,7 @@ public abstract class State {
     public abstract boolean mustBreak();
     /** Counts how many lines a particular formatting took. */
     public abstract int numLines();
+
     /**
      * Counts how many times reached a branch, where multiple formattings would be considered. Expected runtime is
      * exponential in this number.

palantir-java-format/src/main/java/com/palantir/javaformat/java/JavaInputAstVisitor.java

@@ -683,12 +683,15 @@ public Void visitTypeCast(TypeCastTree node, Void unused) {
     @Override
     public Void visitNewClass(NewClassTree node, Void unused) {
         sync(node);
-        builder.open(
-                ZERO,
-                BreakBehaviours.preferBreakingLastInnerLevel(true),
-                node.getClassBody() != null
-                        ? LastLevelBreakability.ACCEPT_INLINE_CHAIN
-                        : LastLevelBreakability.CHECK_INNER);
+        LastLevelBreakability breakabilityIfLastLevel = node.getClassBody() != null
+                ? LastLevelBreakability.ACCEPT_INLINE_CHAIN
+                : LastLevelBreakability.ACCEPT_INLINE_CHAIN_IF_SIMPLE_OTHERWISE_CHECK_INNER;
+        builder.open(OpenOp.builder()
+                .debugName("visitNewClass")
+                .plusIndent(ZERO)
+                .breakBehaviour(BreakBehaviours.preferBreakingLastInnerLevel(true))
+                .breakabilityIfLastLevel(breakabilityIfLastLevel)
+                .build());
         if (node.getEnclosingExpression() != null) {
             scan(node.getEnclosingExpression(), null);
             builder.breakOp();
@@ -2625,8 +2628,10 @@ void visitDot(ExpressionTree node0) {
                         .debugName("visitDot")
                         .plusIndent(plusFour)
                         .breakBehaviour(BreakBehaviours.preferBreakingLastInnerLevel(true))
-                        .breakabilityIfLastLevel(LastLevelBreakability.ACCEPT_INLINE_CHAIN)
+                        .breakabilityIfLastLevel(
+                                LastLevelBreakability.ACCEPT_INLINE_CHAIN_IF_SIMPLE_OTHERWISE_CHECK_INNER)
                         .columnLimitBeforeLastBreak(METHOD_CHAIN_COLUMN_LIMIT)
+                        .isSimple(false)
                         .build());
                 scan(getArrayBase(node), null);
                 builder.breakOp();

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

@@ -1,7 +1,9 @@
 class Test {
     public void testPrimitiveFields() {
         ConverterVerifier.verifyFieldsAreAnnotatedWithToPojo(
-                Contact.class, DomainContact.class, new ImmutableList.Builder<String>()
+                Contact.class,
+                DomainContact.class,
+                new ImmutableList.Builder<String>()
                         .add(
                                 "address1",
                                 "address2",

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

@@ -1,7 +1,7 @@
 class B32114928 {
     {
-        Class<T> tClass = (Class<T>)
-                verifyNotNull((ParameterizedType) getClass().getGenericSuperclass())
+        Class<T> tClass =
+                (Class<T>) verifyNotNull((ParameterizedType) getClass().getGenericSuperclass())
                         .getActualTypeArguments()[0];
     }
 }

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

@@ -1,9 +1,8 @@
 class Palantir1 {
     void foo() {
         for (Double d : Lists.newArrayList(5.0, 10.0, 1.0, 0.0)) {
-            list.add(
-                    new BinaryNumericConstantColumnInfoV1(
-                            operation, SparkIntegrationTestRuleSets.column("x"), d, Optional.empty()));
+            list.add(new BinaryNumericConstantColumnInfoV1(
+                    operation, SparkIntegrationTestRuleSets.column("x"), d, Optional.empty()));
         }
     }
 }

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

@@ -2,17 +2,14 @@ class Palantir2 {
     void foo() {
         operations.put(
                 FunctionDefinition.LOREM_IPSUM_DOLOR_,
-                (func, helper) -> new Column(
-                        new GetArrayItem(
-                                helper.accessColumnFuncti(func.arguments().get(0))
-                                        .expr(),
-                                helper.accessColumnFuncti(
-                                                // Subtract 1 from the provided index to move from 1-indexed to
-                                                // 0-indexed.
-                                                FooBarClassName1.of(
-                                                        func.arguments().get(1),
-                                                        FooBarClassName1.BinaryOperation.SUBTRACT,
-                                                        IntegerLiteral.of(1)))
-                                        .expr())));
+                (func, helper) -> new Column(new GetArrayItem(
+                        helper.accessColumnFuncti(func.arguments().get(0)).expr(),
+                        helper.accessColumnFuncti(
+                                        // Subtract 1 from the provided index to move from 1-indexed to 0-indexed.
+                                        FooBarClassName1.of(
+                                                func.arguments().get(1),
+                                                FooBarClassName1.BinaryOperation.SUBTRACT,
+                                                IntegerLiteral.of(1)))
+                                .expr())));
     }
 }

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

@@ -1,20 +1,17 @@
 class Palantir3 {
     void foo() {
         operations.put(FunctionDefinition.ARRAY_REGEX_REPLACE, (func, helper) -> {
-            return new Column(
-                    new ArrayTransform(
-                            helper.expressionAsColumn(func.arguments().get(0)).expr(),
-                            new LambdaFunction(
-                                    new RegExpReplace(
-                                            var,
-                                            helper.expressionAsColumn(
-                                                            func.arguments().get(1))
-                                                    .expr(),
-                                            helper.expressionAsColumn(
-                                                            func.arguments().get(2))
-                                                    .expr()),
-                                    JavaConversions.asScalaBuffer(Arrays.asList(var)),
-                                    false)));
+            return new Column(new ArrayTransform(
+                    helper.expressionAsColumn(func.arguments().get(0)).expr(),
+                    new LambdaFunction(
+                            new RegExpReplace(
+                                    var,
+                                    helper.expressionAsColumn(func.arguments().get(1))
+                                            .expr(),
+                                    helper.expressionAsColumn(func.arguments().get(2))
+                                            .expr()),
+                            JavaConversions.asScalaBuffer(Arrays.asList(var)),
+                            false)));
         });
     }
 }

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

@@ -1,9 +1,9 @@
 class Palantir7 {
     void foo() {
         when(graphService.getPathFromRootToNode(any(GraphId.class), any(NodeId.class)))
-                .thenAnswer(invocation -> invocation.getArguments().length == 2
-                                && invocation.getArguments()[1] != null
-                        ? ImmutableList.of(invocation.getArgument(1))
-                        : ImmutableList.of(NodeId.create()));
+                .thenAnswer(invocation ->
+                        invocation.getArguments().length == 2 && invocation.getArguments()[1] != null
+                                ? ImmutableList.of(invocation.getArgument(1))
+                                : ImmutableList.of(NodeId.create()));
     }
 }

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

@@ -0,0 +1,9 @@
+class Palantir8 {
+    void foo() {
+        assertThatServiceExceptionThrownBy(
+                new Aggregation.Builder()
+                        .max(max)
+                        .aggregations(ImmutableMap.of("bountiful", bountifulAggregation))::build)
+                .hasType(MY_ERROR_TYPE);
+    }
+}

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

@@ -0,0 +1,8 @@
+class Palantir8 {
+    void foo() {
+        assertThatServiceExceptionThrownBy(new Aggregation.Builder()
+                        .max(max)
+                        .aggregations(ImmutableMap.of("bountiful", bountifulAggregation))::build)
+                .hasType(MY_ERROR_TYPE);
+    }
+}

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

@@ -1,6 +1,7 @@
 public class Foo {
     public static void main(String[] args) {
-        foo.bar().baz(argargargargargargargargargargargargargargargargargargargargarg, param ->
-                doStuff().method1().method2(arljhfaldjfhalfjhalfjahflahflahfaljfhadlfjhljsahljhfsaljfhlfajhfljh));
+        foo.bar().baz(argargargargargargargargargargargargargargargargargargargargarg, param -> doStuff()
+                .method1()
+                .method2(arljhfaldjfhalfjhalfjahflahflahfaljfhadlfjhljsahljhfsaljfhlfajhfljh));
     }
 }

palantir-java-format/src/test/resources/com/palantir/javaformat/java/testdata/palantir-lambda-inlining-prefers-break.input

@@ -0,0 +1,11 @@
+class PalantirLambdaInliningPrefersBreak {
+    void foo() {
+        endpointDefinition.getDeprecated().ifPresent(documentation -> endpointBuilder.addMethod(
+                MethodSpec.methodBuilder("deprecated")
+                        .addModifiers(Modifier.PUBLIC)
+                        .addAnnotation(Override.class)
+                        .returns(ParameterizedTypeName.get(ClassName.get(Optional.class), ClassName.get(String.class)))
+                        .addStatement("return $1T.of($2S)", Optional.class, documentation)
+                        .build()));
+    }
+}

palantir-java-format/src/test/resources/com/palantir/javaformat/java/testdata/palantir-lambda-inlining-prefers-break.output

@@ -0,0 +1,12 @@
+class PalantirLambdaInliningPrefersBreak {
+    void foo() {
+        endpointDefinition
+                .getDeprecated()
+                .ifPresent(documentation -> endpointBuilder.addMethod(MethodSpec.methodBuilder("deprecated")
+                        .addModifiers(Modifier.PUBLIC)
+                        .addAnnotation(Override.class)
+                        .returns(ParameterizedTypeName.get(ClassName.get(Optional.class), ClassName.get(String.class)))
+                        .addStatement("return $1T.of($2S)", Optional.class, documentation)
+                        .build()));
+    }
+}