FIR IDE: use symbols for reference resolve

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/highlighter/visitors/FunctionCallHighlightingVisitor.kt

@@ -10,10 +10,10 @@ import com.intellij.openapi.editor.colors.TextAttributesKey
 import org.jetbrains.kotlin.idea.frontend.api.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtAnonymousFunctionSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtConstructorSymbol
-import org.jetbrains.kotlin.idea.refactoring.fqName.getKotlinFqName
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolKind
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.isExtensionDeclaration
 import org.jetbrains.kotlin.serialization.deserialization.KOTLIN_SUSPEND_BUILT_IN_FUNCTION_FQ_NAME
 import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors as Colors
 
@@ -30,7 +30,7 @@ internal class FunctionCallHighlightingVisitor(
                 analysisSession.resolveCall(expression)
                     ?.takeIf { callInfo ->
                         // ignore arithmetic-like operator calls
-                        (callInfo.targetFunction as? KtNamedFunction)?.hasModifier(KtTokens.OPERATOR_KEYWORD) != true
+                        (callInfo.targetFunction as? KtFunctionSymbol)?.isOperator != true
                     }
                     ?.let { callInfo ->
                         getTextAttributesForCal(callInfo)?.let { attributes ->

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtSymbolBasedReference.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2020 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.frontend.api
+
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+import org.jetbrains.kotlin.idea.references.KtReference
+import org.jetbrains.kotlin.idea.references.KtSimpleReference
+
+interface KtSymbolBasedReference : KtReference {
+    fun resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol>
+}
+
+fun KtReference.resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol> {
+    check(this is KtSymbolBasedReference) { "To get reference symbol the one should be KtSymbolBasedReference" }
+    return resolveToSymbols(analysisSession)
+}
+
+fun KtSimpleReference<*>.resolveToSymbol(analysisSession: FrontendAnalysisSession): KtSymbol? {
+    check(this is KtSymbolBasedReference) { "To get reference symbol the one should be KtSymbolBasedReference" }
+    return resolveToSymbols(analysisSession).singleOrNull()
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirKtReference.kt

@@ -6,13 +6,13 @@
 package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.idea.frontend.api.fir.FirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.FrontendAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 
 interface FirKtReference : KtReference {
-    fun resolveToSymbols(analysisSession: FirAnalysisSession): Collection<KtSymbol>
+    fun resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol>
 
-    fun getResolvedToPsi(analysisSession: FirAnalysisSession): Collection<PsiElement> =
+    fun getResolvedToPsi(analysisSession: FrontendAnalysisSession): Collection<PsiElement> =
         resolveToSymbols(analysisSession).mapNotNull(KtSymbol::psi)
 
     override val resolver get() = KtFirReferenceResolver

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
@@ -19,34 +18,42 @@ import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 
-object FirReferenceResolveHelper {
-    fun FirResolvedTypeRef.toTargetPsi(session: FirSession): PsiElement? {
+internal object FirReferenceResolveHelper {
+    fun FirResolvedTypeRef.toTargetSymbol(session: FirSession, symbolBuilder: KtSymbolByFirBuilder): KtSymbol? {
         val type = type as? ConeLookupTagBasedType ?: return null
-        return (type.lookupTag.toSymbol(session) as? AbstractFirBasedSymbol<*>)?.fir?.findPsi(session)
+        val symbol = type.lookupTag.toSymbol(session) as? AbstractFirBasedSymbol<*>
+        return symbol?.fir?.buildSymbol(symbolBuilder)
     }
 
-    fun ClassId.toTargetPsi(session: FirSession, calleeReference: FirReference? = null): PsiElement? {
+    private fun ClassId.toTargetPsi(
+        session: FirSession,
+        symbolBuilder: KtSymbolByFirBuilder,
+        calleeReference: FirReference? = null
+    ): KtSymbol? {
         val classLikeDeclaration = ConeClassLikeLookupTagImpl(this).toSymbol(session)?.fir
         if (classLikeDeclaration is FirRegularClass) {
             if (calleeReference is FirResolvedNamedReference) {
                 val callee = calleeReference.resolvedSymbol.fir as? FirCallableMemberDeclaration
                 // TODO: check callee owner directly?
                 if (callee !is FirConstructor && callee?.isStatic != true) {
-                    classLikeDeclaration.companionObject?.let { return it.findPsi(session) }
+                    classLikeDeclaration.companionObject?.let { return it.buildSymbol(symbolBuilder) }
                 }
             }
         }
-        return classLikeDeclaration?.findPsi(session)
+        return classLikeDeclaration?.buildSymbol(symbolBuilder)
     }
 
-    fun FirReference.toTargetPsi(session: FirSession): PsiElement? {
+    fun FirReference.toTargetSymbol(session: FirSession, symbolBuilder: KtSymbolByFirBuilder): KtSymbol? {
         return when (this) {
             is FirResolvedNamedReference -> {
-                val targetFir = when (val symbol = resolvedSymbol) {
+                val fir = when (val symbol = resolvedSymbol) {
                     is SyntheticPropertySymbol -> {
                         val syntheticProperty = symbol.fir as FirSyntheticProperty
                         if (syntheticProperty.getter.delegate.symbol.callableId == symbol.accessorId) {
@@ -55,18 +62,18 @@ object FirReferenceResolveHelper {
                             syntheticProperty.setter!!.delegate
                         }
                     }
-                    else -> symbol.fir
+                    else -> symbol.fir as? FirDeclaration
                 }
-                targetFir.findPsi(session)
+                fir?.buildSymbol(symbolBuilder)
             }
             is FirResolvedCallableReference -> {
-                resolvedSymbol.fir.findPsi(session)
+                resolvedSymbol.fir.buildSymbol(symbolBuilder)
             }
             is FirThisReference -> {
-                boundSymbol?.fir?.findPsi(session)
+                boundSymbol?.fir?.buildSymbol(symbolBuilder)
             }
             is FirSuperReference -> {
-                (superTypeRef as? FirResolvedTypeRef)?.toTargetPsi(session)
+                (superTypeRef as? FirResolvedTypeRef)?.toTargetSymbol(session, symbolBuilder)
             }
             else -> {
                 null
@@ -74,7 +81,10 @@ object FirReferenceResolveHelper {
         }
     }
 
-    internal fun resolveSimpleNameReference(ref: KtSimpleNameReferenceFirImpl): Collection<PsiElement> {
+    internal fun resolveSimpleNameReference(
+        ref: KtSimpleNameReferenceFirImpl,
+        symbolBuilder: KtSymbolByFirBuilder
+    ): Collection<KtSymbol> {
         val expression = ref.expression
         val fir = expression.getOrBuildFir()
         val session = expression.session
@@ -91,10 +101,10 @@ object FirReferenceResolveHelper {
                         // }
                         (fir.dispatchReceiver as FirQualifiedAccessExpression).calleeReference
                     } else fir.calleeReference
-                return listOfNotNull(calleeReference.toTargetPsi(session))
+                return listOfNotNull(calleeReference.toTargetSymbol(session, symbolBuilder))
             }
             is FirResolvedTypeRef -> {
-                return listOfNotNull(fir.toTargetPsi(session))
+                return listOfNotNull(fir.toTargetSymbol(session, symbolBuilder))
             }
             is FirResolvedQualifier -> {
                 val classId = fir.classId ?: return emptyList()
@@ -109,15 +119,15 @@ object FirReferenceResolveHelper {
                     }
                     val parentFir = selectorExpression.getOrBuildFir()
                     if (parentFir is FirQualifiedAccess) {
-                        return listOfNotNull(classId.toTargetPsi(session, parentFir.calleeReference))
+                        return listOfNotNull(classId.toTargetPsi(session, symbolBuilder, parentFir.calleeReference))
                     }
                     parent = parent.parent as? KtDotQualifiedExpression
                 }
-                return listOfNotNull(classId.toTargetPsi(session))
+                return listOfNotNull(classId.toTargetPsi(session, symbolBuilder))
             }
             is FirAnnotationCall -> {
                 val type = fir.typeRef as? FirResolvedTypeRef ?: return emptyList()
-                return listOfNotNull(type.toTargetPsi(session))
+                return listOfNotNull(type.toTargetSymbol(session, symbolBuilder))
             }
             is FirResolvedImport -> {
                 var parent = expression.parent
@@ -126,30 +136,36 @@ object FirReferenceResolveHelper {
                         // Special: package reference in the middle of import directive
                         // import a.<caret>b.c.SomeClass
                         // TODO: return reference to PsiPackage
-                        return listOf(expression)
+//                        return listOf(expression)
+                        return emptyList()
                     }
                     parent = parent.parent
                 }
                 val classId = fir.resolvedClassId
                 if (classId != null) {
-                    return listOfNotNull(classId.toTargetPsi(session))
+                    return listOfNotNull(classId.toTargetPsi(session, symbolBuilder))
                 }
                 val name = fir.importedName ?: return emptyList()
                 val symbolProvider = session.firSymbolProvider
-                return symbolProvider.getTopLevelCallableSymbols(fir.packageFqName, name)
-                    .mapNotNull { it.fir.findPsi(session) } +
-                        listOfNotNull(
-                            symbolProvider
-                                .getClassLikeSymbolByFqName(ClassId(fir.packageFqName, name))
-                                ?.fir?.findPsi(session)
-                        )
+
+                @OptIn(ExperimentalStdlibApi::class)
+                return buildList {
+                    symbolProvider.getTopLevelCallableSymbols(fir.packageFqName, name)
+                        .mapTo(this) { it.fir.buildSymbol(symbolBuilder) }
+                    symbolProvider
+                        .getClassLikeSymbolByFqName(ClassId(fir.packageFqName, name))
+                        ?.fir
+                        ?.buildSymbol(symbolBuilder)
+                        ?.let(::add)
+                }
             }
             is FirFile -> {
                 if (expression.getNonStrictParentOfType<KtPackageDirective>() != null) {
                     // Special: package reference in the middle of package directive
-                    return listOf(expression)
+                    return emptyList()
+//                    return listOf(expression)
                 }
-                return listOfNotNull(fir.findPsi(session))
+                return listOf(symbolBuilder.buildSymbol(fir))
             }
             is FirArrayOfCall -> {
                 // We can't yet find PsiElement for arrayOf, intArrayOf, etc.
@@ -157,7 +173,7 @@ object FirReferenceResolveHelper {
             }
             is FirReturnExpression -> {
                 return if (expression is KtLabelReferenceExpression) {
-                    listOfNotNull(fir.target.labeledElement.findPsi(session))
+                    listOf(fir.target.labeledElement.buildSymbol(symbolBuilder))
                 } else emptyList()
             }
             is FirErrorNamedReference -> {
@@ -182,7 +198,7 @@ object FirReferenceResolveHelper {
                             unresolvedCounter--
                             classId = classId?.outerClassId
                         }
-                        return listOfNotNull(classId?.toTargetPsi(session))
+                        return listOfNotNull(classId?.toTargetPsi(session, symbolBuilder))
                     }
                     parent = parent.parent as? KtDotQualifiedExpression
                     unresolvedCounter++

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt

@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.frontend.api.FrontendAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.FirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtImportAlias
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -14,7 +16,10 @@ import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 internal class KtSimpleNameReferenceFirImpl(
     expression: KtSimpleNameExpression
 ) : KtSimpleNameReference(expression), FirKtReference {
-    override fun getResolvedToPsi(analysisSession: FirAnalysisSession) = FirReferenceResolveHelper.resolveSimpleNameReference(this)
+    override fun resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol> {
+        check(analysisSession is FirAnalysisSession)
+        return FirReferenceResolveHelper.resolveSimpleNameReference(this, analysisSession.firSymbolBuilder)
+    }
 
     override fun doCanBeReferenceTo(candidateTarget: PsiElement): Boolean {
         return true // TODO

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtDestructuringDeclarationReferenceFirImpl.kt

@@ -5,11 +5,14 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.expressions.FirComponentCall
-import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.fir.getCalleeSymbol
+import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.frontend.api.FrontendAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.FirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.psi.KtDestructuringDeclarationEntry
 
 class KtDestructuringDeclarationReferenceFirImpl(
@@ -17,11 +20,10 @@ class KtDestructuringDeclarationReferenceFirImpl(
 ) : KtDestructuringDeclarationReference(element), FirKtReference {
     override fun canRename(): Boolean = false //todo
 
-    override fun getResolvedToPsi(
-        analysisSession: FirAnalysisSession
-    ): Collection<PsiElement> {
+    override fun resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol> {
+        check(analysisSession is FirAnalysisSession)
         val fir = expression.getOrBuildFirSafe<FirProperty>() ?: return emptyList()
         val componentFunctionSymbol = (fir.initializer as? FirComponentCall)?.getCalleeSymbol() ?: return emptyList()
-        return listOfNotNull(componentFunctionSymbol.fir.findPsi(element.project))
+        return listOfNotNull(componentFunctionSymbol.fir.buildSymbol(analysisSession.firSymbolBuilder))
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtForLoopInReferenceFirImpl.kt

@@ -5,21 +5,22 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
-import org.jetbrains.kotlin.idea.fir.findPsi
 import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.fir.getResolvedSymbolOfNameReference
+import org.jetbrains.kotlin.idea.frontend.api.FrontendAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.FirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.psi.KtForExpression
 
 open class KtForLoopInReferenceFirImpl(expression: KtForExpression) : KtForLoopInReference(expression), FirKtReference {
-    override fun getResolvedToPsi(
-        analysisSession: FirAnalysisSession
-    ): Collection<PsiElement> {
+
+    override fun resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol> {
+        check(analysisSession is FirAnalysisSession)
         val firLoop = expression.getOrBuildFirSafe<FirWhileLoop>() ?: return emptyList()
         val condition = firLoop.condition as? FirFunctionCall
         val iterator = run {
@@ -28,11 +29,10 @@ open class KtForLoopInReferenceFirImpl(expression: KtForExpression) : KtForLoopI
         }
         val hasNext = condition?.calleeReference?.getResolvedSymbolOfNameReference()
         val next = (firLoop.block.statements.firstOrNull() as? FirProperty?)?.getInitializerFunctionCall()
-        val project = element.project
         return listOfNotNull(
-            iterator?.fir?.findPsi(project),
-            hasNext?.fir?.findPsi(project),
-            next?.fir?.findPsi(project),
+            iterator?.fir?.buildSymbol(analysisSession.firSymbolBuilder),
+            hasNext?.fir?.buildSymbol(analysisSession.firSymbolBuilder),
+            next?.fir?.buildSymbol(analysisSession.firSymbolBuilder),
         )
     }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtInvokeFunctionReferenceFirImpl.kt

@@ -5,9 +5,10 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.frontend.api.FrontendAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.VariableAsFunctionLikeCallInfo
 import org.jetbrains.kotlin.idea.frontend.api.fir.FirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtExpression
 
@@ -16,7 +17,7 @@ class KtInvokeFunctionReferenceFirImpl(expression: KtCallExpression) : KtInvokeF
         TODO("Not yet implemented")
     }
 
-    override fun getResolvedToPsi(analysisSession: FirAnalysisSession): Collection<PsiElement> {
+    override fun resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol> {
         val call = analysisSession.resolveCall(expression) ?: return emptyList()
         if (call is VariableAsFunctionLikeCallInfo) {
             return listOf(call.invokeFunction)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtPropertyDelegationMethodsReferenceFirImpl.kt

@@ -5,31 +5,31 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.frontend.api.FrontendAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.FirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtPropertyDelegate
 
 class KtPropertyDelegationMethodsReferenceFirImpl(
     element: KtPropertyDelegate
 ) : KtPropertyDelegationMethodsReference(element), FirKtReference {
-    override fun getResolvedToPsi(
-        analysisSession: FirAnalysisSession
-    ): Collection<PsiElement> {
+    override fun resolveToSymbols(analysisSession: FrontendAnalysisSession): Collection<KtSymbol> {
+        check(analysisSession is FirAnalysisSession)
         val property = (expression.parent as? KtElement)?.getOrBuildFirSafe<FirProperty>() ?: return emptyList()
         if (property.delegate == null) return emptyList()
         val getValueSymbol = (property.getter?.singleStatementOfType<FirReturnExpression>()?.result as? FirFunctionCall)?.getCalleeSymbol()
         val setValueSymbol = property.setter?.singleStatementOfType<FirFunctionCall>()?.getCalleeSymbol()
-        val project = element.project
         return listOfNotNull(
-            getValueSymbol?.fir?.findPsi(project),
-            setValueSymbol?.fir?.findPsi(project)
+            getValueSymbol?.fir?.buildSymbol(analysisSession.firSymbolBuilder),
+            setValueSymbol?.fir?.buildSymbol(analysisSession.firSymbolBuilder)
         )
     }