"map.put() to assignment": handle case with labeled return inside

So #KT-23194 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/ReplacePutWithAssignmentInspection.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.callExpression
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
@@ -31,7 +32,7 @@ import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getExplicitReceiverVa
 import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
 
 class ReplacePutWithAssignmentInspection : AbstractApplicabilityBasedInspection<KtDotQualifiedExpression>(
-        KtDotQualifiedExpression::class.java
+    KtDotQualifiedExpression::class.java
 ) {
 
     override fun isApplicable(element: KtDotQualifiedExpression): Boolean {
@@ -52,10 +53,20 @@ class ReplacePutWithAssignmentInspection : AbstractApplicabilityBasedInspection<
     override fun applyTo(element: PsiElement, project: Project, editor: Editor?) {
         val expression = element.getParentOfType<KtDotQualifiedExpression>(strict = false) ?: return
         val valueArguments = expression.callExpression?.valueArguments ?: return
-        expression.replace(KtPsiFactory(expression).createExpressionByPattern("$0[$1] = $2",
-                                                                              expression.receiverExpression,
-                                                                              valueArguments[0]?.getArgumentExpression() ?: return,
-                                                                              valueArguments[1]?.getArgumentExpression() ?: return))
+        val firstArg = valueArguments[0]?.getArgumentExpression() ?: return
+        val secondArg = valueArguments[1]?.getArgumentExpression() ?: return
+        val label = if (secondArg is KtLambdaExpression) {
+            val returnLabel = secondArg.findDescendantOfType<KtReturnExpression>()?.getLabelName()
+            compatibleNames.firstOrNull { it == returnLabel }?.plus("@") ?: ""
+        } else ""
+        expression.replace(
+            KtPsiFactory(expression).createExpressionByPattern(
+                "$0[$1] = $label$2",
+                expression.receiverExpression,
+                firstArg,
+                secondArg
+            )
+        )
     }
 
     override fun inspectionTarget(element: KtDotQualifiedExpression) = element.callExpression?.calleeExpression ?: element

idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+
+fun test(b: Boolean) {
+    val map = mutableMapOf<String, () -> Unit>()
+    map.<caret>put("") {
+        if (b) {
+            return@put
+        }
+    }
+}

idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn.kt.after

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+
+fun test(b: Boolean) {
+    val map = mutableMapOf<String, () -> Unit>()
+    map[""] = put@{
+        if (b) {
+            return@put
+        }
+    }
+}

idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn2.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+
+fun test() {
+    val map = mutableMapOf<String, () -> Unit>()
+    map.<caret>put("") {
+        listOf(1).forEach {
+            return@forEach
+        }
+    }
+}

idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn2.kt.after

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+
+fun test() {
+    val map = mutableMapOf<String, () -> Unit>()
+    map[""] = {
+        listOf(1).forEach {
+            return@forEach
+        }
+    }
+}

idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn3.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+
+fun test() {
+    val map = mutableMapOf<String, () -> Unit>()
+    map.<caret>put("") label@{
+        return@label
+    }
+}

idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn3.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+
+fun test() {
+    val map = mutableMapOf<String, () -> Unit>()
+    map[""] = label@{
+        return@label
+    }
+}

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

@@ -3951,6 +3951,24 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("putLambdaHasLabeledReturn.kt")
+        public void testPutLambdaHasLabeledReturn() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("putLambdaHasLabeledReturn2.kt")
+        public void testPutLambdaHasLabeledReturn2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn2.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("putLambdaHasLabeledReturn3.kt")
+        public void testPutLambdaHasLabeledReturn3() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/replacePutWithAssignment/putLambdaHasLabeledReturn3.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("putOnParameter.kt")
         public void testPutOnParameter() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/replacePutWithAssignment/putOnParameter.kt");