FIR IDE: construct function type from a function as PsiType

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.impl.FirOuterClassTypeParameterRef
 import org.jetbrains.kotlin.fir.declarations.utils.expandedConeType
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
@@ -49,7 +48,7 @@ import org.jetbrains.kotlin.types.SmartcastStability
 fun List<FirQualifierPart>.toTypeProjections(): Array<ConeTypeProjection> =
     asReversed().flatMap { it.typeArgumentList.typeArguments.map { typeArgument -> typeArgument.toConeTypeProjection() } }.toTypedArray()
 
-fun FirFunction.constructFunctionalTypeRef(isSuspend: Boolean = false): FirResolvedTypeRef {
+fun FirFunction.constructFunctionalType(isSuspend: Boolean = false): ConeLookupTagBasedType {
     val receiverTypeRef = when (this) {
         is FirSimpleFunction -> receiverTypeRef
         is FirAnonymousFunction -> receiverTypeRef
@@ -65,11 +64,13 @@ fun FirFunction.constructFunctionalTypeRef(isSuspend: Boolean = false): FirResol
     }
     val rawReturnType = (this as FirTypedDeclaration).returnTypeRef.coneType
 
-    val functionalType = createFunctionalType(parameters, receiverTypeRef?.coneType, rawReturnType, isSuspend = isSuspend)
+    return createFunctionalType(parameters, receiverTypeRef?.coneType, rawReturnType, isSuspend = isSuspend)
+}
 
+fun FirFunction.constructFunctionalTypeRef(isSuspend: Boolean = false): FirResolvedTypeRef {
     return buildResolvedTypeRef {
         source = this@constructFunctionalTypeRef.source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef)
-        type = functionalType
+        type = constructFunctionalType(isSuspend)
     }
 }
 

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

@@ -7,15 +7,13 @@ package org.jetbrains.kotlin.idea.frontend.api.components
 
 import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
-import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtDoubleColonExpression
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtTypeReference
 
 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,
@@ -37,6 +35,9 @@ public interface KtPsiTypeProviderMixIn : KtAnalysisSessionMixIn {
     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)
 

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

@@ -7,11 +7,14 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 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
@@ -24,10 +27,7 @@ import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 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.KtDeclaration
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtDoubleColonExpression
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtTypeReference
 
 internal class KtFirPsiTypeProvider(
     override val analysisSession: KtFirAnalysisSession,
@@ -74,6 +74,14 @@ internal class KtFirPsiTypeProvider(
         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