FIR IDE: always try to find symbol origin in overridden symbols

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

@@ -9,9 +9,11 @@ import com.google.common.collect.MapMaker
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
-import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.firProvider
+import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
@@ -32,10 +34,12 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirFlexibleType
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirIntersectionType
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirTypeArgumentWithVariance
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirTypeParameterType
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.threadLocal
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.stubindex.PackageIndexUtil
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import java.util.concurrent.ConcurrentMap
 
@@ -43,8 +47,6 @@ import java.util.concurrent.ConcurrentMap
  * Maps FirElement to KtSymbol & ConeType to KtType, thread safe
  */
 internal class KtSymbolByFirBuilder private constructor(
-    firProvider: FirSymbolProvider,
-    typeCheckerContext: ConeTypeCheckerContext,
     private val project: Project,
     resolveState: FirModuleResolveState,
     override val token: ValidityToken,
@@ -52,16 +54,21 @@ internal class KtSymbolByFirBuilder private constructor(
     private val symbolsCache: BuilderCache<FirDeclaration, KtSymbol>,
     private val typesCache: BuilderCache<ConeKotlinType, KtType>
 ) : ValidityTokenOwner {
+    private val typeCheckerContext by threadLocal {
+        ConeTypeCheckerContext(
+            isErrorTypeEqualsToAnything = true,
+            isStubTypeEqualsToAnything = true,
+            resolveState.firIdeLibrariesSession
+        )
+    }
+
+    private val firProvider get() = resolveState.firIdeSourcesSession.firSymbolProvider
 
     constructor(
-        firProvider: FirSymbolProvider,
-        typeCheckerContext: ConeTypeCheckerContext,
         resolveState: FirModuleResolveState,
         project: Project,
         token: ValidityToken
     ) : this(
-        firProvider = firProvider,
-        typeCheckerContext = typeCheckerContext,
         project = project,
         token = token,
         resolveState = resolveState,
@@ -70,15 +77,11 @@ internal class KtSymbolByFirBuilder private constructor(
         typesCache = BuilderCache()
     )
 
-    private val firProvider by weakRef(firProvider)
     private val resolveState by weakRef(resolveState)
-    private val typeCheckerContext by weakRef(typeCheckerContext)
 
     fun createReadOnlyCopy(newResolveState: FirModuleResolveState): KtSymbolByFirBuilder {
         check(!withReadOnlyCaching) { "Cannot create readOnly KtSymbolByFirBuilder from a readonly one" }
         return KtSymbolByFirBuilder(
-            firProvider,
-            typeCheckerContext,
             project,
             token = token,
             resolveState = newResolveState,
@@ -126,8 +129,8 @@ internal class KtSymbolByFirBuilder private constructor(
     fun buildFirConstructorParameter(fir: FirValueParameterImpl) =
         symbolsCache.cache(fir) { KtFirConstructorValueParameterSymbol(fir, resolveState, token, this) }
 
-    fun buildFunctionSymbol(fir: FirSimpleFunction, forcedOrigin: FirDeclarationOrigin? = null) = symbolsCache.cache(fir) {
+    fun buildFunctionSymbol(fir: FirSimpleFunction) = symbolsCache.cache(fir) {
-        KtFirFunctionSymbol(fir, resolveState, token, this, forcedOrigin)
+        KtFirFunctionSymbol(fir, resolveState, token, this)
     }
 
     fun buildConstructorSymbol(fir: FirConstructor) = symbolsCache.cache(fir) { KtFirConstructorSymbol(fir, resolveState, token, this) }
@@ -139,10 +142,10 @@ internal class KtSymbolByFirBuilder private constructor(
     fun buildAnonymousFunctionSymbol(fir: FirAnonymousFunction) =
         symbolsCache.cache(fir) { KtFirAnonymousFunctionSymbol(fir, resolveState, token, this) }
 
-    fun buildVariableSymbol(fir: FirProperty, forcedOrigin: FirDeclarationOrigin? = null): KtVariableSymbol = symbolsCache.cache(fir) {
+    fun buildVariableSymbol(fir: FirProperty): KtVariableSymbol = symbolsCache.cache(fir) {
         when {
             fir.isLocal -> KtFirLocalVariableSymbol(fir, resolveState, token, this)
-            else -> KtFirPropertySymbol(fir, resolveState, token, this, forcedOrigin)
+            else -> KtFirPropertySymbol(fir, resolveState, token, this)
         }
     }
 
@@ -154,6 +157,10 @@ internal class KtSymbolByFirBuilder private constructor(
         (firProvider.getSymbolByLookupTag(lookupTag) as? FirTypeParameterSymbol)?.fir?.let(::buildTypeParameterSymbol)
     }
 
+    fun buildClassLikeSymbolByClassId(classId: ClassId, firSession: FirSession): FirRegularClass? = withValidityAssertion {
+        firSession.firProvider.getClassLikeSymbolByFqName(classId)?.fir as? FirRegularClass
+    }
+
 
     fun createPackageSymbolIfOneExists(packageFqName: FqName): KtFirPackageSymbol? {
         val exists = PackageIndexUtil.packageExists(packageFqName, GlobalSearchScope.allScope(project), project)

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

@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScope
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassLikeSymbol
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.name.Name
 
@@ -60,27 +59,18 @@ internal abstract class KtFirDelegatingScope<S>(
 
 }
 
-private fun FirBasedSymbol<*>.realDeclarationOrigin(): FirDeclarationOrigin? {
-    if (this !is PossiblyFirFakeOverrideSymbol<*, *> || !isFakeOverride) return null
-    var current: FirBasedSymbol<*>? = this.overriddenSymbol
-    while (current is PossiblyFirFakeOverrideSymbol<*, *> && current.isFakeOverride) {
-        current = current.overriddenSymbol
-    }
-    return (current?.fir as? FirDeclaration)?.origin
-}
-
 internal fun FirScope.getCallableSymbols(callableNames: Collection<Name>, builder: KtSymbolByFirBuilder) = sequence {
     callableNames.forEach { name ->
         val callables = mutableListOf<KtCallableSymbol>()
         processFunctionsByName(name) { firSymbol ->
             (firSymbol.fir as? FirSimpleFunction)?.let { fir ->
-                callables.add(builder.buildFunctionSymbol(fir, firSymbol.realDeclarationOrigin()))
+                callables.add(builder.buildFunctionSymbol(fir))
             }
         }
         processPropertiesByName(name) { firSymbol ->
             val symbol = when {
                 firSymbol is FirPropertySymbol && firSymbol.isFakeOverride -> {
-                    builder.buildVariableSymbol(firSymbol.fir, firSymbol.realDeclarationOrigin())
+                    builder.buildVariableSymbol(firSymbol.fir)
                 }
                 else -> builder.buildCallableSymbol(firSymbol.fir)
             }

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

@@ -7,30 +7,29 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRef
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.overriddenDeclaration
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolOrigin
-import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 
 internal interface KtFirSymbol<F : FirDeclaration> : KtSymbol, ValidityTokenOwner {
     val firRef: FirRef<F>
 
-    override val origin: KtSymbolOrigin get() = firRef.withFir { it.origin.asKtSymbolOrigin() }
+    override val origin: KtSymbolOrigin get() = firRef.withFir { it.ktSymbolOrigin() }
 }
 
-internal fun FirDeclarationOrigin.asKtSymbolOrigin() = when (this) {
+
+private tailrec fun FirDeclaration.ktSymbolOrigin(): KtSymbolOrigin = when (origin) {
     FirDeclarationOrigin.Source -> KtSymbolOrigin.SOURCE
     FirDeclarationOrigin.Library -> KtSymbolOrigin.LIBRARY
     FirDeclarationOrigin.Java -> KtSymbolOrigin.JAVA
     FirDeclarationOrigin.SamConstructor -> KtSymbolOrigin.SAM_CONSTRUCTOR
-    FirDeclarationOrigin.Synthetic -> throw InvalidFirDeclarationOriginForSymbol(this)
+    FirDeclarationOrigin.Enhancement -> KtSymbolOrigin.JAVA
-    FirDeclarationOrigin.FakeOverride -> throw InvalidFirDeclarationOriginForSymbol(this)
+    else -> {
-    FirDeclarationOrigin.ImportedFromObject -> throw InvalidFirDeclarationOriginForSymbol(this)
+        val overridden = overriddenDeclaration ?: throw InvalidFirDeclarationOriginForSymbol(origin)
-    FirDeclarationOrigin.IntersectionOverride -> throw InvalidFirDeclarationOriginForSymbol(this)
+        overridden.ktSymbolOrigin()
-    FirDeclarationOrigin.Enhancement -> KtSymbolOrigin.JAVA // TODO
+    }
-    is FirDeclarationOrigin.Plugin -> throw InvalidFirDeclarationOriginForSymbol(this)
 }
 
 class InvalidFirDeclarationOriginForSymbol(origin: FirDeclarationOrigin) :

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

@@ -0,0 +1,15 @@
+/*
+ * 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.utils
+
+import org.jetbrains.kotlin.fir.FirSymbolOwner
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+
+internal inline val FirDeclaration.overriddenDeclaration: FirDeclaration?
+    get() {
+        val symbol = (this as? FirSymbolOwner<*>)?.symbol ?: return null
+        return (symbol as? FirCallableSymbol)?.overriddenSymbol?.fir as? FirDeclaration
+    }