Fix comments in lambdas

src/main/java/spoon/support/compiler/jdt/JDTCommentBuilder.java

@@ -18,6 +18,7 @@
 import spoon.reflect.code.CtComment;
 import spoon.reflect.code.CtConditional;
 import spoon.reflect.code.CtIf;
+import spoon.reflect.code.CtLambda;
 import spoon.reflect.code.CtNewArray;
 import spoon.reflect.code.CtStatement;
 import spoon.reflect.code.CtStatementList;
@@ -397,6 +398,20 @@ public void visitCtAnonymousExecutable(CtAnonymousExecutable e) {
 				e.addComment(comment);
 			}
 
+			@Override
+			public <T> void visitCtLambda(CtLambda<T> e) {
+				if (e.getExpression() != null) {
+					CtParameter<?> lastParameter = e.getParameters().get(e.getParameters().size() - 1);
+					if (comment.getPosition().getSourceStart() > lastParameter.getPosition().getSourceEnd()) {
+						e.getExpression().addComment(comment);
+					} else {
+						e.addComment(comment);
+					}
+				} else if (e.getBody() != null) {
+					e.addComment(comment);
+				}
+			}
+
 			@Override
 			public <T> void visitCtNewArray(CtNewArray<T> e) {
 				addCommentToNear(comment, new ArrayList<>(e.getElements()));

src/main/java/spoon/support/compiler/jdt/PositionBuilder.java

@@ -222,6 +222,12 @@ SourcePosition buildPositionCtElement(CtElement e, ASTNode node) {
 				}
 			}
 
+			// Handle lambda parameters without explicit type
+			if (variableDeclaration instanceof Argument && variableDeclaration.type == null) {
+				declarationSourceStart = findPrevNonWhitespace(contents, 0, declarationSourceStart - 1);
+				declarationSourceEnd = findNextNonWhitespace(contents, contents.length - 1, declarationSourceEnd + 1);
+			}
+
 			if (modifiersSourceStart <= 0) {
 				modifiersSourceStart = findNextNonWhitespace(contents, contents.length - 1, declarationSourceStart);
 			}

src/test/java/spoon/test/comment/CommentTest.java

@@ -34,6 +34,7 @@
 import spoon.reflect.code.CtInvocation;
 import spoon.reflect.code.CtJavaDoc;
 import spoon.reflect.code.CtJavaDocTag;
+import spoon.reflect.code.CtLambda;
 import spoon.reflect.code.CtLiteral;
 import spoon.reflect.code.CtLocalVariable;
 import spoon.reflect.code.CtNewArray;
@@ -1116,7 +1117,7 @@ public void testCommentGetRawContent() {
 				" */", type.getComments().get(0).getRawContent());
 	}
 
