FIR IDE: introduce TypeInfo as a wrapper for types in high level API

Needed for correct handling types lifecycle

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/highlighter/visitors/ExpressionsSmartcastHighlightingVisitor.kt

@@ -29,7 +29,7 @@ internal class ExpressionsSmartcastHighlightingVisitor(
                     KotlinIdeaAnalysisBundle.message(
                         "0.smart.cast.to.1",
                         receiverName,
-                        analysisSession.renderType(type)
+                        type.asDenotableTypeStringRepresentation()
                     )
                 ).textAttributes = org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors.SMART_CAST_RECEIVER
             }
@@ -39,7 +39,7 @@ internal class ExpressionsSmartcastHighlightingVisitor(
                 getSmartCastTarget(expression),
                 KotlinIdeaAnalysisBundle.message(
                     "smart.cast.to.0",
-                    analysisSession.renderType(type)
+                    type.asDenotableTypeStringRepresentation()
                 )
             ).textAttributes = org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors.SMART_CAST_VALUE
         }

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

@@ -5,20 +5,19 @@
 
 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.*
-import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 
 abstract class FrontendAnalysisSession : Invalidatable {
+    abstract fun getSmartCastedToTypes(expression: KtExpression): Collection<TypeInfo>?
 
     abstract fun getImplicitReceiverSmartCasts(expression: KtExpression): Collection<ImplicitReceiverSmartCast>
 
-    abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KotlinTypeMarker?
+    abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): TypeInfo?
 
-    abstract fun renderType(type: KotlinTypeMarker): String
+    abstract fun getKtExpressionType(expression: KtExpression): TypeInfo?
-
-    abstract fun getKtExpressionType(expression: KtExpression): KotlinTypeMarker?
 
     abstract fun isSubclassOf(klass: KtClassOrObject, superClassId: ClassId): Boolean
 

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

@@ -5,9 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api
 
