Rename: Fix renaming of function by label reference inside of lambda argument

#KT-7520 Fixed

idea/src/META-INF/plugin.xml

@@ -474,6 +474,7 @@
     <renameHandler implementation="org.jetbrains.kotlin.idea.refactoring.rename.KotlinDirectoryAsPackageRenameHandler"/>
     <renameHandler implementation="org.jetbrains.kotlin.idea.refactoring.rename.RenameBackingFieldReferenceHandler"/>
     <renameHandler implementation="org.jetbrains.kotlin.idea.refactoring.rename.RenameClassByCompanionObjectShortReferenceHandler"/>
+    <renameHandler implementation="org.jetbrains.kotlin.idea.refactoring.rename.RenameFunctionByLabeledReferenceInLambdaArgumentHandler"/>
     <automaticRenamerFactory implementation="org.jetbrains.kotlin.idea.refactoring.rename.AutomaticVariableRenamerFactory"/>
     <automaticRenamerFactory implementation="org.jetbrains.kotlin.idea.refactoring.rename.AutomaticVariableRenamerFactoryForJavaClass"/>
     <automaticRenamerFactory implementation="org.jetbrains.kotlin.idea.refactoring.rename.AutomaticVariableInJavaRenamerFactory"/>

idea/src/org/jetbrains/kotlin/idea/refactoring/rename/AbstractReferenceSubstitutionRenameHandler.kt

@@ -23,9 +23,11 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiWhiteSpace
 import com.intellij.refactoring.rename.PsiElementRenameHandler
 import com.intellij.refactoring.rename.inplace.MemberInplaceRenameHandler
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -37,7 +39,11 @@ abstract class AbstractReferenceSubstitutionRenameHandler<D : DeclarationDescrip
     protected fun getReferenceExpression(dataContext: DataContext): KtSimpleNameExpression? {
         val caret = CommonDataKeys.CARET.getData(dataContext) ?: return null
         val ktFile = CommonDataKeys.PSI_FILE.getData(dataContext) as? KtFile ?: return null
-        return ktFile.findElementAt(caret.offset)?.getNonStrictParentOfType<KtSimpleNameExpression>() ?: return null
+        var elementAtCaret = ktFile.findElementAt(caret.offset)
+        if (elementAtCaret is PsiWhiteSpace) {
+            elementAtCaret = CodeInsightUtils.getElementAtOffsetIgnoreWhitespaceAfter(ktFile, caret.offset)
+        }
+        return elementAtCaret?.getNonStrictParentOfType<KtSimpleNameExpression>()
     }
 
     protected abstract fun getTargetDescriptor(dataContext: DataContext): D?

idea/src/org/jetbrains/kotlin/idea/refactoring/rename/RenameFunctionByLabeledReferenceInLambdaArgumentHandler.kt

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.refactoring.rename
+
+import com.intellij.openapi.actionSystem.DataContext
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.references.getCalleeByLambdaArgument
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtLabelReferenceExpression
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+
+class RenameFunctionByLabeledReferenceInLambdaArgumentHandler : AbstractReferenceSubstitutionRenameHandler<FunctionDescriptor>() {
+    override fun getTargetDescriptor(dataContext: DataContext): FunctionDescriptor? {
+        val refExpr = getReferenceExpression(dataContext) as? KtLabelReferenceExpression ?: return null
+        val context = refExpr.analyze(BodyResolveMode.PARTIAL)
+        val lambda = context[BindingContext.LABEL_TARGET, refExpr] as? KtFunction ?: return null
+        val calleeExpression = lambda.getCalleeByLambdaArgument() ?: return null
+        return context[BindingContext.REFERENCE_TARGET, calleeExpression] as? FunctionDescriptor
+    }
+}

idea/testData/refactoring/rename/funWithLabeledReturns/after/test.kt

@@ -0,0 +1,7 @@
+fun <R> bar(f: () -> R) = f()
+
+fun test() {
+    bar { return@bar false }
+
+    bar(fun(): Boolean { return@bar false })
+}

idea/testData/refactoring/rename/funWithLabeledReturns/before/test.kt

