Fixed bug in DeprecatedLambdaSyntaxFix

idea/src/org/jetbrains/kotlin/idea/quickfix/DeprecatedLambdaSyntaxFix.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.idea.util.psiModificationUtil.getFunctionLiteralArgu
 import org.jetbrains.kotlin.idea.util.psiModificationUtil.moveInsideParenthesesAndReplaceWith
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.utils.sure
 import java.util.ArrayList
@@ -152,12 +153,12 @@ private class LambdaToFunctionExpression(
         return null
     }
 
-    private fun JetElement.replaceWithReturn(psiFactory: JetPsiFactory) {
+    private fun JetExpression.replaceWithReturn(psiFactory: JetPsiFactory) {
         if (this is JetReturnExpression) {
             return
         }
         else {
-            replace(psiFactory.createReturn(getText()))
+            replace(psiFactory.createReturn(this))
         }
     }
 
@@ -183,8 +184,7 @@ private class LambdaToFunctionExpression(
 
         val functionWithEmptyBody = psiFactory.createFunction(functionDeclaration + " {}")
 
-        val blockExpression = functionLiteral.getBodyExpression()
-        if (blockExpression == null) return functionWithEmptyBody
+        val blockExpression = functionLiteral.getBodyExpression() ?: return functionWithEmptyBody
 
         val statements = blockExpression.getStatements()
         if (statements.isEmpty()) return functionWithEmptyBody
@@ -194,9 +194,16 @@ private class LambdaToFunctionExpression(
         }
 
         // many statements
-        if (returnType != null) statements.last().replaceWithReturn(psiFactory)
+        if (returnType != null) statements.filterIsInstance<JetExpression>().lastOrNull()?.replaceWithReturn(psiFactory)
 
-        return psiFactory.createFunction(functionDeclaration + "{ " + blockExpression.getText() + "}")
+        val fromElement = functionLiteral.getArrowNode()?.getPsi() ?: functionLiteral.getLBrace()
+        val toElement = functionLiteral.getRBrace()
+        // to include comments in the start/end of the body
+        val bodyText = fromElement.siblings(withItself = false)
+                .takeWhile { it != toElement }
+                .map { it.getText() }
+                .joinToString("")
+        return psiFactory.createFunction(functionDeclaration + "{ " + bodyText + "}")
     }
 
 }

idea/testData/quickfix/migration/lambdaSyntax/lambdaSyntaxMultiple.after.kt

@@ -1,6 +1,7 @@
 // "Migrate lambda syntax in whole project" "true"
 
 val a = fun (): Int {
+
     val b = fun (): Int = 5
 
     return b()

idea/testData/quickfix/migration/lambdaSyntax/lastStatementIsComment.kt

@@ -0,0 +1,8 @@
+// "Migrate lambda syntax" "true"
+
+
+val a = { <caret>(p: Int): String ->
+    val v = p + 1
+    v.toString()
+    // returns v.toString()
+}

idea/testData/quickfix/migration/lambdaSyntax/lastStatementIsComment.kt.after

@@ -0,0 +1,8 @@
+// "Migrate lambda syntax" "true"
+
+
+val a = fun (p: Int): String {
+    val v = p + 1
+    return v.toString()
+    // returns v.toString()
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -3075,6 +3075,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("lastStatementIsComment.kt")
+            public void testLastStatementIsComment() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/migration/lambdaSyntax/lastStatementIsComment.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("paranthesizedParameters.kt")
             public void testParanthesizedParameters() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/migration/lambdaSyntax/paranthesizedParameters.kt");