Rename: Support labeled expressions

#KT-7107 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/KtLabeledExpression.java

@@ -17,11 +17,15 @@
 package org.jetbrains.kotlin.psi;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiNameIdentifierOwner;
+import com.intellij.psi.search.LocalSearchScope;
+import com.intellij.psi.search.SearchScope;
+import com.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-public class KtLabeledExpression extends KtExpressionWithLabel {
+public class KtLabeledExpression extends KtExpressionWithLabel implements PsiNameIdentifierOwner {
-
     public KtLabeledExpression(@NotNull ASTNode node) {
         super(node);
     }
@@ -35,4 +39,34 @@ public class KtLabeledExpression extends KtExpressionWithLabel {
     public <R, D> R accept(@NotNull KtVisitor<R, D> visitor, D data) {
         return visitor.visitLabeledExpression(this, data);
     }
+
+    @Override
+    public String getName() {
+        return getLabelName();
+    }
+
+    @Override
+    public PsiElement setName(@NotNull String name) throws IncorrectOperationException {
+        KtSimpleNameExpression currentLabel = getTargetLabel();
+        if (currentLabel != null) {
+            KtSimpleNameExpression newLabel = new KtPsiFactory(getProject()).createLabeledExpression(name).getTargetLabel();
+            //noinspection ConstantConditions
+            currentLabel.replace(newLabel);
+        }
+        return this;
+    }
+
+    @Nullable
+    @Override
+    public PsiElement getNameIdentifier() {
+        KtSimpleNameExpression targetLabel = getTargetLabel();
+        if (targetLabel == null) return null;
+        return targetLabel.getIdentifier();
+    }
+
+    @NotNull
+    @Override
+    public SearchScope getUseScope() {
+        return new LocalSearchScope(this);
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt

@@ -536,4 +536,15 @@ fun KtDeclaration.hasBody() = when (this) {
 
 
 fun KtExpression.referenceExpression(): KtReferenceExpression? =
-        (if (this is KtCallExpression) calleeExpression else this) as? KtReferenceExpression
+        (if (this is KtCallExpression) calleeExpression else this) as? KtReferenceExpression
+
+fun KtExpression.getLabeledParent(labelName: String): KtLabeledExpression? {
+    parents.forEach {
+        when (it) {
+            is KtLabeledExpression -> if (it.getLabelName() == labelName) return it
+            is KtParenthesizedExpression, is KtAnnotatedExpression, is KtLambdaExpression -> return@forEach
+            else -> return null
+        }
+    }
+    return null
+}

idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt

@@ -30,10 +30,7 @@ import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getAssignmentByLHS
+import org.jetbrains.kotlin.psi.psiUtil.*
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
-import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.isReallySuccess
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
@@ -98,12 +95,25 @@ fun PsiReference.matchesTarget(candidateTarget: PsiElement): Boolean {
         return true
     }
 
-    if (element is KtLabelReferenceExpression && (element.parent as? KtContainerNode)?.parent is KtReturnExpression) {
+    val element = element
-        targets.forEach {
+    if (element is KtLabelReferenceExpression) {
-            if (it !is KtFunctionLiteral && !(it is KtNamedFunction && it.name.isNullOrEmpty())) return@forEach
+        val labelParent = (element.parent as? KtContainerNode)?.parent
+        when (labelParent) {
+            is KtReturnExpression -> targets.forEach {
+                if (it !is KtFunctionLiteral && !(it is KtNamedFunction && it.name.isNullOrEmpty())) return@forEach
+                it as KtFunction
 
-            val calleeReference = (it as KtFunction).getCalleeByLambdaArgument()?.mainReference ?: return@forEach
+                val labeledExpression = it.getLabeledParent(element.getReferencedName())
-            if (calleeReference.matchesTarget(candidateTarget)) return true
+                if (labeledExpression != null) {
+                    if (candidateTarget == labeledExpression) return true else return@forEach
+                }
+                val calleeReference = it.getCalleeByLambdaArgument()?.mainReference ?: return@forEach
+                if (calleeReference.matchesTarget(candidateTarget)) return true
+            }
+            is KtBreakExpression, is KtContinueExpression -> targets.forEach {
+                val labeledExpression = (it as? KtExpression)?.getLabeledParent(element.getReferencedName()) ?: return@forEach
+                if (candidateTarget == labeledExpression) return true
+            }
         }
     }
 

idea/src/META-INF/plugin.xml

@@ -497,7 +497,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"/>
+    <renameHandler implementation="org.jetbrains.kotlin.idea.refactoring.rename.RenameByLabeledReferenceInLambdaArgumentHandler"/>
     <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"/>
@@ -755,6 +755,8 @@
         language="kotlin"
         implementationClass="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinTargetElementEvaluator" />
 
+    <targetElementUtilExtender implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinTargetElementEvaluator"/>
+
     <refactoring.pullUpHelperFactory
         language="kotlin"
         implementationClass="org.jetbrains.kotlin.idea.refactoring.pullUp.KotlinPullUpHelperFactory"/>

idea/src/org/jetbrains/kotlin/idea/findUsages/KotlinElementDescriptionProvider.kt

@@ -124,6 +124,7 @@ class KotlinElementDescriptionProvider : ElementDescriptionProvider {
             is KtParameter -> "parameter"
             is KtDestructuringDeclarationEntry -> "variable"
             is KtTypeAlias -> "type alias"
+            is KtLabeledExpression -> "label"
             is RenameJavaSyntheticPropertyHandler.SyntheticPropertyWrapper -> "property"
             is KtLightClassForFacade -> "facade class"
             is RenameKotlinPropertyProcessor.PropertyMethodWrapper -> "property accessor"

idea/src/org/jetbrains/kotlin/idea/findUsages/KotlinFindUsagesProvider.kt

@@ -79,6 +79,7 @@ class KotlinFindUsagesProvider : FindUsagesProvider {
                 val funDescription = "$name$paramsDescription" + (returnTypeDescription?.let { ": $it" } ?: "")
                 return funDescription + (element.containerDescription?.let { " of $it" } ?: "")
             }
+            is KtLabeledExpression -> element.getLabelName() ?: ""
             is KtLightElement<*, *> -> element.kotlinOrigin?.let { getDescriptiveName(it) } ?: ""
             else -> ""
         }

idea/src/org/jetbrains/kotlin/idea/refactoring/KotlinRefactoringSupportProvider.kt

@@ -71,6 +71,7 @@ class KotlinRefactoringSupportProvider : RefactoringSupportProvider() {
                     return grandparent is KtCatchClause || grandparent is KtFunctionLiteral
                 }
             }
+            is KtLabeledExpression -> return true
         }
         return false
     }

idea/src/org/jetbrains/kotlin/idea/refactoring/dataContextUtils.kt

@@ -0,0 +1,24 @@
+/*
+ * 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
+
+import com.intellij.openapi.actionSystem.CommonDataKeys
+import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.project.Project
+
+val DataContext.project: Project
+    get() = CommonDataKeys.PROJECT.getData(this)!!

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

@@ -25,18 +25,17 @@ 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.RenameHandler
 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
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 
-abstract class AbstractReferenceSubstitutionRenameHandler<D : DeclarationDescriptor> : PsiElementRenameHandler() {
+abstract class AbstractReferenceSubstitutionRenameHandler(
-    private val inplaceRenameHandler = MemberInplaceRenameHandler()
+        private val delegateHandler: RenameHandler = MemberInplaceRenameHandler()
-
+) : PsiElementRenameHandler() {
-    protected fun getReferenceExpression(dataContext: DataContext): KtSimpleNameExpression? {
+        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
         var elementAtCaret = ktFile.findElementAt(caret.offset)
@@ -46,25 +45,21 @@ abstract class AbstractReferenceSubstitutionRenameHandler<D : DeclarationDescrip
         return elementAtCaret?.getNonStrictParentOfType<KtSimpleNameExpression>()
     }
 
-    protected abstract fun getTargetDescriptor(dataContext: DataContext): D?
+    protected abstract fun getElementToRename(dataContext: DataContext): PsiElement?
-
-    protected open fun getElementToRename(project: Project, descriptor: D): PsiElement? {
-        return DescriptorToSourceUtilsIde.getAnyDeclaration(project, descriptor)
-    }
 
     override fun isAvailableOnDataContext(dataContext: DataContext): Boolean {
-        return CommonDataKeys.EDITOR.getData(dataContext) != null && getTargetDescriptor(dataContext) != null
+        return CommonDataKeys.EDITOR.getData(dataContext) != null && getElementToRename(dataContext) != null
     }
 
     override fun invoke(project: Project, editor: Editor?, file: PsiFile?, dataContext: DataContext) {
-        val descriptor = getTargetDescriptor(dataContext) ?: return
+        val elementToRename = getElementToRename(dataContext) ?: return
         val wrappingContext = DataContext { id ->
-            if (CommonDataKeys.PSI_ELEMENT.`is`(id)) return@DataContext getElementToRename(project, descriptor)
+            if (CommonDataKeys.PSI_ELEMENT.`is`(id)) return@DataContext elementToRename
             dataContext.getData(id)
         }
         // Can't provide new name for inplace refactoring in unit test mode
-        if (!ApplicationManager.getApplication().isUnitTestMode && inplaceRenameHandler.isAvailableOnDataContext(wrappingContext)) {
+        if (!ApplicationManager.getApplication().isUnitTestMode && delegateHandler.isAvailableOnDataContext(wrappingContext)) {
-            inplaceRenameHandler.invoke(project, editor, file, wrappingContext)
+            delegateHandler.invoke(project, editor, file, wrappingContext)
         }
         else {
             super.invoke(project, editor, file, wrappingContext)

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

@@ -17,20 +17,33 @@
 package org.jetbrains.kotlin.idea.refactoring.rename
 
 import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.psi.PsiElement
+import com.intellij.refactoring.rename.inplace.VariableInplaceRenameHandler
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
+import org.jetbrains.kotlin.idea.refactoring.project
 import org.jetbrains.kotlin.idea.references.getCalleeByLambdaArgument
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtLabelReferenceExpression
+import org.jetbrains.kotlin.psi.psiUtil.getLabeledParent
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
-class RenameFunctionByLabeledReferenceInLambdaArgumentHandler : AbstractReferenceSubstitutionRenameHandler<FunctionDescriptor>() {
+class RenameByLabeledReferenceInLambdaArgumentHandler :
-    override fun getTargetDescriptor(dataContext: DataContext): FunctionDescriptor? {
+        AbstractReferenceSubstitutionRenameHandler(VariableInplaceRenameHandler()) {
+    override fun getElementToRename(dataContext: DataContext): PsiElement? {
         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
+        val labeledParent = lambda.getLabeledParent(refExpr.getReferencedName())
-        return context[BindingContext.REFERENCE_TARGET, calleeExpression] as? FunctionDescriptor
+        return if (labeledParent != null) {
+            labeledParent
+        }
+        else {
+            val calleeExpression = lambda.getCalleeByLambdaArgument() ?: return null
+            val descriptor = context[BindingContext.REFERENCE_TARGET, calleeExpression] as? FunctionDescriptor ?: return null
+            DescriptorToSourceUtilsIde.getAnyDeclaration(dataContext.project, descriptor)
+        }
     }
 }

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

@@ -17,14 +17,17 @@
 package org.jetbrains.kotlin.idea.refactoring.rename
 
 import com.intellij.openapi.actionSystem.DataContext
-import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
+import org.jetbrains.kotlin.idea.refactoring.project
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
-class RenameClassByCompanionObjectShortReferenceHandler : AbstractReferenceSubstitutionRenameHandler<ClassifierDescriptor>() {
+class RenameClassByCompanionObjectShortReferenceHandler : AbstractReferenceSubstitutionRenameHandler() {
-    override fun getTargetDescriptor(dataContext: DataContext): ClassifierDescriptor? {
+    override fun getElementToRename(dataContext: DataContext): PsiElement? {
         val refExpr = getReferenceExpression(dataContext) ?: return null
-        return refExpr.analyze(BodyResolveMode.PARTIAL)[BindingContext.SHORT_REFERENCE_TO_COMPANION_OBJECT, refExpr]
+        val descriptor = refExpr.analyze(BodyResolveMode.PARTIAL)[BindingContext.SHORT_REFERENCE_TO_COMPANION_OBJECT, refExpr] ?: return null
+        return DescriptorToSourceUtilsIde.getAnyDeclaration(dataContext.project, descriptor)
     }
 }

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

@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.idea.refactoring.rename
 
 import com.intellij.openapi.actionSystem.DataContext
-import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiMethod
@@ -27,6 +26,7 @@ import com.intellij.psi.search.SearchScope
 import com.intellij.refactoring.rename.RenamePsiElementProcessor
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.refactoring.project
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -34,7 +34,7 @@ import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
 
-class RenameJavaSyntheticPropertyHandler : AbstractReferenceSubstitutionRenameHandler<SyntheticJavaPropertyDescriptor>() {
+class RenameJavaSyntheticPropertyHandler : AbstractReferenceSubstitutionRenameHandler() {
     class Processor : RenamePsiElementProcessor() {
         override fun prepareRenaming(element: PsiElement, newName: String, allRenames: MutableMap<PsiElement, String>, scope: SearchScope) {
             val propertyWrapper = element as? SyntheticPropertyWrapper ?: return
@@ -68,12 +68,10 @@ class RenameJavaSyntheticPropertyHandler : AbstractReferenceSubstitutionRenameHa
         }
     }
 
-    override fun getTargetDescriptor(dataContext: DataContext): SyntheticJavaPropertyDescriptor? {
+    override fun getElementToRename(dataContext: DataContext): PsiElement? {
         val refExpr = getReferenceExpression(dataContext) ?: return null
-        return refExpr.analyze(BodyResolveMode.PARTIAL)[BindingContext.REFERENCE_TARGET, refExpr] as? SyntheticJavaPropertyDescriptor
+        val descriptor = refExpr.analyze(BodyResolveMode.PARTIAL)[BindingContext.REFERENCE_TARGET, refExpr] as? SyntheticJavaPropertyDescriptor
-    }
+                         ?: return null
-
+        return SyntheticPropertyWrapper(PsiManager.getInstance(dataContext.project), descriptor)
-    override fun getElementToRename(project: Project, descriptor: SyntheticJavaPropertyDescriptor): PsiElement? {
-        return SyntheticPropertyWrapper(PsiManager.getInstance(project), descriptor)
     }
 }

idea/src/org/jetbrains/kotlin/idea/search/ideaExtensions/KotlinTargetElementEvaluator.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.search.ideaExtensions
 import com.intellij.codeInsight.JavaTargetElementEvaluator
 import com.intellij.codeInsight.TargetElementEvaluatorEx
 import com.intellij.codeInsight.TargetElementUtil
+import com.intellij.codeInsight.TargetElementUtilExtender
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiReference
@@ -27,17 +28,17 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithSource
 import org.jetbrains.kotlin.idea.intentions.isAutoCreatedItUsage
 import org.jetbrains.kotlin.idea.references.KtDestructuringDeclarationReference
+import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.idea.references.resolveMainReferenceToDescriptors
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.*
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
-import org.jetbrains.kotlin.psi.psiUtil.isAbstract
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.source.getPsi
 
-class KotlinTargetElementEvaluator : TargetElementEvaluatorEx {
+class KotlinTargetElementEvaluator : TargetElementEvaluatorEx, TargetElementUtilExtender {
     companion object {
+        val DO_NOT_UNWRAP_LABELED_EXPRESSION = 0x100
+
         // Place caret after the open curly brace in lambda for generated 'it'
         fun findLambdaOpenLBraceForGeneratedIt(ref: PsiReference): PsiElement? {
             val element: PsiElement = ref.element
@@ -66,9 +67,23 @@ class KotlinTargetElementEvaluator : TargetElementEvaluatorEx {
         }
     }
 
+    override fun getAdditionalDefinitionSearchFlags() = 0
+
+    override fun getAdditionalReferenceSearchFlags() = DO_NOT_UNWRAP_LABELED_EXPRESSION
+
+    override fun getAllAdditionalFlags() = DO_NOT_UNWRAP_LABELED_EXPRESSION
+
     override fun includeSelfInGotoImplementation(element: PsiElement): Boolean = !(element is KtClass && element.isAbstract())
 
     override fun getElementByReference(ref: PsiReference, flags: Int): PsiElement? {
+        if (ref is KtSimpleNameReference && ref.expression is KtLabelReferenceExpression) {
+            val refTarget = ref.resolve() as? KtExpression ?: return null
+            if (!BitUtil.isSet(flags, DO_NOT_UNWRAP_LABELED_EXPRESSION)) {
+                return refTarget.getLabeledParent(ref.expression.getReferencedName()) ?: refTarget
+            }
+            return refTarget
+        }
+
         // prefer destructing declaration entry to its target if element name is accepted
         if (ref is KtDestructuringDeclarationReference && BitUtil.isSet(flags, TargetElementUtil.ELEMENT_NAME_ACCEPTED)) {
             return ref.element

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

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

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

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

idea/testData/refactoring/rename/labeledAnonymousFunByLabel/labeledLambdaByLabel.test

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

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

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

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

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

idea/testData/refactoring/rename/labeledLambdaByLabel/labeledLambdaByLabel.test

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

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

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

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

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

idea/testData/refactoring/rename/labeledLambdaByLabelRef/labeledLambdaByLabelRef.test

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

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

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

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

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

idea/testData/refactoring/rename/labeledLambdaByLabelWithParens/labeledLambdaByLabelWithParens.test

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

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

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

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

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

idea/testData/refactoring/rename/labeledLambdaByLabelWithSameName/labeledLambdaByLabelWithSameName.test

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

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

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

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

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

idea/testData/refactoring/rename/labeledLambdaWithMultipleLabels/labeledLambdaWithMultipleLabels.test

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

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

@@ -0,0 +1,6 @@
+fun test() {
+    bar@ for (n in 1..10) {
+        if (n == 5) continue@bar
+        if (n > 8) break@bar
+    }
+}

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

@@ -0,0 +1,6 @@
+fun test() {
+    /*rename*/foo@ for (n in 1..10) {
+        if (n == 5) continue@foo
+        if (n > 8) break@foo
+    }
+}

idea/testData/refactoring/rename/labeledLoopByLabel/labeledLoopByLabel.test

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

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

@@ -0,0 +1,6 @@
+fun test() {
+    bar@ for (n in 1..10) {
+        if (n == 5) continue@bar
+        if (n > 8) break@bar
+    }
+}

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

@@ -0,0 +1,6 @@
+fun test() {
+    foo@ for (n in 1..10) {
+        if (n == 5) continue@foo
+        if (n > 8) break@/*rename*/foo
+    }
+}

idea/testData/refactoring/rename/labeledLoopByLabelRefInBreak/labeledLoopByLabelRefInBreak.test

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

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

@@ -0,0 +1,6 @@
+fun test() {
+    bar@ for (n in 1..10) {
+        if (n == 5) continue@bar
+        if (n > 8) break@bar
+    }
+}

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

@@ -0,0 +1,6 @@
+fun test() {
+    foo@ for (n in 1..10) {
+        if (n == 5) continue@/*rename*/foo
+        if (n > 8) break@foo
+    }
+}

idea/testData/refactoring/rename/labeledLoopByLabelRefInContinue/labeledLoopByLabelRefInContinue.test

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

idea/testData/usageHighlighter/labeledAnonymousFun.kt

@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo (<info descr="null">~bar</info>@ fun(): Boolean { return@<info descr="null">bar</info> false })
+}

idea/testData/usageHighlighter/labeledLambda.kt

@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo <info descr="null">~bar</info>@ { return@<info descr="null">bar</info> false }
+}

idea/testData/usageHighlighter/labeledLoop.kt

@@ -0,0 +1,6 @@
+fun test() {
+    <info descr="null">~foo</info>@ for (n in 1..10) {
+        if (n == 5) continue@<info descr="null">foo</info>
+        if (n > 8) break@<info descr="null">foo</info>
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/highlighter/UsageHighlightingTestGenerated.java

@@ -36,6 +36,24 @@ public class UsageHighlightingTestGenerated extends AbstractUsageHighlightingTes
         KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/usageHighlighter"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
     }
 
+    @TestMetadata("labeledAnonymousFun.kt")
+    public void testLabeledAnonymousFun() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/usageHighlighter/labeledAnonymousFun.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLambda.kt")
+    public void testLabeledLambda() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/usageHighlighter/labeledLambda.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLoop.kt")
+    public void testLabeledLoop() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/usageHighlighter/labeledLoop.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("localVal.kt")
     public void testLocalVal() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/usageHighlighter/localVal.kt");

idea/tests/org/jetbrains/kotlin/idea/refactoring/rename/AbstractRenameTest.kt

@@ -26,6 +26,7 @@ import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.editor.EditorFactory
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.fileEditor.impl.text.TextEditorPsiDataProvider
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ModuleRootManager
@@ -333,12 +334,17 @@ abstract class AbstractRenameTest : KotlinLightCodeInsightFixtureTestCase() {
             assert(marker != -1)
 
             editor.caretModel.moveToOffset(marker)
+            val currentCaret = editor.caretModel.currentCaret
+
+            val textEditorPsiDataProvider = TextEditorPsiDataProvider()
 
             val dataContext = DataContext { dataId ->
                 when (dataId) {
+                    CommonDataKeys.PROJECT.name -> project
                     CommonDataKeys.EDITOR.name -> editor
-                    CommonDataKeys.CARET.name -> editor.caretModel.currentCaret
+                    CommonDataKeys.CARET.name,
-                    CommonDataKeys.PSI_FILE.name -> psiFile
+                    CommonDataKeys.PSI_ELEMENT.name,
+                    CommonDataKeys.PSI_FILE.name -> textEditorPsiDataProvider.getData(dataId, editor, currentCaret)
                     PsiElementRenameHandler.DEFAULT_NAME.name -> newName
                     else -> null
                 }

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

@@ -204,6 +204,60 @@ public class RenameTestGenerated extends AbstractRenameTest {
         doTest(fileName);
     }
 
+    @TestMetadata("labeledAnonymousFunByLabel/labeledLambdaByLabel.test")
+    public void testLabeledAnonymousFunByLabel_LabeledLambdaByLabel() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledAnonymousFunByLabel/labeledLambdaByLabel.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLambdaByLabel/labeledLambdaByLabel.test")
+    public void testLabeledLambdaByLabel_LabeledLambdaByLabel() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLambdaByLabel/labeledLambdaByLabel.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLambdaByLabelRef/labeledLambdaByLabelRef.test")
+    public void testLabeledLambdaByLabelRef_LabeledLambdaByLabelRef() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLambdaByLabelRef/labeledLambdaByLabelRef.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLambdaByLabelWithParens/labeledLambdaByLabelWithParens.test")
+    public void testLabeledLambdaByLabelWithParens_LabeledLambdaByLabelWithParens() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLambdaByLabelWithParens/labeledLambdaByLabelWithParens.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLambdaByLabelWithSameName/labeledLambdaByLabelWithSameName.test")
+    public void testLabeledLambdaByLabelWithSameName_LabeledLambdaByLabelWithSameName() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLambdaByLabelWithSameName/labeledLambdaByLabelWithSameName.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLambdaWithMultipleLabels/labeledLambdaWithMultipleLabels.test")
+    public void testLabeledLambdaWithMultipleLabels_LabeledLambdaWithMultipleLabels() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLambdaWithMultipleLabels/labeledLambdaWithMultipleLabels.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLoopByLabel/labeledLoopByLabel.test")
+    public void testLabeledLoopByLabel_LabeledLoopByLabel() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLoopByLabel/labeledLoopByLabel.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLoopByLabelRefInBreak/labeledLoopByLabelRefInBreak.test")
+    public void testLabeledLoopByLabelRefInBreak_LabeledLoopByLabelRefInBreak() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLoopByLabelRefInBreak/labeledLoopByLabelRefInBreak.test");
+        doTest(fileName);
+    }
+
+    @TestMetadata("labeledLoopByLabelRefInContinue/labeledLoopByLabelRefInContinue.test")
+    public void testLabeledLoopByLabelRefInContinue_LabeledLoopByLabelRefInContinue() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledLoopByLabelRefInContinue/labeledLoopByLabelRefInContinue.test");
+        doTest(fileName);
+    }
+
     @TestMetadata("labeledReturnInAnonymousFun/labeledReturnInAnonymousFun.test")
     public void testLabeledReturnInAnonymousFun_LabeledReturnInAnonymousFun() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/rename/labeledReturnInAnonymousFun/labeledReturnInAnonymousFun.test");