@@ -0,0 +1,7 @@
+fun <R> /*rename*/foo(f: () -> R) = f()
+
+fun test() {
+    foo { return@foo false }
+
+    foo(fun(): Boolean { return@foo false })
+}

idea/testData/refactoring/rename/funWithLabeledReturns/funWithLabeledReturns.test

@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "bar"
+}

idea/testData/refactoring/rename/labeledReturnInAnonymousFun/after/test.kt

@@ -0,0 +1,5 @@
+fun <R> bar(f: () -> R) = f()
+
+fun test() {
+    bar(fun(): Boolean { return@bar false })
+}

idea/testData/refactoring/rename/labeledReturnInAnonymousFun/before/test.kt

@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo(fun(): Boolean { return@/*rename*/foo false })
+}

idea/testData/refactoring/rename/labeledReturnInAnonymousFun/labeledReturnInAnonymousFun.test

@@ -0,0 +1,5 @@
+{
+  "type": "AUTO_DETECT",
+  "mainFile": "test.kt",
+  "newName": "bar"
+}

idea/testData/refactoring/rename/labeledReturnInAnonymousFunBeforeWhitespace/after/test.kt

@@ -0,0 +1,5 @@
+fun <R> bar(f: () -> R) = f()
+
+fun test() {
+    bar(fun(): Boolean { return@bar false })
+}

idea/testData/refactoring/rename/labeledReturnInAnonymousFunBeforeWhitespace/before/test.kt

@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo(fun(): Boolean { return@foo/*rename*/ false })
+}

idea/testData/refactoring/rename/labeledReturnInAnonymousFunBeforeWhitespace/labeledReturnInAnonymousFunBeforeWhitespace.test

@@ -0,0 +1,5 @@
+{
+  "type": "AUTO_DETECT",
+  "mainFile": "test.kt",
+  "newName": "bar"
+}

idea/testData/refactoring/rename/labeledReturnInLambda/after/test.kt

@@ -0,0 +1,7 @@
+fun <R> bar(f: () -> R) = f()
+
+fun test() {
+    bar {
+        return@bar false
+    }
+}

idea/testData/refactoring/rename/labeledReturnInLambda/before/test.kt

@@ -0,0 +1,7 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo {
+        return@/*rename*/foo false
+    }
+}

idea/testData/refactoring/rename/labeledReturnInLambda/labeledReturnInLambda.test

@@ -0,0 +1,5 @@
+{
+  "type": "AUTO_DETECT",
+  "mainFile": "test.kt",
+  "newName": "bar"
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/rename/RenameTestGenerated.java

@@ -204,6 +204,30 @@ public class RenameTestGenerated extends AbstractRenameTest {
         doTest(fileName);
     }
 
+    @TestMetadata("funWithLabeledReturns/funWithLabeledReturns.test")
+    public void testFunWithLabeledReturns_FunWithLabeledReturns() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/funWithLabeledReturns/funWithLabeledReturns.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledReturnInAnonymousFun/labeledReturnInAnonymousFun.test")
+    public void testLabeledReturnInAnonymousFun_LabeledReturnInAnonymousFun() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledReturnInAnonymousFun/labeledReturnInAnonymousFun.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledReturnInAnonymousFunBeforeWhitespace/labeledReturnInAnonymousFunBeforeWhitespace.test")
+    public void testLabeledReturnInAnonymousFunBeforeWhitespace_LabeledReturnInAnonymousFunBeforeWhitespace() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledReturnInAnonymousFunBeforeWhitespace/labeledReturnInAnonymousFunBeforeWhitespace.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledReturnInLambda/labeledReturnInLambda.test")
+    public void testLabeledReturnInLambda_LabeledReturnInLambda() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledReturnInLambda/labeledReturnInLambda.test");
+        doTest(fileName);
+    }
+
     @TestMetadata("lambdaParameterRedeclaration/lambdaParameterRedeclaration.test")
     public void testLambdaParameterRedeclaration_LambdaParameterRedeclaration() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/lambdaParameterRedeclaration/lambdaParameterRedeclaration.test");