FIR: make inference context mandatory in withNullability() etc.

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrTypeConverter.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.classId
 import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.typeContext
@@ -222,6 +223,9 @@ class Fir2IrTypeConverter(
     private fun approximateType(type: ConeKotlinType): ConeKotlinType {
         if (type is ConeClassLikeType && type.typeArguments.isEmpty()) return type
         val substitutor = object : AbstractConeSubstitutor() {
+            override val typeInferenceContext: ConeInferenceContext
+                get() = session.inferenceComponents.ctx
+
             override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
                 return if (type is ConeIntersectionType) {
                     type.alternativeType?.let { substituteOrSelf(it) }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/CallAndReferenceGenerator.kt

@@ -665,6 +665,9 @@ class CallAndReferenceGenerator(
     }
 
     private val starProjectionApproximator = object : AbstractConeSubstitutor() {
+        override val typeInferenceContext: ConeInferenceContext
+            get() = session.inferenceComponents.ctx
+
         override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
             if (type !is ConeClassLikeType || type.typeArguments.none { it == ConeStarProjection }) return null
             val fir = type.lookupTag.toSymbol(session)?.fir as? FirTypeParameterRefsOwner ?: return null

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.fir.references.impl.FirReferencePlaceholderForResolv
 import org.jetbrains.kotlin.fir.resolve.bindSymbolToLookupTag
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedReferenceError
+import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredPropertySymbols
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -274,7 +275,7 @@ private fun computeRawProjection(
     session: FirSession,
     parameter: FirTypeParameter,
     attr: TypeComponentPosition,
-    erasedUpperBound: ConeKotlinType = parameter.getErasedUpperBound()
+    erasedUpperBound: ConeKotlinType = parameter.getErasedUpperBound(session)
 ) = when (attr) {
     // Raw(List<T>) => (List<Any?>..List<*>)
     // Raw(Enum<T>) => (Enum<Enum<*>>..Enum<out Enum<*>>)
@@ -306,6 +307,7 @@ private fun computeRawProjection(
 // ErasedUpperBound(T : A) = A // UpperBound(T) is a type A without arguments
 // ErasedUpperBound(T : F) = UpperBound(F) // UB(T) is another type parameter F
 private fun FirTypeParameter.getErasedUpperBound(
+    session: FirSession,
     // Calculation of `potentiallyRecursiveTypeParameter.upperBounds` may recursively depend on `this.getErasedUpperBound`
     // E.g. `class A<T extends A, F extends A>`
     // To prevent recursive calls return defaultValue() instead
@@ -318,10 +320,11 @@ private fun FirTypeParameter.getErasedUpperBound(
 
     val firstUpperBound = this.bounds.first().coneType
 
-    return getErasedVersionOfFirstUpperBound(firstUpperBound, mutableSetOf(this, potentiallyRecursiveTypeParameter), defaultValue)
+    return getErasedVersionOfFirstUpperBound(session, firstUpperBound, mutableSetOf(this, potentiallyRecursiveTypeParameter), defaultValue)
 }
 
 private fun getErasedVersionOfFirstUpperBound(
+    session: FirSession,
     firstUpperBound: ConeKotlinType,
     alreadyVisitedParameters: MutableSet<FirTypeParameter?>,
     defaultValue: () -> ConeKotlinType
@@ -332,18 +335,18 @@ private fun getErasedVersionOfFirstUpperBound(
 
         is ConeFlexibleType -> {
             val lowerBound =
-                getErasedVersionOfFirstUpperBound(firstUpperBound.lowerBound, alreadyVisitedParameters, defaultValue)
+                getErasedVersionOfFirstUpperBound(session, firstUpperBound.lowerBound, alreadyVisitedParameters, defaultValue)
                     .lowerBoundIfFlexible()
             if (firstUpperBound.upperBound is ConeTypeParameterType) {
                 // Avoid exponential complexity
                 ConeFlexibleType(
                     lowerBound,
-                    lowerBound.withNullability(ConeNullability.NULLABLE)
+                    lowerBound.withNullability(ConeNullability.NULLABLE, session.inferenceComponents.ctx)
                 )
             } else {
                 ConeFlexibleType(
                     lowerBound,
-                    getErasedVersionOfFirstUpperBound(firstUpperBound.upperBound, alreadyVisitedParameters, defaultValue)
+                    getErasedVersionOfFirstUpperBound(session, firstUpperBound.upperBound, alreadyVisitedParameters, defaultValue)
                 )
             }
         }
@@ -352,7 +355,7 @@ private fun getErasedVersionOfFirstUpperBound(
 
             if (alreadyVisitedParameters.add(current)) {
                 val nextUpperBound = current.bounds.first().coneType
-                getErasedVersionOfFirstUpperBound(nextUpperBound, alreadyVisitedParameters, defaultValue)
+                getErasedVersionOfFirstUpperBound(session, nextUpperBound, alreadyVisitedParameters, defaultValue)
             } else {
                 defaultValue()
             }
@@ -433,7 +436,7 @@ private fun FirRegularClass.createRawArguments(
     defaultArgs: List<ConeStarProjection>,
     position: TypeComponentPosition
 ): List<ConeTypeProjection> = typeParameters.filterIsInstance<FirTypeParameter>().map { typeParameter ->
-    val erasedUpperBound = typeParameter.getErasedUpperBound {
+    val erasedUpperBound = typeParameter.getErasedUpperBound(session) {
         defaultType().withArguments(defaultArgs.toTypedArray())
     }
     computeRawProjection(session, typeParameter, position, erasedUpperBound)

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

@@ -299,7 +299,9 @@ fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(nullableReceiver
 
     val expandedReceiverType = if (receiverType is ConeClassLikeType) receiverType.fullyExpandedType(session) else receiverType
 
-    val resolvedTypeRef = typeRef.resolvedTypeFromPrototype(expandedReceiverType.makeConeTypeDefinitelyNotNullOrNotNull())
+    val resolvedTypeRef = typeRef.resolvedTypeFromPrototype(
+        expandedReceiverType.makeConeTypeDefinitelyNotNullOrNotNull(session.inferenceComponents.ctx)
+    )
     replaceTypeRef(resolvedTypeRef)
     session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, source, null)
 }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.declarations.expandedConeType
+import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.typeContext
@@ -62,7 +63,7 @@ fun ConeClassLikeType.directExpansionType(
     val resultType = expandedConeType(typeAlias)?.applyNullabilityFrom(useSiteSession, this) ?: return null
 
     if (resultType.typeArguments.isEmpty()) return resultType
-    return mapTypeAliasArguments(typeAlias, this, resultType) as? ConeClassLikeType
+    return mapTypeAliasArguments(typeAlias, this, resultType, useSiteSession) as? ConeClassLikeType
 }
 
 private fun ConeClassLikeType.applyNullabilityFrom(
@@ -77,6 +78,7 @@ private fun mapTypeAliasArguments(
     typeAlias: FirTypeAlias,
     abbreviatedType: ConeClassLikeType,
     resultingType: ConeClassLikeType,
+    useSiteSession: FirSession,
 ): ConeKotlinType {
     if (typeAlias.typeParameters.isNotEmpty() && abbreviatedType.typeArguments.isEmpty()) {
         return resultingType.lookupTag.constructClassType(emptyArray(), resultingType.isNullable)
@@ -84,6 +86,9 @@ private fun mapTypeAliasArguments(
     val typeAliasMap = typeAlias.typeParameters.map { it.symbol }.zip(abbreviatedType.typeArguments).toMap()
 
     val substitutor = object : AbstractConeSubstitutor() {
+        override val typeInferenceContext: ConeInferenceContext
+            get() = useSiteSession.inferenceComponents.ctx
+
         override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
             return null
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/substitution/Substitutors.kt

@@ -16,6 +16,8 @@ import org.jetbrains.kotlin.types.model.TypeSubstitutorMarker
 import org.jetbrains.kotlin.types.model.typeConstructor
 
 abstract class AbstractConeSubstitutor : ConeSubstitutor() {
+    abstract val typeInferenceContext: ConeInferenceContext
+
     private fun wrapProjection(old: ConeTypeProjection, newType: ConeKotlinType): ConeTypeProjection {
         return when (old) {
             is ConeStarProjection -> old
@@ -35,8 +37,8 @@ abstract class AbstractConeSubstitutor : ConeSubstitutor() {
 
     fun ConeKotlinType?.updateNullabilityIfNeeded(originalType: ConeKotlinType): ConeKotlinType? {
         return when {
-            originalType is ConeDefinitelyNotNullType -> this?.withNullability(ConeNullability.NOT_NULL)
+            originalType is ConeDefinitelyNotNullType -> this?.withNullability(ConeNullability.NOT_NULL, typeInferenceContext)
-            originalType.isMarkedNullable -> this?.withNullability(ConeNullability.NULLABLE)
+            originalType.isMarkedNullable -> this?.withNullability(ConeNullability.NULLABLE, typeInferenceContext)
             else -> this
         }
     }
@@ -44,7 +46,7 @@ abstract class AbstractConeSubstitutor : ConeSubstitutor() {
     override fun substituteOrNull(type: ConeKotlinType): ConeKotlinType? {
         val newType = substituteType(type)
         if (newType != null && type is ConeDefinitelyNotNullType) {
-            return newType.makeConeTypeDefinitelyNotNullOrNotNull()
+            return newType.makeConeTypeDefinitelyNotNullOrNotNull(typeInferenceContext)
         }
         return (newType ?: type.substituteRecursive())
     }
@@ -82,7 +84,7 @@ abstract class AbstractConeSubstitutor : ConeSubstitutor() {
     }
 
     private fun ConeDefinitelyNotNullType.substituteOriginal(): ConeKotlinType? {
-        val substituted = substituteOrNull(original)?.withNullability(ConeNullability.NOT_NULL) ?: return null
+        val substituted = substituteOrNull(original)?.withNullability(ConeNullability.NOT_NULL, typeInferenceContext) ?: return null
         return ConeDefinitelyNotNullType.create(substituted) ?: substituted
     }
 
@@ -157,6 +159,9 @@ data class ConeSubstitutorByMap(
     val substitution: Map<FirTypeParameterSymbol, ConeKotlinType>,
     val useSiteSession: FirSession
 ) : AbstractConeSubstitutor() {
+    override val typeInferenceContext: ConeInferenceContext
+        get() = useSiteSession.inferenceComponents.ctx
+
     override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
         if (type !is ConeTypeParameterType) return null
         val result = substitution[type.lookupTag.symbol].updateNullabilityIfNeeded(type) ?: return null
@@ -171,8 +176,10 @@ data class ConeSubstitutorByMap(
 
 fun createTypeSubstitutorByTypeConstructor(map: Map<TypeConstructorMarker, ConeKotlinType>, context: ConeTypeContext): ConeSubstitutor {
     if (map.isEmpty()) return ConeSubstitutor.Empty
-    return object : AbstractConeSubstitutor(),
+    return object : AbstractConeSubstitutor(), TypeSubstitutorMarker {
-        TypeSubstitutorMarker {
+        override val typeInferenceContext: ConeInferenceContext
+            get() = context.session.inferenceComponents.ctx
+
         override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
             if (type !is ConeLookupTagBasedType && type !is ConeStubType) return null
             val new = map[type.typeConstructor(context)] ?: return null

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeInferenceContext.kt

@@ -228,12 +228,12 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
 
     override fun KotlinTypeMarker.makeDefinitelyNotNullOrNotNull(): KotlinTypeMarker {
         require(this is ConeKotlinType)
-        return makeConeTypeDefinitelyNotNullOrNotNull()
+        return makeConeTypeDefinitelyNotNullOrNotNull(this@ConeInferenceContext)
     }
 
     override fun SimpleTypeMarker.makeSimpleTypeDefinitelyNotNullOrNotNull(): SimpleTypeMarker {
         require(this is ConeKotlinType)
-        return makeConeTypeDefinitelyNotNullOrNotNull() as SimpleTypeMarker
+        return makeConeTypeDefinitelyNotNullOrNotNull(this@ConeInferenceContext) as SimpleTypeMarker
     }
 
     override fun createCapturedType(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
 import org.jetbrains.kotlin.fir.expressions.arguments
 import org.jetbrains.kotlin.fir.resolve.correspondingSupertypesCache
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -128,7 +129,7 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 
     override fun SimpleTypeMarker.withNullability(nullable: Boolean): SimpleTypeMarker {
         require(this is ConeKotlinType)
-        return withNullability(ConeNullability.create(nullable), this as? ConeInferenceContext)
+        return withNullability(ConeNullability.create(nullable), session.inferenceComponents.ctx)
     }
 
     override fun SimpleTypeMarker.typeConstructor(): TypeConstructorMarker {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/TypeUtils.kt

@@ -56,11 +56,11 @@ fun ConeDefinitelyNotNullType.Companion.create(original: ConeKotlinType): ConeDe
     }
 }
 
-fun ConeKotlinType.makeConeTypeDefinitelyNotNullOrNotNull(): ConeKotlinType {
+fun ConeKotlinType.makeConeTypeDefinitelyNotNullOrNotNull(typeContext: ConeInferenceContext): ConeKotlinType {
     if (this is ConeIntersectionType) {
-        return ConeIntersectionType(intersectedTypes.map { it.makeConeTypeDefinitelyNotNullOrNotNull() })
+        return ConeIntersectionType(intersectedTypes.map { it.makeConeTypeDefinitelyNotNullOrNotNull(typeContext) })
     }
-    return ConeDefinitelyNotNullType.create(this) ?: this.withNullability(ConeNullability.NOT_NULL)
+    return ConeDefinitelyNotNullType.create(this) ?: this.withNullability(ConeNullability.NOT_NULL, typeContext)
 }
 
 fun <T : ConeKotlinType> T.withArguments(arguments: Array<out ConeTypeProjection>): T {
@@ -95,7 +95,7 @@ fun <T : ConeKotlinType> T.withAttributes(attributes: ConeAttributes): T {
 
 fun <T : ConeKotlinType> T.withNullability(
     nullability: ConeNullability,
-    typeContext: ConeInferenceContext? = null,
+    typeContext: ConeInferenceContext,
     attributes: ConeAttributes = this.attributes,
 ): T {
     if (this.nullability == nullability && this.attributes == attributes) {
@@ -123,7 +123,7 @@ fun <T : ConeKotlinType> T.withNullability(
         is ConeCapturedType -> ConeCapturedType(captureStatus, lowerType, nullability, constructor, attributes)
         is ConeIntersectionType -> when (nullability) {
             ConeNullability.NULLABLE -> this.mapTypes {
-                it.withNullability(nullability)
+                it.withNullability(nullability, typeContext)
             }
             ConeNullability.UNKNOWN -> this // TODO: is that correct?
             ConeNullability.NOT_NULL -> this
@@ -140,7 +140,7 @@ fun <T : ConeKotlinType> T.withNullability(
 }
 
 fun coneFlexibleOrSimpleType(
-    typeContext: ConeInferenceContext?,
+    typeContext: ConeInferenceContext,
     lowerBound: ConeKotlinType,
     upperBound: ConeKotlinType,
 ): ConeKotlinType {
@@ -151,15 +151,8 @@ fun coneFlexibleOrSimpleType(
         return coneFlexibleOrSimpleType(typeContext, lowerBound, upperBound.upperBound)
     }
     return when {
-        typeContext != null && AbstractStrictEqualityTypeChecker.strictEqualTypes(typeContext, lowerBound, upperBound) -> {
+        AbstractStrictEqualityTypeChecker.strictEqualTypes(typeContext, lowerBound, upperBound) -> lowerBound
-            lowerBound
+        else -> ConeFlexibleType(lowerBound, upperBound)
-        }
-        typeContext == null && lowerBound == upperBound -> {
-            lowerBound
-        }
-        else -> {
-            ConeFlexibleType(lowerBound, upperBound)
-        }
     }
 }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/firLightUtils.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.backend.jvm.jvmTypeMapper
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.isPrimitiveType
+import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
@@ -76,8 +77,13 @@ internal fun KtNamedClassOrObjectSymbol.typeForClassSymbol(psiElement: PsiElemen
     return types.coneType.asPsiType(session, firRef.resolveState, TypeMappingMode.DEFAULT, psiElement)
 }
 
-private class AnonymousTypesSubstitutor(private val session: FirSession, private val state: FirModuleResolveState) :
+private class AnonymousTypesSubstitutor(
-    AbstractConeSubstitutor() {
+    private val session: FirSession,
+    private val state: FirModuleResolveState
+) : AbstractConeSubstitutor() {
+    override val typeInferenceContext: ConeInferenceContext
+        get() = session.inferenceComponents.ctx
+
     override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
 
         if (type !is ConeClassLikeType) return null