Change necessary labels in "Call on collection type may be reduced"

#KT-24492 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/collections/RenameUselessCallFix.kt

@@ -20,9 +20,8 @@ import com.intellij.codeInspection.LocalQuickFix
 import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.idea.core.replaced
-import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
-import org.jetbrains.kotlin.psi.KtQualifiedExpression
 
 class RenameUselessCallFix(val newName: String) : LocalQuickFix {
     override fun getName() = "Change call to '$newName'"
@@ -35,6 +34,24 @@ class RenameUselessCallFix(val newName: String) : LocalQuickFix {
             val selectorCallExpression = it.selectorExpression as? KtCallExpression
             val calleeExpression = selectorCallExpression?.calleeExpression ?: return
             calleeExpression.replaced(factory.createExpression(newName))
+            selectorCallExpression.renameGivenReturnLabels(factory, calleeExpression.text, newName)
+        }
+    }
+
+    private fun KtCallExpression.renameGivenReturnLabels(factory: KtPsiFactory, labelName: String, newName: String) {
+        val lambdaExpression = lambdaArguments.firstOrNull()?.getLambdaExpression() ?: return
+        val bodyExpression = lambdaExpression.bodyExpression ?: return
+
+        bodyExpression.forEachDescendantOfType<KtReturnExpression> {
+            if (it.getLabelName() != labelName) return@forEachDescendantOfType
+
+            it.replaced(
+                factory.createExpressionByPattern(
+                    "return@$0 $1",
+                    newName,
+                    it.returnedExpression ?: ""
+                )
+            )
         }
     }
 }

idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda2.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// FIX: Change call to 'map'
+
+fun foo(c: Collection<String>) {
+    c.<caret>mapNotNull {
+        return@mapNotNull ""
+    }
+}

idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda2.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// FIX: Change call to 'map'
+
+fun foo(c: Collection<String>) {
+    c.map {
+        return@map ""
+    }
+}

idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda3.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// FIX: Change call to 'map'
+
+fun foo(c: Collection<String>, f: Boolean) {
+    c.<caret>mapNotNull {
+        if (f) {
+            return@mapNotNull
+        }
+    }
+}

idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda3.kt.after

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// FIX: Change call to 'map'
+
+fun foo(c: Collection<String>, f: Boolean) {
+    c.<caret>map {
+        if (f) {
+            return@map
+        }
+    }
+}

idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda4.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// FIX: Change call to 'map'
+
+fun foo(c: Collection<String>) {
+    c.<caret>mapNotNull label@{
+        return@label ""
+    }
+}

idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda4.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// FIX: Change call to 'map'
+
+fun foo(c: Collection<String>) {
+    c.map label@{
+        return@label ""
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -1294,6 +1294,21 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda.kt");
             }
 
+            @TestMetadata("MapNotNullWithLambda2.kt")
+            public void testMapNotNullWithLambda2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda2.kt");
+            }
+
+            @TestMetadata("MapNotNullWithLambda3.kt")
+            public void testMapNotNullWithLambda3() throws Exception {
+                runTest("idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda3.kt");
+            }
+
+            @TestMetadata("MapNotNullWithLambda4.kt")
+            public void testMapNotNullWithLambda4() throws Exception {
+                runTest("idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambda4.kt");
+            }
+
             @TestMetadata("MapNotNullWithLambdaFake.kt")
             public void testMapNotNullWithLambdaFake() throws Exception {
                 runTest("idea/testData/inspectionsLocal/collections/uselessCallOnCollection/MapNotNullWithLambdaFake.kt");