FIR IDE: make API functions that returns types return non-null value

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.frontend.api
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
@@ -15,15 +14,14 @@ abstract class FrontendAnalysisSession : Invalidatable {
 
     abstract fun getImplicitReceiverSmartCasts(expression: KtExpression): Collection<ImplicitReceiverSmartCast>
 
-    abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): TypeInfo?
+    abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): TypeInfo
 
-    abstract fun getKtExpressionType(expression: KtExpression): TypeInfo?
+    abstract fun getKtExpressionType(expression: KtExpression): TypeInfo
 
     abstract fun isSubclassOf(klass: KtClassOrObject, superClassId: ClassId): Boolean
 
     abstract fun getDiagnosticsForElement(element: KtElement): Collection<Diagnostic>
 
     abstract fun resolveCall(call: KtCallExpression): CallInfo?
-
     abstract fun resolveCall(call: KtBinaryExpression): CallInfo?
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt

@@ -20,16 +20,30 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import kotlin.reflect.KClass
 
 fun KtElement.getOrBuildFir(
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
 ) = LowLevelFirApiFacade.getOrBuildFirFor(this, phase)
 
-
 inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
 ) = LowLevelFirApiFacade.getOrBuildFirFor(this, phase) as? E
 
+inline fun <reified E : FirElement> KtElement.getOrBuildFirOfType(
+    phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
+): E {
+    val fir = LowLevelFirApiFacade.getOrBuildFirFor(this, phase)
+    if (fir is E) return fir
+    throw InvalidFirElementTypeException(this, E::class, fir::class)
+}
+
+class InvalidFirElementTypeException(
+    ktElement: KtElement,
+    expectedFirClass: KClass<out FirElement>,
+    actualFirClass: KClass<out FirElement>
+) : IllegalStateException("For $ktElement with text `${ktElement.text}` the $expectedFirClass expected, but $actualFirClass found")
+
 
 val KtElement.session
     get() = parentOfType<KtDeclaration>()?.getOrBuildFirSafe<FirDeclaration>()?.session

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

@@ -24,11 +24,8 @@ 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.idea.fir.getOrBuildFir
+import org.jetbrains.kotlin.idea.fir.*
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
-import org.jetbrains.kotlin.idea.fir.isImplicitFunctionCall
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
-import org.jetbrains.kotlin.idea.fir.session
 import org.jetbrains.kotlin.idea.frontend.api.*
 import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper
 import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper.toTargetPsi
@@ -81,21 +78,21 @@ class FirAnalysisSession(
     }
 
 
-    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): TypeInfo? {
+    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): TypeInfo {
         assertIsValid()
-        val firDeclaration = declaration.toFir<FirCallableDeclaration<*>>() ?: return null
+        val firDeclaration = declaration.getOrBuildFirOfType<FirCallableDeclaration<*>>()
-        return firDeclaration.returnTypeRef.coneTypeSafe<ConeKotlinType>()?.asTypeInfo(declaration.session)
+        return firDeclaration.returnTypeRef.coneTypeUnsafe<ConeKotlinType>().asTypeInfo(declaration.session)
     }
 
-    override fun getKtExpressionType(expression: KtExpression): TypeInfo? {
+    override fun getKtExpressionType(expression: KtExpression): TypeInfo {
         assertIsValid()
-        return expression.toFir<FirExpression>()?.typeRef?.coneTypeSafe<ConeKotlinType>()?.asTypeInfo(expression.session)
+        return expression.getOrBuildFirOfType<FirExpression>().typeRef.coneTypeUnsafe<ConeKotlinType>().asTypeInfo(expression.session)
     }
 
     override fun isSubclassOf(klass: KtClassOrObject, superClassId: ClassId): Boolean {
         assertIsValid()
         var result = false
-        forEachSuperClass(klass.toFir() ?: return false) { type ->
+        forEachSuperClass(klass.getOrBuildFirSafe() ?: return false) { type ->
             result = result || type.firClassLike(klass.session)?.symbol?.classId == superClassId
         }
         return result
@@ -108,13 +105,13 @@ class FirAnalysisSession(
 
     override fun resolveCall(call: KtBinaryExpression): CallInfo? {
         assertIsValid()
-        val firCall = call.toFir<FirFunctionCall>() ?: return null
+        val firCall = call.getOrBuildFirSafe<FirFunctionCall>() ?: return null
         return resolveCall(firCall, call)
     }
 
     override fun resolveCall(call: KtCallExpression): CallInfo? {
         assertIsValid()
-        val firCall = call.toFir<FirFunctionCall>() ?: return null
+        val firCall = call.getOrBuildFirSafe<FirFunctionCall>() ?: return null
         return resolveCall(firCall, call)
     }
 
@@ -156,9 +153,6 @@ class FirAnalysisSession(
         else -> null
     }
 
-    private inline fun <reified F : FirElement> KtElement.toFir(phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE): F? =
-        getOrBuildFir(phase) as? F
-
     private fun forEachSuperClass(firClass: FirClass<*>, action: (FirResolvedTypeRef) -> Unit) {
         firClass.superTypeRefs.forEach { superType ->
             (superType as? FirResolvedTypeRef)?.let(action)