FIR IDE: introduce symbol provider

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

@@ -7,14 +7,16 @@ package org.jetbrains.kotlin.idea.frontend.api
 
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
 
 abstract class FrontendAnalysisSession(project: Project) : Invalidatable {
-    protected val validityToken = ReadActionConfinementValidityToken(project)
+    protected open val validityToken = ReadActionConfinementValidityToken(project)
     override fun isValid(): Boolean = validityToken.isValid()
     override fun invalidationReason(): String = validityToken.invalidationReason()
 
+    abstract val symbolProvider: KtSymbolProvider
 
     abstract fun getSmartCastedToTypes(expression: KtExpression): Collection<TypeInfo>?
 

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

@@ -5,10 +5,11 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir
 
-import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.diagnostics.Diagnostic
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.isSuspend
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
@@ -16,10 +17,18 @@ import org.jetbrains.kotlin.fir.resolve.transformers.firClassLike
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeCheckerContext
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 import org.jetbrains.kotlin.idea.fir.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.frontend.api.*
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.FirKtSymbolProvider
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolProvider
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtVariableLikeSymbol
 import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper
 import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper.toTargetSymbol
 import org.jetbrains.kotlin.name.ClassId
@@ -27,11 +36,17 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 
 class FirAnalysisSession(
-    project: Project
+    element: KtElement
-) : FrontendAnalysisSession(project) {
+) : FrontendAnalysisSession(element.project) {
-   constructor(element: KtElement) : this(element.project)
+    internal val token get() = validityToken
 
-    internal val symbolBuilder = KtSymbolByFirBuilder(validityToken)
+    internal val firSymbolBuilder = KtSymbolByFirBuilder(validityToken)
+    override val symbolProvider: KtSymbolProvider =
+        FirKtSymbolProvider(
+            this,
+            element.session.firSymbolProvider,
+            firSymbolBuilder
+        )
 
     init {
         assertIsValid()
@@ -112,17 +127,17 @@ class FirAnalysisSession(
 
     private fun resolveCall(firCall: FirFunctionCall, callExpression: KtExpression): CallInfo? {
         val session = callExpression.session
-        val resolvedFunctionSymbol = firCall.calleeReference.toTargetSymbol(session, symbolBuilder)
+        val resolvedFunctionSymbol = firCall.calleeReference.toTargetSymbol(session, firSymbolBuilder)
         val resolvedCalleeSymbol = (firCall.calleeReference as? FirResolvedNamedReference)?.resolvedSymbol
         return when {
             resolvedCalleeSymbol is FirConstructorSymbol -> {
                 val fir = resolvedCalleeSymbol.fir
-                FunctionCallInfo(symbolBuilder.buildFirConstructorSymbol(fir))
+                FunctionCallInfo(firSymbolBuilder.buildConstructorSymbol(fir))
             }
             firCall.dispatchReceiver is FirQualifiedAccessExpression && firCall.isImplicitFunctionCall() -> {
                 val target = with(FirReferenceResolveHelper) {
                     val calleeReference = (firCall.dispatchReceiver as FirQualifiedAccessExpression).calleeReference
-                    calleeReference.toTargetSymbol(session, symbolBuilder)
+                    calleeReference.toTargetSymbol(session, firSymbolBuilder)
                 }
                 when (target) {
                     null -> null
@@ -132,7 +147,7 @@ class FirAnalysisSession(
                         when (functionSymbol?.callableId) {
                             null -> null
                             in kotlinFunctionInvokeCallableIds -> VariableAsFunctionCallInfo(target, functionSymbol.fir.isSuspend)
-                            else -> (resolvedFunctionSymbol as? KtSimpleFunctionSymbol)
+                            else -> (resolvedFunctionSymbol as? KtFunctionSymbol)
                                 ?.let { VariableAsFunctionLikeCallInfo(target, it) }
                         }
                     }
@@ -144,7 +159,7 @@ class FirAnalysisSession(
     }
 
     private fun KtSymbol.asSimpleFunctionCall() =
-        (this as? KtSimpleFunctionSymbol)?.let(::FunctionCallInfo)
+        (this as? KtFunctionSymbol)?.let(::FunctionCallInfo)
 
     private fun forEachSuperClass(firClass: FirClass<*>, action: (FirResolvedTypeRef) -> Unit) {
         firClass.superTypeRefs.forEach { superType ->

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

@@ -31,19 +31,19 @@ internal class KtSymbolByFirBuilder(private val validityToken: Invalidatable) {
     fun buildSymbol(fir: FirDeclaration): KtSymbol = when (fir) {
         is FirRegularClass -> buildClassSymbol(fir)
         is FirSimpleFunction -> buildFunctionSymbol(fir)
-        is FirProperty -> buildPropertySymbol(fir)
+        is FirProperty -> buildVariableSymbol(fir)
         is FirValueParameterImpl -> buildParameterSymbol(fir)
-        is FirConstructor -> buildFirConstructorSymbol(fir)
+        is FirConstructor -> buildConstructorSymbol(fir)
-        is FirTypeParameter -> buildFirTypeParameterSymbol(fir)
+        is FirTypeParameter -> buildTypeParameterSymbol(fir)
-        is FirTypeAlias -> buildFirTypeAliasSymbol(fir)
+        is FirTypeAlias -> buildTypeAliasSymbol(fir)
-        is FirEnumEntry -> buildFirEnumEntrySymbol(fir)
+        is FirEnumEntry -> buildEnumEntrySymbol(fir)
-        is FirField -> buildFirFieldSymbol(fir)
+        is FirField -> buildFieldSymbol(fir)
-        is FirAnonymousFunction -> buildFirAnonymousFunction(fir)
+        is FirAnonymousFunction -> buildAnonymousFunctionSymbol(fir)
         else ->
             TODO(fir::class.toString())
     }
 
-    fun buildClassSymbol(fir: FirRegularClass) = KtFirClassOrObjectSymbol(fir, validityToken)
+    fun buildClassSymbol(fir: FirRegularClass) = KtFirClassOrObjectSymbol(fir, validityToken, this)
 
     // TODO it can be a constructor parameter, which may be split into parameter & property
     // we should handle them both
@@ -51,14 +51,14 @@ internal class KtSymbolByFirBuilder(private val validityToken: Invalidatable) {
     fun buildFirConstructorParameter(fir: FirValueParameterImpl) = KtFirConstructorValueParameterSymbol(fir, validityToken)
 
     fun buildFunctionSymbol(fir: FirSimpleFunction) = KtFirFunctionSymbol(fir, validityToken, this)
-    fun buildFirConstructorSymbol(fir: FirConstructor) = KtFirConstructorSymbol(fir, validityToken, this)
+    fun buildConstructorSymbol(fir: FirConstructor) = KtFirConstructorSymbol(fir, validityToken, this)
-    fun buildFirTypeParameterSymbol(fir: FirTypeParameter) = KtFirTypeParameterSymbol(fir, validityToken)
+    fun buildTypeParameterSymbol(fir: FirTypeParameter) = KtFirTypeParameterSymbol(fir, validityToken)
-    fun buildFirTypeAliasSymbol(fir: FirTypeAlias) = KtFirTypeAliasSymbol(fir, validityToken)
+    fun buildTypeAliasSymbol(fir: FirTypeAlias) = KtFirTypeAliasSymbol(fir, validityToken)
-    fun buildFirEnumEntrySymbol(fir: FirEnumEntry) = KtFirEnumEntrySymbol(fir, validityToken)
+    fun buildEnumEntrySymbol(fir: FirEnumEntry) = KtFirEnumEntrySymbol(fir, validityToken)
-    fun buildFirFieldSymbol(fir: FirField) = KtFirFieldSymbol(fir, validityToken)
+    fun buildFieldSymbol(fir: FirField) = KtFirFieldSymbol(fir, validityToken)
-    fun buildFirAnonymousFunction(fir: FirAnonymousFunction) = KtFirAnonymousFunctionSymbol(fir, validityToken, this)
+    fun buildAnonymousFunctionSymbol(fir: FirAnonymousFunction) = KtFirAnonymousFunctionSymbol(fir, validityToken, this)
 
-    fun buildPropertySymbol(fir: FirProperty): KtVariableSymbol {
+    fun buildVariableSymbol(fir: FirProperty): KtVariableSymbol {
         return when {
             fir.isLocal -> KtFirLocalVariableSymbol(fir, validityToken)
             else -> KtFirPropertySymbol(fir, validityToken)

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

@@ -0,0 +1,71 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.frontend.api.Invalidatable
+import org.jetbrains.kotlin.idea.frontend.api.InvalidatableByValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
+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.withValidityAssertion
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.*
+
+internal class FirKtSymbolProvider(
+    override val token: Invalidatable,
+    firSymbolProvider: FirSymbolProvider,
+    private val firSymbolBuilder: KtSymbolByFirBuilder
+) : KtSymbolProvider(), InvalidatableByValidityToken {
+    private val firSymbolProvider by weakRef(firSymbolProvider)
+
+    override fun getParameterSymbol(psi: KtParameter): KtParameterSymbol = withValidityAssertion {
+        firSymbolBuilder.buildParameterSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getFunctionSymbol(psi: KtNamedFunction): KtFunctionSymbol = withValidityAssertion {
+        firSymbolBuilder.buildFunctionSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol = withValidityAssertion {
+        firSymbolBuilder.buildConstructorSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getTypeParameterSymbol(psi: KtTypeParameter): KtTypeParameterSymbol = withValidityAssertion {
+        firSymbolBuilder.buildTypeParameterSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol = withValidityAssertion {
+        firSymbolBuilder.buildTypeAliasSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol = withValidityAssertion {
+        firSymbolBuilder.buildEnumEntrySymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol = withValidityAssertion {
+        firSymbolBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getAnonymousFunctionSymbol(psi: KtLambdaExpression): KtAnonymousFunctionSymbol = withValidityAssertion {
+        firSymbolBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getVariableSymbol(psi: KtProperty): KtVariableSymbol = withValidityAssertion {
+        firSymbolBuilder.buildVariableSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol = withValidityAssertion {
+        firSymbolBuilder.buildClassSymbol(psi.getOrBuildFirOfType())
+    }
+
+    override fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol? = withValidityAssertion {
+        val symbol = firSymbolProvider.getClassLikeSymbolByFqName(classId) as? FirRegularClassSymbol ?: return null
+        firSymbolBuilder.buildClassSymbol(symbol.fir)
+    }
+}