Rename: Support labeled expressions

#KT-7107 Fixed
2017-04-24 20:22:07 +03:00
parent 63d32a524c
commit 3c94184de9
46 changed files with 395 additions and 52 deletions
@@ -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);
+    }
 }
@@ -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
+}
@@ -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.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
-import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
+import org.jetbrains.kotlin.psi.psiUtil.*
 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) {
-        targets.forEach {
-            if (it !is KtFunctionLiteral && !(it is KtNamedFunction && it.name.isNullOrEmpty())) return@forEach
+    val element = element
+    if (element is KtLabelReferenceExpression) {
+        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
-            if (calleeReference.matchesTarget(candidateTarget)) return true
+                val labeledExpression = it.getLabeledParent(element.getReferencedName())
+                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
+            }
        }
    }

@@ -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"/>
@@ -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"
@@ -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 -> ""
        }
@@ -71,6 +71,7 @@ class KotlinRefactoringSupportProvider : RefactoringSupportProvider() {
                    return grandparent is KtCatchClause || grandparent is KtFunctionLiteral
                }
            }
+            is KtLabeledExpression -> return true
        }
        return false
    }
@@ -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)!!
@@ -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() {
-    private val inplaceRenameHandler = MemberInplaceRenameHandler()
-
-    protected fun getReferenceExpression(dataContext: DataContext): KtSimpleNameExpression? {
+abstract class AbstractReferenceSubstitutionRenameHandler(
+        private val delegateHandler: RenameHandler = MemberInplaceRenameHandler()
+) : PsiElementRenameHandler() {
+        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 open fun getElementToRename(project: Project, descriptor: D): PsiElement? {
-        return DescriptorToSourceUtilsIde.getAnyDeclaration(project, descriptor)
-    }
+    protected abstract fun getElementToRename(dataContext: DataContext): PsiElement?

    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)) {
-            inplaceRenameHandler.invoke(project, editor, file, wrappingContext)
+        if (!ApplicationManager.getApplication().isUnitTestMode && delegateHandler.isAvailableOnDataContext(wrappingContext)) {
+            delegateHandler.invoke(project, editor, file, wrappingContext)
        }
        else {
            super.invoke(project, editor, file, wrappingContext)
@@ -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>() {
-    override fun getTargetDescriptor(dataContext: DataContext): FunctionDescriptor? {
+class RenameByLabeledReferenceInLambdaArgumentHandler :
+        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
-        return context[BindingContext.REFERENCE_TARGET, calleeExpression] as? FunctionDescriptor
+        val labeledParent = lambda.getLabeledParent(refExpr.getReferencedName())
+        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)
+        }
    }
 }
@@ -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>() {
-    override fun getTargetDescriptor(dataContext: DataContext): ClassifierDescriptor? {
+class RenameClassByCompanionObjectShortReferenceHandler : AbstractReferenceSubstitutionRenameHandler() {
+    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)
    }
 }
@@ -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
-    }
-
-    override fun getElementToRename(project: Project, descriptor: SyntheticJavaPropertyDescriptor): PsiElement? {
-        return SyntheticPropertyWrapper(PsiManager.getInstance(project), descriptor)
+        val descriptor = refExpr.analyze(BodyResolveMode.PARTIAL)[BindingContext.REFERENCE_TARGET, refExpr] as? SyntheticJavaPropertyDescriptor
+                         ?: return null
+        return SyntheticPropertyWrapper(PsiManager.getInstance(dataContext.project), descriptor)
    }
 }
@@ -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.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
-import org.jetbrains.kotlin.psi.psiUtil.isAbstract
+import org.jetbrains.kotlin.psi.psiUtil.*
 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
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo (baz@ fun(): Boolean { return@baz false })
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo (/*rename*/bar@ fun(): Boolean { return@bar false })
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "baz"
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo baz@ { return@baz false }
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo /*rename*/bar@ { return@bar false }
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "baz"
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo baz@ { return@baz false }
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo bar@ { return@/*rename*/bar false }
+}
@@ -0,0 +1,5 @@
+{
+  "type": "AUTO_DETECT",
+  "mainFile": "test.kt",
+  "newName": "baz"
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo (baz@ { return@baz false })
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo (/*rename*/bar@ { return@bar false })
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "baz"
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo baz@ { return@baz false }
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo /*rename*/foo@ { return@foo false }
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "baz"
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo (bar2@ baz@ { return@baz false })
+}
@@ -0,0 +1,5 @@
+fun <R> foo(f: () -> R) = f()
+
+fun test() {
+    foo (bar2@ /*rename*/bar@ { return@bar false })
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "baz"
+}
@@ -0,0 +1,6 @@
+fun test() {
+    bar@ for (n in 1..10) {
+        if (n == 5) continue@bar
+        if (n > 8) break@bar
+    }
+}
@@ -0,0 +1,6 @@
+fun test() {
+    /*rename*/foo@ for (n in 1..10) {
+        if (n == 5) continue@foo
+        if (n > 8) break@foo
+    }
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "bar"
+}
@@ -0,0 +1,6 @@
+fun test() {
+    bar@ for (n in 1..10) {
+        if (n == 5) continue@bar
+        if (n > 8) break@bar
+    }
+}
@@ -0,0 +1,6 @@
+fun test() {
+    foo@ for (n in 1..10) {
+        if (n == 5) continue@foo
+        if (n > 8) break@/*rename*/foo
+    }
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "bar"
+}
@@ -0,0 +1,6 @@
+fun test() {
+    bar@ for (n in 1..10) {
+        if (n == 5) continue@bar
+        if (n > 8) break@bar
+    }
+}
@@ -0,0 +1,6 @@
+fun test() {
+    foo@ for (n in 1..10) {
+        if (n == 5) continue@/*rename*/foo
+        if (n > 8) break@foo
+    }
+}
@@ -0,0 +1,5 @@
+{
+  "type": "MARKED_ELEMENT",
+  "mainFile": "test.kt",
+  "newName": "bar"
+}
@@ -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 })
+}
@@ -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 }
+}
@@ -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>
+    }
+}
@@ -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");
@@ -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.PSI_FILE.name -> psiFile
+                    CommonDataKeys.CARET.name,
+                    CommonDataKeys.PSI_ELEMENT.name,
+                    CommonDataKeys.PSI_FILE.name -> textEditorPsiDataProvider.getData(dataId, editor, currentCaret)
                    PsiElementRenameHandler.DEFAULT_NAME.name -> newName
                    else -> null
                }
@@ -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");