From e49b201472fb2bfb768d505bbf40a6b53987e183 Mon Sep 17 00:00:00 2001
From: Jordan Kiesel <jordantkiesel@gmail.com>
Date: Wed, 6 Sep 2023 00:20:43 -0600
Subject: [PATCH] fix(printer): properly break and indent lambda with comments

closes #581
---
 .../src/utils/printSingleLambdaInvocation.ts  | 21 ++++++-----
 .../test/unit-test/lambda/_input.java         | 20 +++++++++++
 .../test/unit-test/lambda/_output.java        | 35 +++++++++++++++++--
 3 files changed, 62 insertions(+), 14 deletions(-)

diff --git a/packages/prettier-plugin-java/src/utils/printSingleLambdaInvocation.ts b/packages/prettier-plugin-java/src/utils/printSingleLambdaInvocation.ts
index 614c1a7b..3c64a6ca 100644
--- a/packages/prettier-plugin-java/src/utils/printSingleLambdaInvocation.ts
+++ b/packages/prettier-plugin-java/src/utils/printSingleLambdaInvocation.ts
@@ -2,7 +2,7 @@ import { ArgumentListCstNode, IToken } from "java-parser";
 import { builders } from "prettier/doc";
 import { isSingleArgumentLambdaExpressionWithBlock } from "./expressions-utils.js";
 import { printTokenWithComments } from "../printers/comments/format-comments.js";
-import { concat, dedent, indent } from "../printers/prettier-builder.js";
+import { concat, dedent } from "../printers/prettier-builder.js";
 import { putIntoBraces } from "../printers/printer-utils.js";
 
 const { softline, ifBreak } = builders;
@@ -18,14 +18,13 @@ export default function printSingleLambdaInvocation(
     isInsideMethodInvocationSuffix: true
   });
 
-  const formattedRBrace = isSingleArgumentLambdaExpressionWithBlock(
-    argumentListCtx
-  )
-    ? ifBreak(
-        indent(concat([softline, rBrace])),
-        printTokenWithComments(rBrace),
-        { groupId: lambdaParametersGroupId }
-      )
-    : indent(concat([softline, rBrace]));
-  return dedent(putIntoBraces(argumentList, "", lBrace, formattedRBrace));
+  if (isSingleArgumentLambdaExpressionWithBlock(argumentListCtx)) {
+    const formattedRBrace = ifBreak(
+      concat([softline, rBrace]),
+      printTokenWithComments(rBrace),
+      { groupId: lambdaParametersGroupId }
+    );
+    return putIntoBraces(dedent(argumentList), "", lBrace, formattedRBrace);
+  }
+  return putIntoBraces(argumentList, softline, lBrace, rBrace);
 }
diff --git a/packages/prettier-plugin-java/test/unit-test/lambda/_input.java b/packages/prettier-plugin-java/test/unit-test/lambda/_input.java
index 1de5d55a..b619d1f8 100644
--- a/packages/prettier-plugin-java/test/unit-test/lambda/_input.java
+++ b/packages/prettier-plugin-java/test/unit-test/lambda/_input.java
@@ -151,6 +151,26 @@ public void testConstructor() {
           }
       );
     }
+
+    void lambdaWithLeadingComments() {
+        System.out.println(
+            List.of(1, 2, 3).stream().map(
+                // a very long comment which explains the beatifullness of multiplication by 2
+                // yes this is very important
+                v -> v * 2
+            ).collect(Collectors.summingInt(v -> v))
+        );
+    }
+
+    void lambdaWithTrailingComments() {
+        System.out.println(
+            List.of(1, 2, 3).stream().map(
+                v -> v * 2
+                // a very long comment which explains the beatifullness of multiplication by 2
+                // yes this is very important
+            ).collect(Collectors.summingInt(v -> v))
+        );
+    }
 }
 
 class T {
diff --git a/packages/prettier-plugin-java/test/unit-test/lambda/_output.java b/packages/prettier-plugin-java/test/unit-test/lambda/_output.java
index aab01499..acddd589 100644
--- a/packages/prettier-plugin-java/test/unit-test/lambda/_output.java
+++ b/packages/prettier-plugin-java/test/unit-test/lambda/_output.java
@@ -46,9 +46,10 @@ public void onlyOneMethodInBody() {
   }
 
   public void lambdaWithoutBracesWhichBreak() {
-    call(x ->
-      foo.isVeryVeryVeryLongConditionTrue() &&
-      foo.isAnotherVeryVeryLongConditionTrue()
+    call(
+      x ->
+        foo.isVeryVeryVeryLongConditionTrue() &&
+        foo.isAnotherVeryVeryLongConditionTrue()
     );
   }
 
@@ -143,6 +144,34 @@ public void testConstructor() {
       }
     );
   }
+
+  void lambdaWithLeadingComments() {
+    System.out.println(
+      List
+        .of(1, 2, 3)
+        .stream()
+        .map(
+          // a very long comment which explains the beatifullness of multiplication by 2
+          // yes this is very important
+          v -> v * 2
+        )
+        .collect(Collectors.summingInt(v -> v))
+    );
+  }
+
+  void lambdaWithTrailingComments() {
+    System.out.println(
+      List
+        .of(1, 2, 3)
+        .stream()
+        .map(
+          v -> v * 2
+          // a very long comment which explains the beatifullness of multiplication by 2
+          // yes this is very important
+        )
+        .collect(Collectors.summingInt(v -> v))
+    );
+  }
 }
 
 class T {