FIR IDE: simplify HL APIs for PsiType

The idea is to convert `KtType` to `PsiType`, along with the given
`PsiElement` context, instead of providing individual API to convert
different psi source to `PsiType`. To that end, we need to
migrate/rewrite some APIs/implementations that retrieve `KtType` from
different source of `PsiElement`.

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtExpressionTypeProvider.kt

@@ -9,10 +9,13 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFunction
 
 public abstract class KtExpressionTypeProvider : KtAnalysisSessionComponent() {
-    public abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType
     public abstract fun getKtExpressionType(expression: KtExpression): KtType
+    public abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType
+    public abstract fun getFunctionalTypeForKtFunction(declaration: KtFunction): KtType
+
     public abstract fun getExpectedType(expression: PsiElement): KtType?
     public abstract fun isDefinitelyNull(expression: KtExpression): Boolean
     public abstract fun isDefinitelyNotNull(expression: KtExpression): Boolean
@@ -22,9 +25,25 @@ public interface KtExpressionTypeProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtExpression.getKtType(): KtType =
         analysisSession.expressionTypeProvider.getKtExpressionType(this)
 
+    /**
+     * Returns the return type of the given [KtDeclaration] as [KtType]
+     */
     public fun KtDeclaration.getReturnKtType(): KtType =
         analysisSession.expressionTypeProvider.getReturnTypeForKtDeclaration(this)
 
+    /**
+     * Returns the functional type of the given [KtFunction].
+     *
+     * For a regular function, it would be kotlin.FunctionN<Ps, R> where
+     *   N is the number of value parameters in the function;
+     *   Ps are types of value parameters;
+     *   R is the return type of the function.
+     * Depending on the function's attributes, such as `suspend` or reflective access, different functional type,
+     * such as `SuspendFunction`, `KFunction`, or `KSuspendFunction`, will be constructed.
+     */
+    public fun KtFunction.getFunctionalType(): KtType =
+        analysisSession.expressionTypeProvider.getFunctionalTypeForKtFunction(this)
+
     /**
      * Returns the expected [KtType] of this [PsiElement] if it is an expression. The returned value should not be a
      * [org.jetbrains.kotlin.idea.frontend.api.types.KtClassErrorType].

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtPsiTypeProvider.kt

@@ -5,42 +5,19 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.components
 
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiType
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
-import org.jetbrains.kotlin.psi.*
 
 public abstract class KtPsiTypeProvider : KtAnalysisSessionComponent() {
-    public abstract fun commonSuperType(expressions: Collection<KtExpression>, context: KtExpression, mode: TypeMappingMode): PsiType?
-    public abstract fun getPsiTypeForKtExpression(expression: KtExpression, mode: TypeMappingMode): PsiType
-    public abstract fun getPsiTypeForKtDeclaration(ktDeclaration: KtDeclaration, mode: TypeMappingMode): PsiType
-    public abstract fun getPsiTypeForKtFunction(ktFunction: KtFunction, mode: TypeMappingMode): PsiType
-    public abstract fun getPsiTypeForKtTypeReference(ktTypeReference: KtTypeReference, mode: TypeMappingMode): PsiType
-    public abstract fun getPsiTypeForReceiverOfDoubleColonExpression(
-        ktDoubleColonExpression: KtDoubleColonExpression,
-        mode: TypeMappingMode
-    ): PsiType?
+    public abstract fun asPsiType(type: KtType, context: PsiElement, mode: TypeMappingMode): PsiType
 }
 
 public interface KtPsiTypeProviderMixIn : KtAnalysisSessionMixIn {
-    public fun commonSuperType(
-        expressions: Collection<KtExpression>,
-        context: KtExpression,
+    public fun KtType.asPsiType(
+        context: PsiElement,
         mode: TypeMappingMode = TypeMappingMode.DEFAULT
     ): PsiType? =
-        analysisSession.psiTypeProvider.commonSuperType(expressions, context, mode)
-
-    public fun KtExpression.getPsiType(mode: TypeMappingMode = TypeMappingMode.DEFAULT): PsiType =
-        analysisSession.psiTypeProvider.getPsiTypeForKtExpression(this, mode)
-
-    public fun KtDeclaration.getPsiType(mode: TypeMappingMode = TypeMappingMode.DEFAULT): PsiType =
-        analysisSession.psiTypeProvider.getPsiTypeForKtDeclaration(this, mode)
-
-    public fun KtFunction.getPsiType(mode: TypeMappingMode = TypeMappingMode.DEFAULT): PsiType =
-        analysisSession.psiTypeProvider.getPsiTypeForKtFunction(this, mode)
-
-    public fun KtTypeReference.getPsiType(mode: TypeMappingMode = TypeMappingMode.DEFAULT): PsiType =
-        analysisSession.psiTypeProvider.getPsiTypeForKtTypeReference(this, mode)
-
-    public fun KtDoubleColonExpression.getReceiverPsiType(mode: TypeMappingMode = TypeMappingMode.DEFAULT): PsiType? =
-        analysisSession.psiTypeProvider.getPsiTypeForReceiverOfDoubleColonExpression(this, mode)
+        analysisSession.psiTypeProvider.asPsiType(this, context, mode)
 }

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtTypeProvider.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtNamedClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtTypeNullability
+import org.jetbrains.kotlin.psi.KtDoubleColonExpression
 import org.jetbrains.kotlin.psi.KtTypeReference
 
 public abstract class KtTypeProvider : KtAnalysisSessionComponent() {
@@ -18,8 +19,12 @@ public abstract class KtTypeProvider : KtAnalysisSessionComponent() {
 
     public abstract fun buildSelfClassType(symbol: KtNamedClassOrObjectSymbol): KtType
 
+    public abstract fun commonSuperType(types: Collection<KtType>): KtType?
+
     public abstract fun getKtType(ktTypeReference: KtTypeReference): KtType
 
+    public abstract fun getReceiverTypeForDoubleColonExpression(expression: KtDoubleColonExpression): KtType?
+
     public abstract fun withNullability(type: KtType, newNullability: KtTypeNullability): KtType
 }
 
@@ -41,6 +46,14 @@ public interface KtTypeProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtNamedClassOrObjectSymbol.buildSelfClassType(): KtType =
         analysisSession.typeProvider.buildSelfClassType(this)
 
+    /**
+     * Computes the common super type of the given collection of [KtType].
+     *
+     * If the collection is empty, it returns `null`.
+     */
+    public fun commonSuperType(types: Collection<KtType>): KtType? =
+        analysisSession.typeProvider.commonSuperType(types)
+
     /**
      * Resolve [KtTypeReference] and return corresponding [KtType] if resolved.
      *
@@ -49,6 +62,14 @@ public interface KtTypeProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtTypeReference.getKtType(): KtType =
         analysisSession.typeProvider.getKtType(this)
 
+    /**
+     * Resolve [KtDoubleColonExpression] and return [KtType] of its receiver.
+     *
+     * Return `null` if the resolution fails or the resolved callable reference is not a reflection type.
+     */
+    public fun KtDoubleColonExpression.getReceiverKtType(): KtType? =
+        analysisSession.typeProvider.getReceiverTypeForDoubleColonExpression(this)
+
     public fun KtType.withNullability(newNullability: KtTypeNullability): KtType =
         analysisSession.typeProvider.withNullability(this, newNullability)
 }

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
@@ -18,11 +17,9 @@ import org.jetbrains.kotlin.idea.frontend.api.KtStarProjectionTypeArgument
 import org.jetbrains.kotlin.idea.frontend.api.KtTypeArgument
 import org.jetbrains.kotlin.idea.frontend.api.KtTypeArgumentWithVariance
 import org.jetbrains.kotlin.fir.types.ConeTypeCheckerContext
