anonymous -> lambda: restore comments (non-javadoc) outside method body (IDEA-157286)

java/java-analysis-impl/src/com/intellij/codeInspection/AnonymousCanBeLambdaInspection.java

@@ -254,6 +254,7 @@ public class AnonymousCanBeLambdaInspection extends BaseJavaBatchLocalInspection
       final PsiCodeBlock body = method.getBody();
       if (body == null) return null;
 
+      final Collection<PsiComment> comments = collectCommentsOutsideMethodBody(anonymousClass, body);
       final Project project = element.getProject();
       final PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project);
 
@@ -277,12 +278,12 @@ public class AnonymousCanBeLambdaInspection extends BaseJavaBatchLocalInspection
         .giveUniqueNames(project, elementFactory, lambdaExpression,
                          usedLocalNames, variables.toArray(new PsiVariable[variables.size()]));
 
-      final PsiExpression singleExpr = RedundantLambdaCodeBlockInspection.isCodeBlockRedundant(
-        lambdaExpression.getBody());
+      final PsiExpression singleExpr = RedundantLambdaCodeBlockInspection.isCodeBlockRedundant(lambdaExpression.getBody());
       if (singleExpr != null) {
         lambdaExpression.getBody().replace(singleExpr);
       }
       ChangeContextUtil.decodeContextInfo(lambdaExpression, null, null);
+      restoreComments(comments, lambdaExpression);
 
       final JavaCodeStyleManager javaCodeStyleManager = JavaCodeStyleManager.getInstance(project);
       if (forceIgnoreTypeCast) {
@@ -310,6 +311,16 @@ public class AnonymousCanBeLambdaInspection extends BaseJavaBatchLocalInspection
     return null;
   }
 
+  private static Collection<PsiComment> collectCommentsOutsideMethodBody(PsiAnonymousClass anonymousClass, PsiCodeBlock body) {
+    final Collection<PsiComment> psiComments = PsiTreeUtil.findChildrenOfType(anonymousClass, PsiComment.class);
+    for (Iterator<PsiComment> iterator = psiComments.iterator(); iterator.hasNext(); ) {
+      if (PsiTreeUtil.isAncestor(body, iterator.next(), false)) {
+        iterator.remove();
+      }
+    }
+    return ContainerUtil.map(psiComments, (comment) -> (PsiComment)comment.copy());
+  }
+
   private static void collectLocalVariablesDefinedInsideLambda(PsiLambdaExpression lambdaExpression,
                                                                final Set<PsiVariable> variables,
                                                                Set<String> namesOfVariablesInTheBlock) {
@@ -452,6 +463,16 @@ public class AnonymousCanBeLambdaInspection extends BaseJavaBatchLocalInspection
     return false;
   }
 
+  public static void restoreComments(Collection<PsiComment> comments, PsiElement lambda) {
+    PsiElement anchor = PsiTreeUtil.getParentOfType(lambda, PsiStatement.class, PsiField.class);
+    if (anchor == null) {
+      anchor = lambda;
+    }
+    for (PsiComment comment : comments) {
+      anchor.getParent().addBefore(comment, anchor);
+    }
+  }
+
   private static class ForbiddenRefsChecker extends JavaRecursiveElementWalkingVisitor {
     private boolean myBodyContainsForbiddenRefs;
 

java/java-analysis-impl/src/com/intellij/codeInspection/LambdaCanBeMethodReferenceInspection.java

@@ -518,13 +518,7 @@ public class LambdaCanBeMethodReferenceInspection extends BaseJavaBatchLocalInsp
           replace = replace.replace(cast);
         }
 
-        PsiElement anchor = PsiTreeUtil.getParentOfType(replace, PsiStatement.class);
-        if (anchor == null) {
-          anchor = replace;
-        }
-        for (PsiComment comment : comments) {
-          anchor.getParent().addBefore(comment, anchor);
-        }
+        AnonymousCanBeLambdaInspection.restoreComments(comments, replace);
         JavaCodeStyleManager.getInstance(project).shortenClassReferences(replace);
       }
     }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/anonymous2lambda/afterCommentsInsideAnonymousClassBody.java

@@ -0,0 +1,6 @@
+// "Replace with lambda" "true"
+
+class Test {
+    //my comment
+    Runnable r = () -> {};
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/anonymous2lambda/beforeCommentsInsideAnonymousClassBody.java

@@ -0,0 +1,8 @@
+// "Replace with lambda" "true"
+
+class Test {
+  Runnable r = new Ru<caret>nnable() {
+    //my comment
+    public void run () {}
+  };
+}