FIR IDE: fix symbol pointer creation for non direct member symbols

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSession.kt

@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForCompletion
@@ -80,6 +82,9 @@ private constructor(
         )
     }
 
+    val rootModuleSession: FirSession get() = firResolveState.rootModuleSession
+    val firSymbolProvider: FirSymbolProvider get() = rootModuleSession.symbolProvider
+
     companion object {
         @Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
         internal fun createForElement(element: KtElement): KtFirAnalysisSession {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirEnumEntrySymbol.kt

@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.KtFirEnumEntrySymbolPointer
-import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.createSignature
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
@@ -46,6 +45,11 @@ internal class KtFirEnumEntrySymbol(
 
     override fun createPointer(): KtSymbolPointer<KtEnumEntrySymbol> {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
-        return KtFirEnumEntrySymbolPointer(containingEnumClassIdIfNonLocal!!, firRef.withFir { it.createSignature() })
+        return firRef.withFir { fir ->
+            KtFirEnumEntrySymbolPointer(
+                fir.symbol.containingClass()?.classId ?: error("Containing class should present for enum entry"),
+                fir.name
+            )
+        }
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirFunctionSymbol.kt

@@ -16,6 +16,8 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.annotations.containsAn
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.annotations.getAnnotationClassIds
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.annotations.toAnnotationsList
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.KtFirMemberFunctionSymbolPointer
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.KtFirMemberPropertySymbolPointer
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.KtFirTopLevelFunctionSymbolPointer
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.createSignature
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
@@ -92,11 +94,16 @@ internal class KtFirFunctionSymbol(
     override fun createPointer(): KtSymbolPointer<KtFunctionSymbol> {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         return when (symbolKind) {
-            KtSymbolKind.TOP_LEVEL -> TODO("Creating symbol for top level fun is not supported yet")
+            KtSymbolKind.TOP_LEVEL -> firRef.withFir { fir ->
-            KtSymbolKind.MEMBER -> KtFirMemberFunctionSymbolPointer(
+                KtFirTopLevelFunctionSymbolPointer(fir.symbol.callableId, fir.createSignature())
-                firRef.withFir { it.containingClass()?.classId ?: error("ClassId should not be null for member function") },
+            }
-                firRef.withFir { it.createSignature() }
+            KtSymbolKind.MEMBER ->  firRef.withFir { fir ->
-            )
+                KtFirMemberFunctionSymbolPointer(
+                    fir.containingClass()?.classId ?: error("ClassId should not be null for member function"),
+                    fir.name,
+                    fir.createSignature()
+                )
+            }
             KtSymbolKind.NON_PROPERTY_PARAMETER -> error("KtFunction could not be a parameter")
             KtSymbolKind.LOCAL -> throw CanNotCreateSymbolPointerForLocalLibraryDeclarationException(
                 callableIdIfNonLocal?.asString() ?: name.asString()

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirKotlinPropertySymbol.kt

@@ -109,12 +109,15 @@ internal class KtFirKotlinPropertySymbol(
     override fun createPointer(): KtSymbolPointer<KtKotlinPropertySymbol> {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         return when (symbolKind) {
-            KtSymbolKind.TOP_LEVEL -> TODO("Creating symbol for top level fun is not supported yet")
+            KtSymbolKind.TOP_LEVEL -> TODO("Creating symbol for top level properties is not supported yet")
             KtSymbolKind.NON_PROPERTY_PARAMETER -> TODO("Creating symbol for top level parameters is not supported yet")
-            KtSymbolKind.MEMBER -> KtFirMemberPropertySymbolPointer(
+            KtSymbolKind.MEMBER -> firRef.withFir { fir ->
-                firRef.withFir { it.containingClass()?.classId ?: error("ClassId should not be null for member property") },
+                KtFirMemberPropertySymbolPointer(
-                firRef.withFir { it.createSignature() }
+                    fir.containingClass()?.classId ?: error("ClassId should not be null for member property"),
-            )
+                    fir.name,
+                    fir.createSignature()
+                )
+            }
             KtSymbolKind.LOCAL -> throw CanNotCreateSymbolPointerForLocalLibraryDeclarationException(name.asString())
         }
     }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/KtFirConstructorSymbolPointer.kt

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtConstructorSymbol
 import org.jetbrains.kotlin.ir.util.IdSignature
@@ -19,10 +19,12 @@ internal class KtFirConstructorSymbolPointer(
     private val signature: IdSignature
 ) : KtFirMemberSymbolPointer<KtConstructorSymbol>(ownerClassId) {
     override fun KtFirAnalysisSession.chooseCandidateAndCreateSymbol(
-        candidates: Collection<FirDeclaration>,
+        candidates: FirScope,
         firSession: FirSession
     ): KtConstructorSymbol? {
-        val firConstructor = candidates.findDeclarationWithSignature<FirConstructor>(signature, firSession) ?: return null
+        val firConstructor =
+            candidates.findDeclarationWithSignature<FirConstructor>(signature, firSession) { processDeclaredConstructors(it) }
+                ?: return null
         if (firConstructor.isPrimary != isPrimary) return null
         return firSymbolBuilder.functionLikeBuilder.buildConstructorSymbol(firConstructor)
     }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/KtFirEnumEntrySymbolPointer.kt

@@ -5,24 +5,43 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers
 
+import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.pointers.KtSymbolPointer
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
 
 internal class KtFirEnumEntrySymbolPointer(
-    ownerClassId: ClassId,
+    private val ownerClassId: ClassId,
-    private val signature: IdSignature
+    private val name: Name
-) : KtFirMemberSymbolPointer<KtEnumEntrySymbol>(ownerClassId) {
+) : KtSymbolPointer<KtEnumEntrySymbol>() {
-    override fun KtFirAnalysisSession.chooseCandidateAndCreateSymbol(
+    override fun restoreSymbol(analysisSession: KtAnalysisSession): KtEnumEntrySymbol? {
-        candidates: Collection<FirDeclaration>,
+        require(analysisSession is KtFirAnalysisSession)
-        firSession: FirSession
+        val enumClass = getEnumClass(analysisSession, ownerClassId)
-    ): KtEnumEntrySymbol? {
+            ?: return null
-        val firProperty = candidates.findDeclarationWithSignature<FirEnumEntry>(signature, firSession) ?: return null
+        val enumEntry = enumClass.enumEntryByName(name)
-        return firSymbolBuilder.buildEnumEntrySymbol(firProperty)
+            ?: return null
+        return analysisSession.firSymbolBuilder.buildEnumEntrySymbol(enumEntry)
     }
+
+    private fun getEnumClass(analysisSession: KtFirAnalysisSession, classId: ClassId): FirRegularClass? {
+        val enumClass = analysisSession.firSymbolProvider.getClassLikeSymbolByFqName(classId)?.fir as? FirRegularClass
+            ?: return null
+        if (enumClass.classKind != ClassKind.ENUM_CLASS) return null
+        return enumClass
+    }
+
+    private fun FirRegularClass.enumEntryByName(name: Name): FirEnumEntry? =
+        declarations.firstOrNull { member ->
+            member is FirEnumEntry && member.name == name
+        } as FirEnumEntry?
 }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/KtFirMemberFunctionSymbolPointer.kt

@@ -6,22 +6,26 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
 
 internal class KtFirMemberFunctionSymbolPointer(
     ownerClassId: ClassId,
+    private val name: Name,
     private val signature: IdSignature
 ) : KtFirMemberSymbolPointer<KtFunctionSymbol>(ownerClassId) {
     override fun KtFirAnalysisSession.chooseCandidateAndCreateSymbol(
-        candidates: Collection<FirDeclaration>,
+        candidates: FirScope,
         firSession: FirSession
     ): KtFunctionSymbol? {
-        val firFunction = candidates.findDeclarationWithSignature<FirSimpleFunction>(signature, firSession) ?: return null
+        val firFunction = candidates.findDeclarationWithSignature<FirSimpleFunction>(signature, firSession) {
+            processFunctionsByName(name, it)
+        } ?: return null
         return firSymbolBuilder.functionLikeBuilder.buildFunctionSymbol(firFunction)
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/KtFirMemberPropertySymbolPointer.kt

@@ -6,22 +6,26 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
 import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtKotlinPropertySymbol
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
 
 internal class KtFirMemberPropertySymbolPointer(
     ownerClassId: ClassId,
+    private val name: Name,
     private val signature: IdSignature
 ) : KtFirMemberSymbolPointer<KtKotlinPropertySymbol>(ownerClassId) {
     override fun KtFirAnalysisSession.chooseCandidateAndCreateSymbol(
-        candidates: Collection<FirDeclaration>,
+        candidates: FirScope,
         firSession: FirSession
     ): KtKotlinPropertySymbol? {
-        val firProperty = candidates.findDeclarationWithSignature<FirProperty>(signature, firSession) ?: return null
+        val firProperty = candidates.findDeclarationWithSignature<FirProperty>(signature, firSession) { processPropertiesByName(name, it) }
+            ?: return null
         return firSymbolBuilder.variableLikeBuilder.buildVariableSymbol(firProperty) as? KtKotlinPropertySymbol
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/KtFirMemberSymbolPointer.kt

@@ -6,9 +6,11 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
@@ -22,11 +24,16 @@ internal abstract class KtFirMemberSymbolPointer<S : KtSymbol>(
         require(analysisSession is KtFirAnalysisSession)
         val owner = analysisSession.getClassLikeSymbol(ownerClassId) as? FirRegularClass
             ?: return null
-        return analysisSession.chooseCandidateAndCreateSymbol(owner.declarations, owner.session)
+        val scope = owner.unsubstitutedScope(
+            analysisSession.firResolveState.rootModuleSession,
+            ScopeSession(),
+            withForcedTypeCalculator = false
+        )
+        return analysisSession.chooseCandidateAndCreateSymbol(scope, owner.session)
     }
 
     protected abstract fun KtFirAnalysisSession.chooseCandidateAndCreateSymbol(
-        candidates: Collection<FirDeclaration>,
+        candidates: FirScope,
         firSession: FirSession
     ): S?
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/KtFirTopLevelCallableSymbolPointer.kt

@@ -0,0 +1,36 @@
+/*
+ * 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.frontend.api.fir.symbols.pointers
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.pointers.KtSymbolPointer
+import org.jetbrains.kotlin.name.CallableId
+
+internal abstract class KtTopLevelCallableSymbolPointer<S : KtCallableSymbol>(
+    private val callableId: CallableId
+) : KtSymbolPointer<S>() {
+    final override fun restoreSymbol(analysisSession: KtAnalysisSession): S? {
+        require(analysisSession is KtFirAnalysisSession)
+        val candidates = analysisSession.getCallableSymbols(callableId)
+        if (candidates.isEmpty()) return null
+        val session = candidates.first().fir.session
+        return analysisSession.chooseCandidateAndCreateSymbol(candidates, session)
+    }
+
+    protected abstract fun KtFirAnalysisSession.chooseCandidateAndCreateSymbol(
+        candidates: Collection<FirCallableSymbol<*>>,
+        firSession: FirSession
+    ): S?
+}
+
+private fun KtFirAnalysisSession.getCallableSymbols(callableId: CallableId) =
+    firResolveState.rootModuleSession.symbolProvider.getTopLevelCallableSymbols(callableId.packageName, callableId.callableName)
+

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/KtFirTopLevelFunctionSymbolPointer.kt

@@ -0,0 +1,30 @@
+/*
+ * 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.fir.symbols.pointers
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol
+import org.jetbrains.kotlin.ir.util.IdSignature
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+
+internal class KtFirTopLevelFunctionSymbolPointer(
+    callableId: CallableId,
+    private val signature: IdSignature
+) : KtTopLevelCallableSymbolPointer<KtFunctionSymbol>(callableId) {
+    override fun KtFirAnalysisSession.chooseCandidateAndCreateSymbol(
+        candidates: Collection<FirCallableSymbol<*>>,
+        firSession: FirSession
+    ): KtFunctionSymbol? {
+        val firFunction = candidates.findDeclarationWithSignatureBySymbols<FirSimpleFunction>(signature, firSession) ?: return null
+        return firSymbolBuilder.functionLikeBuilder.buildFunctionSymbol(firFunction)
+    }
+}
+

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/pointers/pointerUtils.kt

@@ -5,17 +5,39 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers
 
+import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.renderWithType
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.ideSessionComponents
 import org.jetbrains.kotlin.ir.util.IdSignature
 
-internal inline fun <reified D : FirDeclaration> Collection<FirDeclaration>.findDeclarationWithSignature(
+internal inline fun <reified D : FirDeclaration> FirScope.findDeclarationWithSignature(
+    signature: IdSignature,
+    firSession: FirSession,
+    processor: FirScope.((AbstractFirBasedSymbol<*>) -> Unit) -> Unit
+): D? {
+    val signatureComposer = firSession.ideSessionComponents.signatureComposer
+    var foundSymbol: D? = null
+    processor { symbol ->
+        val declaration = symbol.fir
+        if (declaration is D && signatureComposer.composeSignature(declaration) == signature) {
+            foundSymbol = declaration
+        }
+    }
+    return foundSymbol
+}
+
+internal inline fun <reified D : FirDeclaration> Collection<FirCallableSymbol<*>>.findDeclarationWithSignatureBySymbols(
     signature: IdSignature,
     firSession: FirSession
 ): D? {
     val signatureComposer = firSession.ideSessionComponents.signatureComposer
-    for (declaration in this) {
+    for (symbol in this) {
+        val declaration = symbol.fir
         if (declaration is D && signatureComposer.composeSignature(declaration) == signature) {
             return declaration
         }
@@ -26,5 +48,5 @@ internal inline fun <reified D : FirDeclaration> Collection<FirDeclaration>.find
 internal fun FirDeclaration.createSignature(): IdSignature {
     val signatureComposer = session.ideSessionComponents.signatureComposer
     return signatureComposer.composeSignature(this)
-        ?: error("Could not compose signature for declaration, looks like it is private or local")
+        ?: error("Could not compose signature for ${this.renderWithType(FirRenderer.RenderMode.WithResolvePhases)}, looks like it is private or local")
 }