-import org.jetbrains.kotlin.idea.asJava.asPsiType
 import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KT_DIAGNOSTIC_CONVERTER
-import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.types.model.convertVariance
@@ -37,9 +34,6 @@ internal interface KtFirAnalysisSessionComponent {
 
     fun ConeKotlinType.asKtType() = analysisSession.firSymbolBuilder.typeBuilder.buildKtType(this)
 
-    fun ConeKotlinType.asPsiType(mode: TypeMappingMode, psiContext: PsiElement) =
-        asPsiType(rootModuleSession, analysisSession.firResolveState, mode, psiContext)
-
     fun FirPsiDiagnostic.asKtDiagnostic(): KtDiagnosticWithPsi<*> =
         KT_DIAGNOSTIC_CONVERTER.convert(analysisSession, this as FirDiagnostic)
 

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

@@ -7,9 +7,12 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirNamedReference
+import org.jetbrains.kotlin.fir.resolve.constructFunctionalType
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
@@ -30,10 +33,6 @@ internal class KtFirExpressionTypeProvider(
     override val analysisSession: KtFirAnalysisSession,
     override val token: ValidityToken,
 ) : KtExpressionTypeProvider(), KtFirAnalysisSessionComponent {
-    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType = withValidityAssertion {
-        val firDeclaration = declaration.getOrBuildFirOfType<FirCallableDeclaration>(firResolveState)
-        firDeclaration.returnTypeRef.coneType.asKtType()
-    }
 
     override fun getKtExpressionType(expression: KtExpression): KtType = withValidityAssertion {
         when (val fir = expression.unwrap().getOrBuildFir(firResolveState)) {
@@ -44,6 +43,16 @@ internal class KtFirExpressionTypeProvider(
         }
     }
 
+    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType = withValidityAssertion {
+        val firDeclaration = declaration.getOrBuildFirOfType<FirCallableDeclaration>(firResolveState)
+        firDeclaration.returnTypeRef.coneType.asKtType()
+    }
+
+    override fun getFunctionalTypeForKtFunction(declaration: KtFunction): KtType = withValidityAssertion {
+        val firFunction = declaration.getOrBuildFirOfType<FirFunction>(firResolveState)
+        firFunction.constructFunctionalType(firFunction.isSuspend).asKtType()
+    }
+
     override fun getExpectedType(expression: PsiElement): KtType? {
         val expectedType = getExpectedTypeByReturnExpression(expression)
             ?: getExpressionTypeByIfOrBooleanCondition(expression)

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

@@ -5,111 +5,22 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiType
-import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFunction
-import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
-import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.references.FirNamedReference
-import org.jetbrains.kotlin.fir.resolve.constructFunctionalType
-import org.jetbrains.kotlin.fir.typeContext
-import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.throwUnexpectedFirElementError
+import org.jetbrains.kotlin.idea.asJava.asPsiType
 import org.jetbrains.kotlin.idea.frontend.api.components.KtPsiTypeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.getReferencedElementType
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
-import org.jetbrains.kotlin.name.StandardClassIds
-import org.jetbrains.kotlin.psi.*
 
 internal class KtFirPsiTypeProvider(
     override val analysisSession: KtFirAnalysisSession,
     override val token: ValidityToken,
 ) : KtPsiTypeProvider(), KtFirAnalysisSessionComponent {
 
-    override fun commonSuperType(
-        expressions: Collection<KtExpression>,
-        context: KtExpression,
-        mode: TypeMappingMode
-    ): PsiType? = withValidityAssertion {
-        val unitType = analysisSession.rootModuleSession.builtinTypes.unitType.type
-        analysisSession.rootModuleSession.typeContext
-            .commonSuperTypeOrNull(expressions.map { e -> e.getConeType(unitType) { it } })
-            ?.asPsiType(mode, context)
-    }
-
-    override fun getPsiTypeForKtExpression(
-        expression: KtExpression,
-        mode: TypeMappingMode,
-    ): PsiType = withValidityAssertion {
-        expression.getConeType(PsiType.VOID) {
-            it.asPsiType(mode, expression)
-        }
-    }
-
-    private inline fun <T> KtExpression.getConeType(
-        defaultType: T,
-        coneTypeConverter: (ConeKotlinType) -> T,
-    ): T = withValidityAssertion {
-        when (val fir = this.getOrBuildFir(firResolveState)) {
-            is FirExpression -> coneTypeConverter(fir.typeRef.coneType)
-            is FirNamedReference -> coneTypeConverter(fir.getReferencedElementType())
-            is FirStatement -> defaultType
-            else -> throwUnexpectedFirElementError(fir, this)
-        }
-    }
-
-    override fun getPsiTypeForKtDeclaration(
-        ktDeclaration: KtDeclaration,
-        mode: TypeMappingMode
-    ): PsiType = withValidityAssertion {
-        val firDeclaration = ktDeclaration.getOrBuildFirOfType<FirCallableDeclaration>(firResolveState)
-        firDeclaration.returnTypeRef.coneType.asPsiType(mode, ktDeclaration)
-    }
-
-    override fun getPsiTypeForKtFunction(
-        ktFunction: KtFunction,
-        mode: TypeMappingMode
-    ): PsiType = withValidityAssertion {
-        val firFunction = ktFunction.getOrBuildFirOfType<FirFunction>(firResolveState)
-        firFunction.constructFunctionalType(firFunction.isSuspend).asPsiType(mode, ktFunction)
-    }
-
-    override fun getPsiTypeForKtTypeReference(
-        ktTypeReference: KtTypeReference,
-        mode: TypeMappingMode
-    ): PsiType = withValidityAssertion {
-        when (val fir = ktTypeReference.getOrBuildFir(firResolveState)) {
-            // NB: [FirErrorTypeRef] is a subtype of [FirResolvedTypeRef], and the error type in it will be properly handled by [asPsiType].
-            is FirResolvedTypeRef -> fir.coneType.asPsiType(mode, ktTypeReference)
-            else -> throwUnexpectedFirElementError(fir, ktTypeReference)
-        }
-    }
-
-    override fun getPsiTypeForReceiverOfDoubleColonExpression(
-        ktDoubleColonExpression: KtDoubleColonExpression,
-        mode: TypeMappingMode
-    ): PsiType? = withValidityAssertion {
-        val receiver = ktDoubleColonExpression.receiverExpression ?: return null
-        when (val fir = ktDoubleColonExpression.getOrBuildFir(firResolveState)) {
-            is FirGetClassCall ->
-                fir.typeRef.coneType.getReceiverOfReflectionType()?.asPsiType(mode, receiver)
-            is FirCallableReferenceAccess ->
-                fir.typeRef.coneType.getReceiverOfReflectionType()?.asPsiType(mode, receiver)
-            else -> throwUnexpectedFirElementError(fir, ktDoubleColonExpression)
-        }
-    }
-
-    private fun ConeKotlinType.getReceiverOfReflectionType(): ConeKotlinType? {
-        if (this !is ConeClassLikeType) return null
-        if (lookupTag.classId.packageFqName != StandardClassIds.BASE_REFLECT_PACKAGE) return null
-        return typeArguments.firstOrNull()?.type
+    override fun asPsiType(type: KtType, context: PsiElement, mode: TypeMappingMode): PsiType = withValidityAssertion {
+        type.coneType.asPsiType(rootModuleSession, analysisSession.firResolveState, mode, context)
     }
 }

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

@@ -6,15 +6,15 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
+import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
 import org.jetbrains.kotlin.fir.typeContext
-import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
-import org.jetbrains.kotlin.fir.types.withNullability
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.throwUnexpectedFirElementError
 import org.jetbrains.kotlin.idea.frontend.api.components.KtBuiltinTypes
 import org.jetbrains.kotlin.idea.frontend.api.components.KtTypeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
@@ -27,6 +27,8 @@ import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtTypeNullability
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlin.psi.KtDoubleColonExpression
 import org.jetbrains.kotlin.psi.KtTypeReference
 
 internal class KtFirTypeProvider(
@@ -58,10 +60,32 @@ internal class KtFirTypeProvider(
         return type.asKtType()
     }
 
+    override fun commonSuperType(types: Collection<KtType>): KtType? {
+        return analysisSession.rootModuleSession.typeContext
+            .commonSuperTypeOrNull(types.map { it.coneType })
+            ?.asKtType()
+    }
+
     override fun getKtType(ktTypeReference: KtTypeReference): KtType = withValidityAssertion {
         ktTypeReference.getOrBuildFirOfType<FirResolvedTypeRef>(firResolveState).coneType.asKtType()
     }
 
+    override fun getReceiverTypeForDoubleColonExpression(expression: KtDoubleColonExpression): KtType? = withValidityAssertion {
+        when (val fir = expression.getOrBuildFir(firResolveState)) {
+            is FirGetClassCall ->
+                fir.typeRef.coneType.getReceiverOfReflectionType()?.asKtType()
+            is FirCallableReferenceAccess ->
+                fir.typeRef.coneType.getReceiverOfReflectionType()?.asKtType()
+            else -> throwUnexpectedFirElementError(fir, expression)
+        }
+    }
+
+    private fun ConeKotlinType.getReceiverOfReflectionType(): ConeKotlinType? {
+        if (this !is ConeClassLikeType) return null
+        if (lookupTag.classId.packageFqName != StandardClassIds.BASE_REFLECT_PACKAGE) return null
+        return typeArguments.firstOrNull()?.type
+    }
+
     override fun withNullability(type: KtType, newNullability: KtTypeNullability): KtType {
         require(type is KtFirType)
         return type.coneType.withNullability(newNullability.toConeNullability(), rootModuleSession.typeContext).asKtType()