Replace return@forEach with continue in ConvertForEachToForLoopIntention

So #KT-17332 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertForEachToForLoopIntention.kt

@@ -20,7 +20,10 @@ import com.intellij.openapi.editor.Editor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.bindingContextUtil.getTargetFunction
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 
@@ -40,11 +43,11 @@ class ConvertForEachToForLoopIntention : SelfTargetingOffsetIndependentIntention
     }
 
     override fun applyTo(element: KtSimpleNameExpression, editor: Editor?) {
-        val (expressionToReplace, receiver, functionLiteral) = extractData(element)!!
+        val (expressionToReplace, receiver, functionLiteral, context) = extractData(element)!!
 
         val commentSaver = CommentSaver(expressionToReplace)
 
-        val loop = generateLoop(functionLiteral, receiver)
+        val loop = generateLoop(functionLiteral, receiver, context)
         val result = expressionToReplace.replace(loop)
 
         commentSaver.restore(result)
@@ -53,7 +56,8 @@ class ConvertForEachToForLoopIntention : SelfTargetingOffsetIndependentIntention
     private data class Data(
             val expressionToReplace: KtExpression,
             val receiver: KtExpression,
-            val functionLiteral: KtLambdaExpression
+            val functionLiteral: KtLambdaExpression,
+            val context: BindingContext
     )
 
     private fun extractData(nameExpr: KtSimpleNameExpression): Data? {
@@ -64,20 +68,31 @@ class ConvertForEachToForLoopIntention : SelfTargetingOffsetIndependentIntention
             else -> null
         } ?: return null) as KtExpression //TODO: submit bug
 
-        val resolvedCall = expression.getResolvedCall(expression.analyze()) ?: return null
+        val context = expression.analyze()
+        val resolvedCall = expression.getResolvedCall(context) ?: return null
         if (DescriptorUtils.getFqName(resolvedCall.resultingDescriptor).toString() !in FOR_EACH_FQ_NAMES) return null
 
         val receiver = resolvedCall.call.explicitReceiver as? ExpressionReceiver ?: return null
         val argument = resolvedCall.call.valueArguments.singleOrNull() ?: return null
         val functionLiteral = argument.getArgumentExpression() as? KtLambdaExpression ?: return null
-        return Data(expression, receiver.expression, functionLiteral)
+        return Data(expression, receiver.expression, functionLiteral, context)
     }
 
-    private fun generateLoop(functionLiteral: KtLambdaExpression, receiver: KtExpression): KtExpression {
+    private fun generateLoop(functionLiteral: KtLambdaExpression, receiver: KtExpression, context: BindingContext): KtExpression {
         val factory = KtPsiFactory(functionLiteral)
-        val loopRange = KtPsiUtil.safeDeparenthesize(receiver)
+
         val body = functionLiteral.bodyExpression!!
+        val function = functionLiteral.functionLiteral
+
+        body.forEachDescendantOfType<KtReturnExpression> {
+            if (it.getTargetFunction(context) == function) {
+                it.replace(factory.createExpression("continue"))
+            }
+        }
+
+        val loopRange = KtPsiUtil.safeDeparenthesize(receiver)
         val parameter = functionLiteral.valueParameters.singleOrNull()
+
         return factory.createExpressionByPattern("for($0 in $1){ $2 }", parameter ?: "it", loopRange, body)
     }
 }

idea/testData/intentions/convertForEachToForLoop/withNestedReturn.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+fun foo() {
+    listOf(1).<caret>forEach {
+        listOf(1).forEach { return@forEach }
+    }
+}

idea/testData/intentions/convertForEachToForLoop/withNestedReturn.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+fun foo() {
+    for (it in listOf(1)) {
+        listOf(1).forEach { return@forEach }
+    }
+}

idea/testData/intentions/convertForEachToForLoop/withReturn.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+fun foo() {
+    listOf(1).<caret>forEach {
+        return@forEach
+    }
+}

idea/testData/intentions/convertForEachToForLoop/withReturn.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+fun foo() {
+    for (it in listOf(1)) {
+        continue
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -4535,6 +4535,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("withNestedReturn.kt")
+        public void testWithNestedReturn() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertForEachToForLoop/withNestedReturn.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("withReturn.kt")
+        public void testWithReturn() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertForEachToForLoop/withReturn.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("zeroArguments.kt")
         public void testZeroArguments() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertForEachToForLoop/zeroArguments.kt");