Only apply line break to lambdas with statements

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

@@ -0,0 +1,12 @@
+type: fix
+fix:
+  description: |-
+    Only apply line break to lambdas with statements
+
+    ## Before this PR
+    Previous PR #707 was a little aggressive and inserted line breaks for non-statement lambdas.
+
+    ## After this PR
+    Only multi-line lambdas cause a line-break in the method chain.
+  links:
+  - https://github.com/palantir/palantir-java-format/pull/708

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

@@ -107,6 +107,7 @@
 import com.sun.source.tree.IntersectionTypeTree;
 import com.sun.source.tree.LabeledStatementTree;
 import com.sun.source.tree.LambdaExpressionTree;
+import com.sun.source.tree.LambdaExpressionTree.BodyKind;
 import com.sun.source.tree.LiteralTree;
 import com.sun.source.tree.MemberReferenceTree;
 import com.sun.source.tree.MemberSelectTree;
@@ -2784,11 +2785,13 @@ private void visitRegularDot(List<ExpressionTree> items, boolean needDot) {
         int length = needDot0 ? minLength : 0;
         for (ExpressionTree e : items) {
             if (needDot) {
-                // Also break if invoked with a lambda -- palantir-break-lambda-arg
+                // Also break if invoked with a multi-statement lambda -- palantir-break-lambda-arg
                 // foo
-                //     .doSomething(() -> bar())
+                //     .doSomething(() -> {
+                //         bar();
+                //      })
                 //     .doSomethingElse();
-                if (length > minLength || methodHasLambdaArg(e)) {
+                if (length > minLength || methodHasMultiStatementLambdaArg(e)) {
                     builder.breakOp(Break.builder()
                             .fillMode(FillMode.UNIFIED)
                             .flat("")
@@ -2813,12 +2816,15 @@ private void visitRegularDot(List<ExpressionTree> items, boolean needDot) {
         }
     }
 
-    private boolean methodHasLambdaArg(ExpressionTree e) {
+    private boolean methodHasMultiStatementLambdaArg(ExpressionTree e) {
         if (!e.getKind().equals(METHOD_INVOCATION)) {
             return false;
         }
         return ((MethodInvocationTree) e)
-                .getArguments().stream().anyMatch(argExpr -> argExpr.getKind().equals(LAMBDA_EXPRESSION));
+                .getArguments().stream()
+                        .filter(argExpr -> argExpr.getKind().equals(LAMBDA_EXPRESSION))
+                        .map(argExpr -> (LambdaExpressionTree) argExpr)
+                        .anyMatch(lambda -> lambda.getBodyKind().equals(BodyKind.STATEMENT));
     }
 
     // avoid formattings like:

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

@@ -0,0 +1,7 @@
+class PalantirLambdaBreakChain {
+    void foo() {
+        return hello
+            .read(txn -> doSomeWork())
+            .doMoreWork();
+    }
+}

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

@@ -0,0 +1,5 @@
+class PalantirLambdaBreakChain {
+    void foo() {
+        return hello.read(txn -> doSomeWork()).doMoreWork();
+    }
+}