-import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+data class ImplicitReceiverSmartCast(val types: Collection<TypeInfo>, val kind: ImplicitReceiverSmartcastKind)
-
-data class ImplicitReceiverSmartCast(val types: Collection<KotlinTypeMarker>, val kind: ImplicitReceiverSmartcastKind)
 
 enum class ImplicitReceiverSmartcastKind {
     DISPATCH, EXTENSION

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

@@ -0,0 +1,28 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.name.ClassId
+
+abstract class TypeInfo : Invalidatable {
+    abstract fun isClassType(): Boolean
+    abstract fun classIdIfClassTypeOrError(): ClassId
+
+    abstract fun isErrorType(): Boolean
+
+    abstract fun asDenotableTypeStringRepresentation(): String
+
+    abstract fun isEqualTo(other: TypeInfo): Boolean
+    abstract fun isSubTypeOf(superType: TypeInfo): Boolean
+
+    abstract fun isDefinitelyNullable(): Boolean
+    abstract fun isDefinitelyNotNull(): Boolean
+
+    override fun toString(): String = asDenotableTypeStringRepresentation()
+}
+
+class ErrorTypeClassIdAccessException(override val message: String? = null) : IllegalStateException()
+class ClassTypeExpectedException(override val message: String? = null) : IllegalStateException()

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

@@ -0,0 +1,98 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.idea.frontend.api.ClassTypeExpectedException
+import org.jetbrains.kotlin.idea.frontend.api.ErrorTypeClassIdAccessException
+import org.jetbrains.kotlin.idea.frontend.api.Invalidatable
+import org.jetbrains.kotlin.idea.frontend.api.TypeInfo
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.types.AbstractTypeChecker
+import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
+import java.lang.ref.WeakReference
+
+internal class ConeTypeInfo(
+    coneType: ConeKotlinType,
+    private val typeCheckerContext: ConeTypeCheckerContext,
+    private val validityToken: Invalidatable
+) : TypeInfo() {
+    private val coneTypeWeakRef = WeakReference(coneType)
+
+    override fun isClassType(): Boolean {
+        assertIsValid()
+        return coneType is ConeClassLikeType
+    }
+
+    override fun isErrorType(): Boolean {
+        assertIsValid()
+        return coneType is ConeClassErrorType
+    }
+
+    override fun classIdIfClassTypeOrError(): ClassId {
+        assertIsValid()
+        return when (val coneType = coneType) {
+            is ConeClassLikeTypeImpl -> coneType.lookupTag.classId
+            is ConeClassErrorType -> throw ErrorTypeClassIdAccessException()
+            else -> throw ClassTypeExpectedException()
+        }
+    }
+
+    override fun asDenotableTypeStringRepresentation(): String {
+        assertIsValid()
+        return coneType.render()//TODO
+    }
+
+    override fun isEqualTo(other: TypeInfo): Boolean {
+        assertIsValid()
+        other.assertIsValid()
+        check(other is ConeTypeInfo)
+        return AbstractTypeChecker.equalTypes(
+            typeCheckerContext as AbstractTypeCheckerContext,
+            coneType,
+            other.coneType
+        )
+    }
+
+    override fun isSubTypeOf(superType: TypeInfo): Boolean {
+        assertIsValid()
+        superType.assertIsValid()
+        check(superType is ConeTypeInfo)
+        return AbstractTypeChecker.isSubtypeOf(
+            typeCheckerContext as AbstractTypeCheckerContext,
+            coneType,
+            superType.coneType
+        )
+    }
+
+    override fun isDefinitelyNullable(): Boolean {
+        assertIsValid()
+        return coneType.nullability == ConeNullability.NULLABLE
+    }
+
+    override fun isDefinitelyNotNull(): Boolean {
+        assertIsValid()
+        return coneType.nullability == ConeNullability.NOT_NULL
+    }
+
+    override fun isValid(): Boolean {
+        if (coneTypeWeakRef.get() == null) return false
+        return validityToken.isValid()
+    }
+
+    override fun invalidationReason(): String {
+        if (coneTypeWeakRef.get() == null) return "Cone type was garbage collected"
+        return validityToken.invalidationReason()
+    }
+
+    private inline val coneType
+        get() = coneTypeWeakRef.get() ?: if (validityToken.isValid()) {
+            error("Cone type was garbage collected while analysis session is still valid")
+        } else {
+            error("Accessing the invalid coneType")
+        }
+}

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

@@ -12,10 +12,8 @@ import com.intellij.psi.PsiMethod
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.declarations.isSuspend
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
@@ -25,10 +23,7 @@ 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.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneTypeSafe
-import org.jetbrains.kotlin.fir.types.render
 import org.jetbrains.kotlin.idea.fir.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.fir.isImplicitFunctionCall
@@ -40,19 +35,24 @@ import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper.toTargetPs
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 
-internal class FirAnalysisSession(
+class FirAnalysisSession(
     project: Project
-) : FrontendAnalysisSession(), Invalidatable by ReadActionConfinementValidityToken(project) {
+) : FrontendAnalysisSession() {
+    private val validityToken = ReadActionConfinementValidityToken(project)
+    override fun isValid(): Boolean = validityToken.isValid()
+    override fun invalidationReason(): String = validityToken.invalidationReason()
+
+    constructor(element: KtElement) : this(element.project)
+
     init {
         assertIsValid()
     }
 
-    override fun getSmartCastedToTypes(expression: KtExpression): Collection<KotlinTypeMarker>? {
+    override fun getSmartCastedToTypes(expression: KtExpression): Collection<TypeInfo>? {
         assertIsValid()
         // TODO filter out not used smartcasts
-        return expression.getOrBuildFirSafe<FirExpressionWithSmartcast>()?.typesFromSmartCast
+        return expression.getOrBuildFirSafe<FirExpressionWithSmartcast>()?.typesFromSmartCast?.map { it.asTypeInfo(expression.session) }
     }
 
     @OptIn(ExperimentalStdlibApi::class)
@@ -63,35 +63,39 @@ internal class FirAnalysisSession(
         if (qualifiedExpression.dispatchReceiver !is FirExpressionWithSmartcast
             && qualifiedExpression.extensionReceiver !is FirExpressionWithSmartcast
         ) return emptyList()
+        val session = expression.session
         return buildList {
             (qualifiedExpression.dispatchReceiver as? FirExpressionWithSmartcast)?.let { smartCasted ->
-                ImplicitReceiverSmartCast(smartCasted.typesFromSmartCast, ImplicitReceiverSmartcastKind.DISPATCH)
+                ImplicitReceiverSmartCast(
+                    smartCasted.typesFromSmartCast.map { it.asTypeInfo(session) },
+                    ImplicitReceiverSmartcastKind.DISPATCH
+                )
             }?.let(::add)
             (qualifiedExpression.extensionReceiver as? FirExpressionWithSmartcast)?.let { smartCasted ->
-                ImplicitReceiverSmartCast(smartCasted.typesFromSmartCast, ImplicitReceiverSmartcastKind.EXTENSION)
+                ImplicitReceiverSmartCast(
+                    smartCasted.typesFromSmartCast.map { it.asTypeInfo(session) },
+                    ImplicitReceiverSmartcastKind.EXTENSION
+                )
             }?.let(::add)
         }
     }
 
-    override fun renderType(type: KotlinTypeMarker): String {
+
+    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): TypeInfo? {
         assertIsValid()
-        return type.asConeType().render()
+        val firDeclaration = declaration.toFir<FirCallableDeclaration<*>>() ?: return null
+        return firDeclaration.returnTypeRef.coneTypeSafe<ConeKotlinType>()?.asTypeInfo(declaration.session)
     }
 
-    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KotlinTypeMarker? {
+    override fun getKtExpressionType(expression: KtExpression): TypeInfo? {
         assertIsValid()
-        return declaration.toFir<FirCallableDeclaration<*>>()?.returnTypeRef?.coneTypeSafe()
+        return expression.toFir<FirExpression>()?.typeRef?.coneTypeSafe<ConeKotlinType>()?.asTypeInfo(expression.session)
-    }
-
-    override fun getKtExpressionType(expression: KtExpression): ConeKotlinType? {
-        assertIsValid()
-        return expression.toFir<FirExpression>()?.typeRef?.coneTypeSafe()
     }
 
     override fun isSubclassOf(klass: KtClassOrObject, superClassId: ClassId): Boolean {
         assertIsValid()
         var result = false
-        forEachSubClass(klass.toFir() ?: return false) { type ->
+        forEachSuperClass(klass.toFir() ?: return false) { type ->
             result = result || type.firClassLike(klass.session)?.symbol?.classId == superClassId
         }
         return result
@@ -155,16 +159,23 @@ internal class FirAnalysisSession(
     private inline fun <reified F : FirElement> KtElement.toFir(phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE): F? =
         getOrBuildFir(phase) as? F
 
-    private fun forEachSubClass(firClass: FirClass<*>, action: (FirResolvedTypeRef) -> Unit) {
+    private fun forEachSuperClass(firClass: FirClass<*>, action: (FirResolvedTypeRef) -> Unit) {
         firClass.superTypeRefs.forEach { superType ->
             (superType as? FirResolvedTypeRef)?.let(action)
-            (superType.firClassLike(firClass.session) as? FirClass<*>?)?.let { forEachSubClass(it, action) }
+            (superType.firClassLike(firClass.session) as? FirClass<*>?)?.let { forEachSuperClass(it, action) }
         }
     }
 
+    private fun ConeKotlinType.asTypeInfo(session: FirSession) =
+        ConeTypeInfo(this, createTypeCheckingContext(session), validityToken)
+
+    private fun createTypeCheckingContext(session: FirSession) = ConeTypeCheckerContext(
+        isErrorTypeEqualsToAnything = true, // TODO?
+        isStubTypeEqualsToAnything = true,  // TODO?
+        session = session
+    )
+
     companion object {
-        private fun KotlinTypeMarker.asConeType(): ConeKotlinType =
-            this as? ConeKotlinType ?: error("$this should be ConeKotlinType")
 
         private val kotlinFunctionInvokeCallableIds = (0..23).flatMapTo(hashSetOf()) { arity ->
             listOf(