Move KtReference interface to psi module

compiler/psi/src/org/jetbrains/kotlin/idea/references/KDocReference.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
+import org.jetbrains.kotlin.name.Name
+
+abstract class KDocReference(element: KDocName) : KtMultiReference<KDocName>(element) {
+    override fun getRangeInElement(): TextRange = element.getNameTextRange()
+
+    override fun canRename(): Boolean = true
+
+    override fun resolve(): PsiElement? = multiResolve(false).firstOrNull()?.element
+
+    abstract override fun handleElementRename(newElementName: String): PsiElement?
+
+    override fun getCanonicalText(): String = element.getNameText()
+
+    override val resolvesByNames: Collection<Name> get() = listOf(Name.identifier(element.getNameText()))
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KotlinPsiReferenceProvider.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.components.ServiceManager
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiReference
+import com.intellij.util.containers.MultiMap
+import org.jetbrains.kotlin.psi.KtElement
+
+interface KotlinPsiReferenceProvider {
+    fun getReferencesByElement(element: PsiElement): Array<PsiReference>
+}
+
+interface KotlinReferenceProviderContributor {
+    fun registerReferenceProviders(registrar: KotlinPsiReferenceRegistrar)
+
+    companion object {
+        fun getInstance(): KotlinReferenceProviderContributor =
+            ServiceManager.getService(KotlinReferenceProviderContributor::class.java)
+    }
+}
+
+
+class KotlinPsiReferenceRegistrar {
+    val providers: MultiMap<Class<out PsiElement>, KotlinPsiReferenceProvider> = MultiMap.create()
+
+    inline fun <reified E : KtElement> registerProvider(crossinline factory: (E) -> PsiReference?) {
+        registerMultiProvider<E> { element ->
+            factory(element)?.let { reference -> arrayOf(reference) } ?: PsiReference.EMPTY_ARRAY
+        }
+    }
+
+    inline fun <reified E : KtElement>  registerMultiProvider(crossinline factory: (E) -> Array<PsiReference>) {
+        val provider: KotlinPsiReferenceProvider = object : KotlinPsiReferenceProvider {
+            override fun getReferencesByElement(element: PsiElement): Array<PsiReference> {
+                return factory(element as E)
+            }
+        }
+
+        registerMultiProvider(E::class.java, provider)
+    }
+
+    fun registerMultiProvider(klass: Class<out PsiElement>, provider: KotlinPsiReferenceProvider) {
+        providers.putValue(klass, provider)
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtArrayAccessReference.kt

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.google.common.collect.Lists
+import com.intellij.psi.MultiRangeReference
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.lexer.KtToken
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getPossiblyQualifiedCallExpression
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+abstract class KtArrayAccessReference(
+    expression: KtArrayAccessExpression
+) : KtSimpleReference<KtArrayAccessExpression>(expression), MultiRangeReference {
+    override val resolvesByNames: Collection<Name>
+        get() = NAMES
+
+    override fun getRangeInElement() = element.textRange.shiftRight(-element.textOffset)
+
+    private fun getBracketRange(bracketToken: KtToken) =
+        expression.indicesNode.node.findChildByType(bracketToken)?.textRange?.shiftRight(-expression.textOffset)
+
+    override fun getRanges() = listOfNotNull(getBracketRange(KtTokens.LBRACKET), getBracketRange(KtTokens.RBRACKET))
+
+    override fun canRename() = true
+
+    override fun handleElementRename(newElementName: String): PsiElement? {
+        val arrayAccessExpression = expression
+        if (OperatorNameConventions.INVOKE.asString() == newElementName) {
+            val replacement = KtPsiFactory(arrayAccessExpression.project).buildExpression {
+                val arrayExpression = arrayAccessExpression.arrayExpression
+                if (arrayExpression is KtQualifiedExpression) {
+                    appendExpression(arrayExpression.receiverExpression)
+                    appendFixedText(arrayExpression.operationSign.value)
+                    appendExpression(arrayExpression.selectorExpression)
+                } else {
+                    appendExpression(arrayExpression)
+                }
+
+                appendFixedText("(")
+                appendExpressions(arrayAccessExpression.indexExpressions, ",")
+                appendFixedText(")")
+            }
+            val fullCallExpression = arrayAccessExpression.replace(replacement) as KtExpression
+            val callExpression = fullCallExpression.getPossiblyQualifiedCallExpression()
+            if (callExpression != null && canMoveLambdaOutsideParentheses(callExpression)) {
+                moveFunctionLiteralOutsideParentheses(callExpression)
+            }
+            return fullCallExpression
+        }
+
+        return doRenameImplicitConventionalCall(newElementName)
+    }
+
+    protected abstract fun moveFunctionLiteralOutsideParentheses(callExpression: KtCallExpression)
+    protected abstract fun canMoveLambdaOutsideParentheses(callExpression: KtCallExpression): Boolean
+    protected abstract fun doRenameImplicitConventionalCall(newName: String?): KtExpression
+
+    companion object {
+        private val NAMES = Lists.newArrayList(OperatorNameConventions.GET, OperatorNameConventions.SET)
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtCollectionLiteralReference.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.MultiRangeReference
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.ArrayFqNames
+
+abstract class KtCollectionLiteralReference(expression: KtCollectionLiteralExpression) :
+    KtSimpleReference<KtCollectionLiteralExpression>(expression), MultiRangeReference {
+    companion object {
+        private val COLLECTION_LITERAL_CALL_NAMES = ArrayFqNames.PRIMITIVE_TYPE_TO_ARRAY.values + ArrayFqNames.ARRAY_OF_FUNCTION
+    }
+
+    override fun getRangeInElement(): TextRange = element.normalizeRange()
+
+    override fun getRanges(): List<TextRange> {
+        return listOfNotNull(element.leftBracket?.normalizeRange(), element.rightBracket?.normalizeRange())
+    }
+
+    override val resolvesByNames: Collection<Name>
+        get() = COLLECTION_LITERAL_CALL_NAMES
+
+    private fun PsiElement.normalizeRange(): TextRange = this.textRange.shiftRight(-expression.textOffset)
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtConstructorDelegationReference.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+
+abstract class KtConstructorDelegationReference(
+    expression: KtConstructorDelegationReferenceExpression
+) : KtSimpleReference<KtConstructorDelegationReferenceExpression>(expression) {
+    override fun getRangeInElement(): TextRange {
+        return TextRange(0, element.textLength)
+    }
+
+    override val resolvesByNames: Collection<Name>
+        get() = emptyList()
+
+    override fun handleElementRename(newElementName: String): PsiElement? {
+        // Class rename never affects this reference, so there is no need to fail with exception
+        return expression
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtDestructuringDeclarationReference.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import com.intellij.util.IncorrectOperationException
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtDestructuringDeclarationEntry
+
+abstract class KtDestructuringDeclarationReference(
+    element: KtDestructuringDeclarationEntry
+) : AbstractKtReference<KtDestructuringDeclarationEntry>(element) {
+
+    override fun getRangeInElement() = TextRange(0, element.textLength)
+
+    abstract override fun canRename(): Boolean
+
+    override fun handleElementRename(newElementName: String): PsiElement? {
+        if (canRename()) return expression
+        throw IncorrectOperationException()
+    }
+
+    override val resolvesByNames: Collection<Name>
+        get() {
+            val destructuringParent = element.parent as? KtDestructuringDeclaration ?: return emptyList()
+            val componentIndex = destructuringParent.entries.indexOf(element) + 1
+            return listOf(Name.identifier("component$componentIndex"))
+        }
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtForLoopInReference.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+abstract class KtForLoopInReference(element: KtForExpression) : KtMultiReference<KtForExpression>(element) {
+    override fun getRangeInElement(): TextRange {
+        val inKeyword = expression.inKeyword ?: return TextRange.EMPTY_RANGE
+
+        val offset = inKeyword.startOffsetInParent
+        return TextRange(offset, offset + inKeyword.textLength)
+    }
+
+    override val resolvesByNames: Collection<Name>
+        get() = NAMES
+
+    companion object {
+        private val NAMES = listOf(
+            OperatorNameConventions.ITERATOR,
+            OperatorNameConventions.NEXT,
+            OperatorNameConventions.HAS_NEXT
+        )
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtInvokeFunctionReference.kt

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.lang.ASTNode
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.MultiRangeReference
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
+import org.jetbrains.kotlin.util.OperatorNameConventions
+import java.util.*
+
+abstract class KtInvokeFunctionReference(expression: KtCallExpression) : KtSimpleReference<KtCallExpression>(expression), MultiRangeReference {
+    override val resolvesByNames: Collection<Name> get() = NAMES
+
+    override fun getRangeInElement(): TextRange {
+        return element.textRange.shiftRight(-element.textOffset)
+    }
+
+    override fun getRanges(): List<TextRange> {
+        val list = ArrayList<TextRange>()
+        val valueArgumentList = expression.valueArgumentList
+        if (valueArgumentList != null) {
+            if (valueArgumentList.arguments.isNotEmpty()) {
+                val valueArgumentListNode = valueArgumentList.node
+                val lPar = valueArgumentListNode.findChildByType(KtTokens.LPAR)
+                if (lPar != null) {
+                    list.add(getRange(lPar))
+                }
+
+                val rPar = valueArgumentListNode.findChildByType(KtTokens.RPAR)
+                if (rPar != null) {
+                    list.add(getRange(rPar))
+                }
+            } else {
+                list.add(getRange(valueArgumentList.node))
+            }
+        }
+
+        val functionLiteralArguments = expression.lambdaArguments
+        for (functionLiteralArgument in functionLiteralArguments) {
+            val functionLiteralExpression = functionLiteralArgument.getLambdaExpression() ?: continue
+            list.add(getRange(functionLiteralExpression.leftCurlyBrace))
+            val rightCurlyBrace = functionLiteralExpression.rightCurlyBrace
+            if (rightCurlyBrace != null) {
+                list.add(getRange(rightCurlyBrace))
+            }
+        }
+
+        return list
+    }
+
+    private fun getRange(node: ASTNode): TextRange {
+        val textRange = node.textRange
+        return textRange.shiftRight(-expression.textOffset)
+    }
+
+    override fun canRename(): Boolean = true
+
+    override fun handleElementRename(newElementName: String): PsiElement? {
+        val callExpression = expression
+        val fullCallExpression = callExpression.getQualifiedExpressionForSelectorOrThis()
+        if (newElementName == OperatorNameConventions.GET.asString() && callExpression.typeArguments.isEmpty()) {
+            val arrayAccessExpression = KtPsiFactory(callExpression).buildExpression {
+                if (fullCallExpression is KtQualifiedExpression) {
+                    appendExpression(fullCallExpression.receiverExpression)
+                    appendFixedText(fullCallExpression.operationSign.value)
+                }
+                appendExpression(callExpression.calleeExpression)
+                appendFixedText("[")
+                appendExpressions(callExpression.valueArguments.map { it.getArgumentExpression() })
+                appendFixedText("]")
+            }
+            return fullCallExpression.replace(arrayAccessExpression)
+        }
+
+        return doRenameImplicitConventionalCall(newElementName)
+    }
+
+    protected abstract fun doRenameImplicitConventionalCall(newName: String?): KtExpression
+
+    companion object {
+        private val NAMES = listOf(OperatorNameConventions.INVOKE)
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtPropertyDelegationMethodsReference.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtPropertyDelegate
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+abstract class KtPropertyDelegationMethodsReference(element: KtPropertyDelegate) : KtMultiReference<KtPropertyDelegate>(element) {
+    override fun getRangeInElement(): TextRange {
+        val byKeywordNode = expression.byKeywordNode
+        val offset = byKeywordNode.psi!!.startOffsetInParent
+        return TextRange(offset, offset + byKeywordNode.textLength)
+    }
+
+       override val resolvesByNames: Collection<Name> get() = NAMES
+
+    companion object {
+        private val NAMES = listOf(
+            OperatorNameConventions.GET_VALUE,
+            OperatorNameConventions.SET_VALUE,
+            OperatorNameConventions.PROVIDE_DELEGATE
+        )
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtReference.kt

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.psi.*
+import com.intellij.psi.impl.source.resolve.ResolveCache
+import com.intellij.util.IncorrectOperationException
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtReferenceExpression
+
+interface KtReference : PsiPolyVariantReference {
+    val resolver: ResolveCache.PolyVariantResolver<KtReference>
+
+    override fun getElement(): KtElement
+
+    val resolvesByNames: Collection<Name>
+}
+
+abstract class AbstractKtReference<T : KtElement>(element: T) : PsiPolyVariantReferenceBase<T>(element), KtReference {
+    val expression: T
+        get() = element
+
+    override fun multiResolve(incompleteCode: Boolean): Array<ResolveResult> =
+        ResolveCache.getInstance(expression.project).resolveWithCaching(this, resolver, false, incompleteCode)
+
+    override fun getCanonicalText(): String = "<TBD>"
+
+    open fun canRename(): Boolean = false
+    override fun handleElementRename(newElementName: String): PsiElement? = throw IncorrectOperationException()
+
+    override fun bindToElement(element: PsiElement): PsiElement = throw IncorrectOperationException()
+
+    @Suppress("UNCHECKED_CAST")
+    override fun getVariants(): Array<Any> = PsiReference.EMPTY_ARRAY as Array<Any>
+
+    override fun isSoft(): Boolean = false
+
+    override fun toString() = this::class.java.simpleName + ": " + expression.text
+}
+
+abstract class KtSimpleReference<T : KtReferenceExpression>(expression: T) : AbstractKtReference<T>(expression)
+
+abstract class KtMultiReference<T : KtElement>(expression: T) : AbstractKtReference<T>(expression)

compiler/psi/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.extensions.Extensions
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.types.expressions.OperatorConventions
+
+abstract class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSimpleReference<KtSimpleNameExpression>(expression) {
+    protected abstract fun doCanBeReferenceTo(candidateTarget: PsiElement): Boolean
+
+    override fun getRangeInElement(): TextRange {
+        val element = element.getReferencedNameElement()
+        val startOffset = getElement().startOffset
+        return element.textRange.shiftRight(-startOffset)
+    }
+
+    override fun canRename(): Boolean {
+        if (expression.getParentOfTypeAndBranch<KtWhenConditionInRange>(strict = true) { operationReference } != null) return false
+
+        val elementType = expression.getReferencedNameElementType()
+        if (elementType == KtTokens.PLUSPLUS || elementType == KtTokens.MINUSMINUS) return false
+
+        return true
+    }
+
+   abstract override fun handleElementRename(newElementName: String): PsiElement?
+
+    enum class ShorteningMode {
+        NO_SHORTENING,
+        DELAYED_SHORTENING,
+        FORCED_SHORTENING
+    }
+
+    // By default reference binding is delayed
+    override fun bindToElement(element: PsiElement): PsiElement =
+        bindToElement(element, ShorteningMode.DELAYED_SHORTENING)
+
+    abstract fun bindToElement(element: PsiElement, shorteningMode: ShorteningMode): PsiElement
+
+    abstract fun bindToFqName(
+        fqName: FqName,
+        shorteningMode: ShorteningMode = ShorteningMode.DELAYED_SHORTENING,
+        targetElement: PsiElement? = null
+    ): PsiElement
+
+    override fun getCanonicalText(): String = expression.text
+
+    override val resolvesByNames: Collection<Name>
+        get() {
+            val element = element
+
+            if (element is KtOperationReferenceExpression) {
+                val tokenType = element.operationSignTokenType
+                if (tokenType != null) {
+                    val name = OperatorConventions.getNameForOperationSymbol(
+                        tokenType, element.parent is KtUnaryExpression, element.parent is KtBinaryExpression
+                    ) ?: return emptyList()
+                    val counterpart = OperatorConventions.ASSIGNMENT_OPERATION_COUNTERPARTS[tokenType]
+                    return if (counterpart != null) {
+                        val counterpartName = OperatorConventions.getNameForOperationSymbol(counterpart, false, true)!!
+                        listOf(name, counterpartName)
+                    } else {
+                        listOf(name)
+                    }
+                }
+            }
+
+            return listOf(element.getReferencedNameAsName())
+        }
+
+    abstract fun getImportAlias(): KtImportAlias?
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/SyntheticPropertyAccessorReference.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiMethod
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+
+abstract class SyntheticPropertyAccessorReference(
+    expression: KtNameReferenceExpression,
+    val getter: Boolean
+) : KtSimpleReference<KtNameReferenceExpression>(expression) {
+    protected abstract fun additionalIsReferenceToChecker(element: PsiElement): Boolean
+
+    protected fun isAccessorName(name: String): Boolean {
+        if (getter) {
+            return name.startsWith("get") || name.startsWith("is")
+        }
+        return name.startsWith("set")
+    }
+
+    override fun getRangeInElement() = TextRange(0, expression.textLength)
+
+    override fun canRename() = true
+
+    abstract override fun handleElementRename(newElementName: String): PsiElement?
+
+    override val resolvesByNames: Collection<Name>
+        get() = listOf(element.getReferencedNameAsName())
+}

compiler/psi/src/org/jetbrains/kotlin/idea/references/referenceUtils.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtReferenceExpression
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+
+val KtSimpleNameExpression.mainReference: KtSimpleNameReference
+    get() = references.firstIsInstance()
+
+val KtReferenceExpression.mainReference: KtReference
+    get() = if (this is KtSimpleNameExpression) mainReference else references.firstIsInstance()
+
+val KDocName.mainReference: KDocReference
+    get() = references.firstIsInstance()
+
+val KtElement.mainReference: KtReference?
+    get() = when (this) {
+        is KtReferenceExpression -> mainReference
+        is KDocName -> mainReference
+        else -> references.firstIsInstanceOrNull()
+    }