-  @Test
+	@Test
 	public void testEmptyStatementComments() {
 		//contract: model building should not produce NPE, comments should exist
 		Launcher launcher = new Launcher();
@@ -1128,4 +1129,46 @@ public void testEmptyStatementComments() {
 		assertEquals("comment", conditions.get(0).getComments().get(0).getContent());
 		assertEquals("comment", conditions.get(1).getComments().get(0).getContent());
 	}
+
+	@Test
+	public void testLambdaComments() {
+		//contract: comments in lambdas should be properly added to the AST
+		Launcher launcher = new Launcher();
+		launcher.addInputResource("./src/test/java/spoon/test/comment/testclasses/LambdaComments.java");
+		launcher.getEnvironment().setCommentEnabled(true);
+		CtModel model = launcher.buildModel();
+
+		List<CtLambda> lambdas = model.getElements(new TypeFilter<>(CtLambda.class));
+		assertEquals("comment", lambdas.get(0).getExpression().getComments().get(0).getContent());
+		assertEquals("comment", lambdas.get(1).getExpression().getComments().get(0).getContent());
+		assertEquals("comment", lambdas.get(2).getExpression().getComments().get(0).getContent());
+		assertEquals("comment", lambdas.get(3).getParent().getComments().get(0).getContent());
+		assertEquals("", lambdas.get(4).getExpression().getComments().get(0).getContent());
+
+		assertEquals("comment", lambdas.get(5).getBody().getStatement(0).getComments().get(0).getContent());
+		assertEquals("comment", lambdas.get(6).getComments().get(0).getContent());
+		assertEquals("comment", lambdas.get(7).getBody().getStatement(0).getComments().get(0).getContent());
+
+		assertTrue(lambdas.get(8).getParent().getComments().isEmpty());
+		assertTrue(lambdas.get(8).getComments().isEmpty());
+		assertTrue(lambdas.get(8).getExpression().getComments().isEmpty());
+		assertEquals(1, ((CtParameter) lambdas.get(8).getParameters().get(0)).getComments().size());
+		assertEquals(1, ((CtParameter) lambdas.get(8).getParameters().get(1)).getComments().size());
+		assertEquals("param1", ((CtParameter) lambdas.get(8).getParameters().get(0)).getComments().get(0).getContent());
+		assertEquals("param2", ((CtParameter) lambdas.get(8).getParameters().get(1)).getComments().get(0).getContent());
+
+		assertEquals("param1", ((CtParameter) lambdas.get(9).getParameters().get(0)).getComments().get(0).getContent());
+		assertEquals("param2", ((CtParameter) lambdas.get(9).getParameters().get(1)).getComments().get(0).getContent());
+
+		assertEquals("param1", ((CtParameter) lambdas.get(10).getParameters().get(0)).getComments().get(0).getContent());
+		assertEquals("param2", ((CtParameter) lambdas.get(10).getParameters().get(1)).getComments().get(0).getContent());
+
+		assertTrue(((CtParameter) lambdas.get(11).getParameters().get(0)).getComments().isEmpty());
+		assertEquals("param2", ((CtParameter) lambdas.get(11).getParameters().get(1)).getComments().get(0).getContent());
+
+		assertEquals("param1", ((CtParameter) lambdas.get(12).getParameters().get(0)).getComments().get(0).getContent());
+		assertEquals("param1", ((CtParameter) lambdas.get(12).getParameters().get(0)).getComments().get(1).getContent());
+		assertEquals("param2", ((CtParameter) lambdas.get(12).getParameters().get(1)).getComments().get(0).getContent());
+		assertEquals("param2", ((CtParameter) lambdas.get(12).getParameters().get(1)).getComments().get(1).getContent());
+	}
 }

src/test/java/spoon/test/comment/testclasses/LambdaComments.java

@@ -0,0 +1,51 @@
+package spoon.test.comment.testclasses;
+
+import java.util.function.BiFunction;
+
+public class LambdaComments {
+
+    void m1() {
+        BiFunction<Integer, Integer, Integer> lambda1 = (a, b) -> /* comment */
+            a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda2 = (a, b) -> // comment
+            a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda3 = (a, b) ->
+            // comment
+            a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda4 = /* comment */ (a, b) ->
+            a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda5 = (a, b) -> //
+                a + b;
+    }
+
+    void m2() {
+        BiFunction<Integer, Integer, Integer> lambda6 = (a, b) -> { /* comment */
+            return a + b;
+        };
+
+        BiFunction<Integer, Integer, Integer> lambda7 = (a, b) -> /* comment */ {
+            return a + b;
+        };
+
+        BiFunction<Integer, Integer, Integer> lambda8 = (a, b) -> {
+            /* comment */
+            return a + b;
+        };
+    }
+
+    void m3() {
+        BiFunction<Integer, Integer, Integer> lambda9 = (/* param1 */ a, /* param2 */ b) -> a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda10 = (/*param1*/ Integer a, /* param2 */ Integer b) -> a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda11 = (a /* param1 */, b /* param2 */) -> a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda12 = (a, b /* param2 */) -> a + b;
+
+        BiFunction<Integer, Integer, Integer> lambda13 = (/* param1 */ a /* param1 */, /* param2 */ b /* param2 */) -> a + b;
+    }
+}