FIR IDE: do not store cone session in every KtFirType

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -67,9 +67,9 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
     private val scopeNameFilter: KtScopeNameFilter =
         { name -> !name.isSpecial && prefixMatcher.prefixMatches(name.identifier) }
 
-    private fun CompletionResultSet.addSymbolToCompletion(symbol: KtSymbol) {
+    private fun KtAnalysisSession.addSymbolToCompletion(completionResultSet: CompletionResultSet, symbol: KtSymbol) {
         if (symbol !is KtNamedSymbol) return
-        lookupElementFactory.createLookupElement(symbol)?.let(::addElement)
+        with(lookupElementFactory) { createLookupElement(symbol)?.let(completionResultSet::addElement) }
     }
 
     private fun recordOriginalFile(completionParameters: CompletionParameters) {
@@ -103,9 +103,9 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
         }
     }
 
-    private fun collectTypesCompletion(result: CompletionResultSet, implicitScopes: KtScope) {
+    private fun KtAnalysisSession.collectTypesCompletion(result: CompletionResultSet, implicitScopes: KtScope) {
         val availableClasses = implicitScopes.getClassifierSymbols(scopeNameFilter)
-        availableClasses.forEach { result.addSymbolToCompletion(it) }
+        availableClasses.forEach { addSymbolToCompletion(result, it) }
     }
 
     private fun KtAnalysisSession.collectDotCompletion(
@@ -125,11 +125,11 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
             .getCallableSymbols(scopeNameFilter)
             .filter { it.isExtension && it.hasSuitableExtensionReceiver() }
 
-        nonExtensionMembers.forEach { result.addSymbolToCompletion(it) }
+        nonExtensionMembers.forEach { addSymbolToCompletion(result, it) }
-        extensionNonMembers.forEach { result.addSymbolToCompletion(it) }
+        extensionNonMembers.forEach { addSymbolToCompletion(result, it) }
     }
 
-    private fun collectDefaultCompletion(
+    private fun KtAnalysisSession.collectDefaultCompletion(
         result: CompletionResultSet,
         implicitScopes: KtCompositeScope,
         hasSuitableExtensionReceiver: KtCallableSymbol.() -> Boolean,
@@ -142,8 +142,8 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
             .getCallableSymbols(scopeNameFilter)
             .filter { it.isExtension && it.hasSuitableExtensionReceiver() }
 
-        availableNonExtensions.forEach { result.addSymbolToCompletion(it) }
+        availableNonExtensions.forEach { addSymbolToCompletion(result, it) }
-        extensionsWhichCanBeCalled.forEach { result.addSymbolToCompletion(it) }
+        extensionsWhichCanBeCalled.forEach { addSymbolToCompletion(result, it) }
 
         collectTypesCompletion(result, implicitScopes)
     }

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirLookupElementFactory.kt

@@ -14,9 +14,9 @@ import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.editor.Document
-import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
@@ -33,9 +33,9 @@ internal class KotlinFirLookupElementFactory {
     private val functionLookupElementFactory = FunctionLookupElementFactory()
     private val typeParameterLookupElementFactory = TypeParameterLookupElementFactory()
 
-    fun createLookupElement(symbol: KtNamedSymbol): LookupElement? {
+    fun KtAnalysisSession.createLookupElement(symbol: KtNamedSymbol): LookupElement? {
         val elementBuilder = when (symbol) {
-            is KtFunctionSymbol -> functionLookupElementFactory.createLookup(symbol)
+            is KtFunctionSymbol -> with(functionLookupElementFactory) { createLookup(symbol) }
             is KtVariableLikeSymbol -> variableLookupElementFactory.createLookup(symbol)
             is KtClassLikeSymbol -> classLookupElementFactory.createLookup(symbol)
             is KtTypeParameterSymbol -> typeParameterLookupElementFactory.createLookup(symbol)
@@ -78,7 +78,7 @@ private class VariableLookupElementFactory {
 }
 
 private class FunctionLookupElementFactory {
-    fun createLookup(symbol: KtFunctionSymbol): LookupElementBuilder? {
+    fun KtAnalysisSession.createLookup(symbol: KtFunctionSymbol): LookupElementBuilder? {
         return try {
             LookupElementBuilder.create(UniqueLookupObject(), symbol.name.asString())
                 .withTailText(getTailText(symbol), true)
@@ -91,16 +91,16 @@ private class FunctionLookupElementFactory {
         }
     }
 
-    private fun getTailText(symbol: KtFunctionSymbol): String {
+    private fun KtAnalysisSession.getTailText(symbol: KtFunctionSymbol): String {
         return if (insertLambdaBraces(symbol)) " {...}" else ShortNamesRenderer.renderFunctionParameters(symbol)
     }
 
-    private fun insertLambdaBraces(symbol: KtFunctionSymbol): Boolean {
+    private fun KtAnalysisSession.insertLambdaBraces(symbol: KtFunctionSymbol): Boolean {
         val singleParam = symbol.valueParameters.singleOrNull()
-        return singleParam != null && !singleParam.hasDefaultValue && singleParam.type.isBuiltInFunctionalType
+        return singleParam != null && !singleParam.hasDefaultValue && singleParam.type.isBuiltInFunctionalType()
     }
 
-    private fun createInsertHandler(symbol: KtFunctionSymbol): InsertHandler<LookupElement> {
+    private fun KtAnalysisSession.createInsertHandler(symbol: KtFunctionSymbol): InsertHandler<LookupElement> {
         return FunctionInsertionHandler(
             symbol.name,
             inputValueArguments = symbol.valueParameters.isNotEmpty(),

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.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.idea.frontend.api.components.*
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
@@ -58,6 +57,12 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
 
     fun KtDeclaration.getReturnKtType(): KtType = typeProvider.getReturnTypeForKtDeclaration(this)
 
+    fun KtType.isEqualTo(other: KtType): Boolean = typeProvider.isEqualTo(this, other)
+
+    fun KtType.isSubTypeOf(superType: KtType): Boolean = typeProvider.isSubTypeOf(this, superType)
+
+    fun KtType.isBuiltInFunctionalType(): Boolean = typeProvider.isBuiltinFunctionalType(this)
+
     fun KtElement.getDiagnostics(): Collection<Diagnostic> = diagnosticProvider.getDiagnosticsForElement(this)
 
     fun KtFile.collectDiagnosticsForFile(): Collection<Diagnostic> = diagnosticProvider.collectDiagnosticsForFile(this)

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

@@ -12,4 +12,10 @@ import org.jetbrains.kotlin.psi.KtExpression
 abstract class KtTypeProvider : KtAnalysisSessionComponent() {
     abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType
     abstract fun getKtExpressionType(expression: KtExpression): KtType
+
+    abstract fun isEqualTo(first: KtType, second: KtType): Boolean
+    abstract fun isSubTypeOf(subType: KtType, superType: KtType): Boolean
+
+    //TODO get rid of
+    abstract fun isBuiltinFunctionalType(type: KtType): Boolean
 }

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/types/KtType.kt

@@ -14,11 +14,6 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 interface KtType : ValidityTokenOwner {
-    fun isEqualTo(other: KtType): Boolean
-    fun isSubTypeOf(superType: KtType): Boolean
-
-    val isBuiltInFunctionalType: Boolean
-
     fun asStringForDebugging(): String
 }
 

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

@@ -43,13 +43,7 @@ internal class KtSymbolByFirBuilder private constructor(
     private val symbolsCache: BuilderCache<FirDeclaration, KtSymbol>,
     private val typesCache: BuilderCache<ConeKotlinType, KtType>
 ) : ValidityTokenOwner {
-    private val typeCheckerContext by threadLocal {
+    private val resolveState by weakRef(resolveState)
-        ConeTypeCheckerContext(
-            isErrorTypeEqualsToAnything = true,
-            isStubTypeEqualsToAnything = true,
-            resolveState.rootModuleSession
-        )
-    }
 
     private val firProvider get() = resolveState.rootModuleSession.firSymbolProvider
 
@@ -66,7 +60,6 @@ internal class KtSymbolByFirBuilder private constructor(
         typesCache = BuilderCache()
     )
 
-    private val resolveState by weakRef(resolveState)
 
     fun createReadOnlyCopy(newResolveState: FirModuleResolveState): KtSymbolByFirBuilder {
         check(!withReadOnlyCaching) { "Cannot create readOnly KtSymbolByFirBuilder from a readonly one" }
@@ -197,11 +190,11 @@ internal class KtSymbolByFirBuilder private constructor(
 
     fun buildKtType(coneType: ConeKotlinType): KtType = typesCache.cache(coneType) {
         when (coneType) {
-            is ConeClassLikeTypeImpl -> KtFirClassType(coneType, typeCheckerContext, token, this)
+            is ConeClassLikeTypeImpl -> KtFirClassType(coneType, token, this)
-            is ConeTypeParameterType -> KtFirTypeParameterType(coneType, typeCheckerContext, token, this)
+            is ConeTypeParameterType -> KtFirTypeParameterType(coneType, token, this)
-            is ConeClassErrorType -> KtFirErrorType(coneType, typeCheckerContext, token)
+            is ConeClassErrorType -> KtFirErrorType(coneType, token)
-            is ConeFlexibleType -> KtFirFlexibleType(coneType, typeCheckerContext, token, this)
+            is ConeFlexibleType -> KtFirFlexibleType(coneType, token, this)
-            is ConeIntersectionType -> KtFirIntersectionType(coneType, typeCheckerContext, token, this)
+            is ConeIntersectionType -> KtFirIntersectionType(coneType, token, this)
             is ConeDefinitelyNotNullType -> buildKtType(coneType.original)
             else -> TODO(coneType::class.toString())
         }

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

@@ -7,20 +7,21 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
+import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.types.ConeTypeCheckerContext
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
-import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartCast
-import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartcastKind
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
-import org.jetbrains.kotlin.idea.frontend.api.components.KtSmartCastProvider
+import org.jetbrains.kotlin.idea.frontend.api.assertIsValid
 import org.jetbrains.kotlin.idea.frontend.api.components.KtTypeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.types.AbstractTypeChecker
+import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 
 internal class KtFirTypeProvider(
     override val analysisSession: KtFirAnalysisSession,
@@ -34,4 +35,37 @@ internal class KtFirTypeProvider(
     override fun getKtExpressionType(expression: KtExpression): KtType = withValidityAssertion {
         expression.getOrBuildFirOfType<FirExpression>(firResolveState).typeRef.coneType.asKtType()
     }
+
+    override fun isEqualTo(first: KtType, second: KtType): Boolean = withValidityAssertion {
+        second.assertIsValid()
+        check(first is KtFirType)
+        check(second is KtFirType)
+        return AbstractTypeChecker.equalTypes(
+            this.createTypeCheckerContext() as AbstractTypeCheckerContext,
+            first.coneType,
+            second.coneType
+        )
+    }
+
+    override fun isSubTypeOf(subType: KtType, superType: KtType): Boolean = withValidityAssertion {
+        superType.assertIsValid()
+        check(subType is KtFirType)
+        check(superType is KtFirType)
+        return AbstractTypeChecker.isSubtypeOf(
+            this.createTypeCheckerContext() as AbstractTypeCheckerContext,
+            subType.coneType,
+            superType.coneType
+        )
+    }
+
+    override fun isBuiltinFunctionalType(type: KtType): Boolean = withValidityAssertion {
+        check(type is KtFirType)
+        type.coneType.isBuiltinFunctionalType(analysisSession.firResolveState.rootModuleSession) //TODO use correct session here
+    }
+
+    private fun createTypeCheckerContext() = ConeTypeCheckerContext(
+        isErrorTypeEqualsToAnything = true,
+        isStubTypeEqualsToAnything = true,
+        analysisSession.firResolveState.rootModuleSession //TODO use correct session here
+    )
 }

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

@@ -18,47 +18,19 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.KtTypeParameterSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.types.AbstractTypeChecker
-import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 
 internal interface KtFirType : KtType, ValidityTokenOwner {
     val coneType: ConeKotlinType
-    val typeCheckerContext: ConeTypeCheckerContext
 
     override fun asStringForDebugging(): String = withValidityAssertion { coneType.render() }
-
-    override fun isEqualTo(other: KtType): Boolean = withValidityAssertion {
-        other.assertIsValid()
-        check(other is KtFirType)
-        return AbstractTypeChecker.equalTypes(
-            typeCheckerContext as AbstractTypeCheckerContext,
-            coneType,
-            other.coneType
-        )
-    }
-
-    override fun isSubTypeOf(superType: KtType): Boolean = withValidityAssertion {
-        superType.assertIsValid()
-        check(superType is KtFirType)
-        return AbstractTypeChecker.isSubtypeOf(
-            typeCheckerContext as AbstractTypeCheckerContext,
-            coneType,
-            superType.coneType
-        )
-    }
-
-    override val isBuiltInFunctionalType: Boolean
-        get() = coneType.isBuiltinFunctionalType(typeCheckerContext.session)
 }
 
 internal class KtFirClassType(
     coneType: ConeClassLikeTypeImpl,
-    typeCheckerContext: ConeTypeCheckerContext,
     override val token: ValidityToken,
     private val firBuilder: KtSymbolByFirBuilder,
 ) : KtClassType(), KtFirType {
     override val coneType by weakRef(coneType)
-    override val typeCheckerContext by weakRef(typeCheckerContext)
 
     override val classId: ClassId get() = withValidityAssertion { coneType.lookupTag.classId }
     override val classSymbol: KtClassLikeSymbol by cached {
@@ -79,23 +51,19 @@ internal class KtFirClassType(
 
 internal class KtFirErrorType(
     coneType: ConeClassErrorType,
-    typeCheckerContext: ConeTypeCheckerContext,
     override val token: ValidityToken,
 ) : KtErrorType(), KtFirType {
     override val coneType by weakRef(coneType)
-    override val typeCheckerContext by weakRef(typeCheckerContext)
 
     override val error: String get() = withValidityAssertion { coneType.diagnostic.reason }
 }
 
 internal class KtFirTypeParameterType(
     coneType: ConeTypeParameterType,
-    typeCheckerContext: ConeTypeCheckerContext,
     override val token: ValidityToken,
     private val firBuilder: KtSymbolByFirBuilder,
 ) : KtTypeParameterType(), KtFirType {
     override val coneType by weakRef(coneType)
-    override val typeCheckerContext by weakRef(typeCheckerContext)
 
     override val name: Name get() = withValidityAssertion { coneType.lookupTag.name }
     override val symbol: KtTypeParameterSymbol by cached {
@@ -112,12 +80,10 @@ internal class KtFirTypeParameterType(
 
 internal class KtFirFlexibleType(
     coneType: ConeFlexibleType,
-    typeCheckerContext: ConeTypeCheckerContext,
     override val token: ValidityToken,
     private val firBuilder: KtSymbolByFirBuilder,
 ) : KtFlexibleType(), KtFirType {
     override val coneType by weakRef(coneType)
-    override val typeCheckerContext by weakRef(typeCheckerContext)
 
     override val lowerBound: KtType by cached { firBuilder.buildKtType(coneType.lowerBound) }
     override val upperBound: KtType by cached { firBuilder.buildKtType(coneType.upperBound) }
@@ -125,12 +91,10 @@ internal class KtFirFlexibleType(
 
 internal class KtFirIntersectionType(
     coneType: ConeIntersectionType,
-    typeCheckerContext: ConeTypeCheckerContext,
     override val token: ValidityToken,
     private val firBuilder: KtSymbolByFirBuilder,
 ) : KtIntersectionType(), KtFirType {
     override val coneType by weakRef(coneType)
-    override val typeCheckerContext by weakRef(typeCheckerContext)
 
     override val conjuncts: List<KtType> by cached {
         coneType.intersectedTypes.map { conjunct -> firBuilder.buildKtType(conjunct) }