FIR: Refine makesSenseToBeDefinitelyNotNull

Make it work just the same as the analogue from FE 1.0
This change is necessary since many tests start failing after
we began reporting diagnostics after call completion

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirTypeDeserializer.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
@@ -137,7 +138,7 @@ class FirTypeDeserializer(
         if (constructor is ConeTypeParameterLookupTag) {
             return ConeTypeParameterTypeImpl(constructor, isNullable = proto.nullable).let {
                 if (Flags.DEFINITELY_NOT_NULL_TYPE.get(proto.flags))
-                    ConeDefinitelyNotNullType.create(it)
+                    ConeDefinitelyNotNullType.create(it, moduleData.session.typeContext)
                 else
                     it
             }

compiler/fir/fir2ir/jvm-backend/src/org/jetbrains/kotlin/fir/backend/jvm/FirJvmTypeMapper.kt

@@ -233,7 +233,10 @@ class ConeTypeSystemCommonBackendContextForTypeMapping(
             require(it is ConeKotlinType)
         }
         @Suppress("UNCHECKED_CAST")
-        return defaultType().withArguments((arguments as List<ConeKotlinType>).toTypedArray())
+        return defaultType().withArguments(
+            (arguments as List<ConeKotlinType>).toTypedArray(),
+            session.typeContext,
+        )
     }
 
     override fun TypeParameterMarker.representativeUpperBound(): ConeKotlinType {

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

@@ -13,7 +13,6 @@ 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
-import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
@@ -24,6 +23,7 @@ import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.IrStarProjectionImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.Variance
 
@@ -230,10 +230,7 @@ class Fir2IrTypeConverter(
 
     private fun approximateType(type: ConeKotlinType): ConeKotlinType {
         if (type is ConeClassLikeType && type.typeArguments.isEmpty()) return type
-        val substitutor = object : AbstractConeSubstitutor() {
+        val substitutor = object : AbstractConeSubstitutor(session.typeContext) {
-            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

@@ -664,10 +664,7 @@ class CallAndReferenceGenerator(
         return IrTypeOperatorCallImpl(this.startOffset, this.endOffset, samType, IrTypeOperator.SAM_CONVERSION, samType, this)
     }
 
-    private val starProjectionApproximator = object : AbstractConeSubstitutor() {
+    private val starProjectionApproximator = object : AbstractConeSubstitutor(session.typeContext) {
-        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/enhancement/javaTypeUtils.kt

@@ -128,7 +128,11 @@ private fun coneFlexibleOrSimpleType(
                     type is ConeTypeParameterType || type.isNullable
                 }
             ) {
-                return ConeDefinitelyNotNullType.create(lowerBound) ?: lowerBound
+                return ConeDefinitelyNotNullType.create(
+                    lowerBound,
+                    session.typeContext,
+                    useCorrectedNullabilityForFlexibleTypeParameters = true
+                ) ?: lowerBound
             }
         }
         return lowerBound

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

@@ -34,12 +34,8 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.ForbiddenNamedArgumentsTarget
 import org.jetbrains.kotlin.name.*
+import org.jetbrains.kotlin.resolve.ForbiddenNamedArgumentsTarget
 
 fun List<FirQualifierPart>.toTypeProjections(): Array<ConeTypeProjection> =
     asReversed().flatMap { it.typeArgumentList.typeArguments.map { typeArgument -> typeArgument.toConeTypeProjection() } }.toTypedArray()
@@ -326,9 +322,8 @@ fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(nullableReceiver
 
     val expandedReceiverType = if (receiverType is ConeClassLikeType) receiverType.fullyExpandedType(session) else receiverType
 
-    val resolvedTypeRef = typeRef.resolvedTypeFromPrototype(
+    val resolvedTypeRef =
-        expandedReceiverType.makeConeTypeDefinitelyNotNullOrNotNull(session.inferenceComponents.ctx)
+        typeRef.resolvedTypeFromPrototype(expandedReceiverType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext))
-    )
     replaceTypeRef(resolvedTypeRef)
     session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, source, null)
 }

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

@@ -8,7 +8,6 @@ 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
@@ -85,10 +84,7 @@ private fun mapTypeAliasArguments(
     }
     val typeAliasMap = typeAlias.typeParameters.map { it.symbol }.zip(abbreviatedType.typeArguments).toMap()
 
-    val substitutor = object : AbstractConeSubstitutor() {
+    val substitutor = object : AbstractConeSubstitutor(useSiteSession.typeContext) {
-        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/calls/Arguments.kt

@@ -600,7 +600,7 @@ internal fun captureFromTypeParameterUpperBoundIfNeeded(
 
     val capturedType = context.captureFromExpression(chosenSupertype) as ConeKotlinType? ?: return argumentType
     return if (argumentType is ConeDefinitelyNotNullType) {
-        ConeDefinitelyNotNullType.create(capturedType) ?: capturedType
+        ConeDefinitelyNotNullType.create(capturedType, session.typeContext) ?: capturedType
     } else {
         capturedType
     }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.substitution
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
@@ -15,9 +16,7 @@ import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeSubstitutorMarker
 import org.jetbrains.kotlin.types.model.typeConstructor
 
-abstract class AbstractConeSubstitutor : ConeSubstitutor() {
+abstract class AbstractConeSubstitutor(private val typeContext: ConeTypeContext) : ConeSubstitutor() {
-    abstract val typeInferenceContext: ConeInferenceContext
-
     private fun wrapProjection(old: ConeTypeProjection, newType: ConeKotlinType): ConeTypeProjection {
         return when (old) {
             is ConeStarProjection -> old
@@ -37,8 +36,8 @@ abstract class AbstractConeSubstitutor : ConeSubstitutor() {
 
     fun ConeKotlinType?.updateNullabilityIfNeeded(originalType: ConeKotlinType): ConeKotlinType? {
         return when {
-            originalType is ConeDefinitelyNotNullType -> this?.withNullability(ConeNullability.NOT_NULL, typeInferenceContext)
+            originalType is ConeDefinitelyNotNullType -> this?.withNullability(ConeNullability.NOT_NULL, typeContext)
-            originalType.isMarkedNullable -> this?.withNullability(ConeNullability.NULLABLE, typeInferenceContext)
+            originalType.isMarkedNullable -> this?.withNullability(ConeNullability.NULLABLE, typeContext)
             else -> this
         }
     }
@@ -46,7 +45,7 @@ abstract class AbstractConeSubstitutor : ConeSubstitutor() {
     override fun substituteOrNull(type: ConeKotlinType): ConeKotlinType? {
         val newType = substituteType(type)
         if (newType != null && type is ConeDefinitelyNotNullType) {
-            return newType.makeConeTypeDefinitelyNotNullOrNotNull(typeInferenceContext)
+            return newType.makeConeTypeDefinitelyNotNullOrNotNull(typeContext)
         }
         return (newType ?: type.substituteRecursive())
     }
@@ -84,8 +83,8 @@ abstract class AbstractConeSubstitutor : ConeSubstitutor() {
     }
 
     private fun ConeDefinitelyNotNullType.substituteOriginal(): ConeKotlinType? {
-        val substituted = substituteOrNull(original)?.withNullability(ConeNullability.NOT_NULL, typeInferenceContext) ?: return null
+        val substituted = substituteOrNull(original)?.withNullability(ConeNullability.NOT_NULL, typeContext) ?: return null
-        return ConeDefinitelyNotNullType.create(substituted) ?: substituted
+        return ConeDefinitelyNotNullType.create(substituted, typeContext) ?: substituted
     }
 
     private fun ConeFlexibleType.substituteBounds(): ConeFlexibleType? {
@@ -158,10 +157,7 @@ fun ConeSubstitutor.chain(other: ConeSubstitutor): ConeSubstitutor {
 data class ConeSubstitutorByMap(
     val substitution: Map<FirTypeParameterSymbol, ConeKotlinType>,
     val useSiteSession: FirSession
-) : AbstractConeSubstitutor() {
+) : AbstractConeSubstitutor(useSiteSession.typeContext) {
-    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
@@ -176,10 +172,7 @@ data class ConeSubstitutorByMap(
 
 fun createTypeSubstitutorByTypeConstructor(map: Map<TypeConstructorMarker, ConeKotlinType>, context: ConeTypeContext): ConeSubstitutor {
     if (map.isEmpty()) return ConeSubstitutor.Empty
-    return object : AbstractConeSubstitutor(), TypeSubstitutorMarker {
+    return object : AbstractConeSubstitutor(context), 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/resolve/transformers/FirSyntheticCallGenerator.kt

@@ -298,7 +298,8 @@ class FirSyntheticCallGenerator(
 
         val returnType = rightArgumentType.resolvedTypeFromPrototype(
             rightArgumentType.type.withAttributes(
-                ConeAttributes.create(listOf(CompilerConeAttributes.Exact))
+                ConeAttributes.create(listOf(CompilerConeAttributes.Exact)),
+                session.typeContext,
             )
         )
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -517,6 +517,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                         returnExpression.resultType.approximatedIfNeededOrSelf(
                             inferenceComponents.approximator,
                             simpleFunction?.visibilityForApproximation(),
+                            transformer.session.typeContext,
                             simpleFunction?.isInline == true
                         )
                     )
@@ -843,7 +844,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                     withExpectedType(
                         expectedType.approximatedIfNeededOrSelf(
                             inferenceComponents.approximator,
-                            variable.visibilityForApproximation()
+                            variable.visibilityForApproximation(),
+                            inferenceComponents.session.typeContext,
                         )
                     )
                 )

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

@@ -17,9 +17,9 @@ import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.createTypeSubstitutorByTypeConstructor
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
-import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
+import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.DFS
@@ -258,12 +258,12 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
 
     override fun KotlinTypeMarker.removeAnnotations(): KotlinTypeMarker {
         require(this is ConeKotlinType)
-        return withAttributes(ConeAttributes.Empty)
+        return withAttributes(ConeAttributes.Empty, this@ConeInferenceContext)
     }
 
     override fun SimpleTypeMarker.replaceArguments(newArguments: List<TypeArgumentMarker>): SimpleTypeMarker {
         require(this is ConeKotlinType)
-        return this.withArguments(newArguments.cast<List<ConeTypeProjection>>().toTypedArray())
+        return this.withArguments(newArguments.cast<List<ConeTypeProjection>>().toTypedArray(), this@ConeInferenceContext)
     }
 
     override fun KotlinTypeMarker.hasExactAnnotation(): Boolean {
@@ -364,7 +364,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
 
     override fun KotlinTypeMarker.removeExactAnnotation(): KotlinTypeMarker {
         require(this is ConeKotlinType)
-        return withAttributes(attributes.remove(CompilerConeAttributes.Exact))
+        return withAttributes(attributes.remove(CompilerConeAttributes.Exact), this@ConeInferenceContext)
     }
 
     override fun TypeConstructorMarker.toErrorType(): SimpleTypeMarker {

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

@@ -415,7 +415,7 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
     }
 
     private fun FirTypeParameterSymbol.allBoundsAreNullable(): Boolean {
-        return fir.bounds.all { it.coneType.isMarkedNullable }
+        return fir.bounds.all { it.coneType.isNullableType() }
     }
 
     private fun TypeConstructorMarker.toFirRegularClass(): FirRegularClass? {

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

@@ -47,23 +47,32 @@ fun ConeInferenceContext.intersectTypesOrNull(types: List<ConeKotlinType>): Cone
 fun TypeCheckerProviderContext.equalTypes(a: ConeKotlinType, b: ConeKotlinType): Boolean =
     AbstractTypeChecker.equalTypes(this, a, b)
 
-fun ConeDefinitelyNotNullType.Companion.create(original: ConeKotlinType): ConeDefinitelyNotNullType? {
+fun ConeDefinitelyNotNullType.Companion.create(
+    original: ConeKotlinType,
+    typeContext: ConeTypeContext,
+    useCorrectedNullabilityForFlexibleTypeParameters: Boolean = false,
+): ConeDefinitelyNotNullType? {
     return when {
         original is ConeDefinitelyNotNullType -> original
-        makesSenseToBeDefinitelyNotNull(original) ->
+        typeContext
-            ConeDefinitelyNotNullType(original.lowerBoundIfFlexible())
+            .newBaseTypeCheckerContext(errorTypesEqualToAnything = false, stubTypesEqualToAnything = false)
+            .makesSenseToBeDefinitelyNotNull(original, useCorrectedNullabilityForFlexibleTypeParameters) ->
+
+            ConeDefinitelyNotNullType(
+                original.lowerBoundIfFlexible()
+            )
         else -> null
     }
 }
 
-fun ConeKotlinType.makeConeTypeDefinitelyNotNullOrNotNull(typeContext: ConeInferenceContext): ConeKotlinType {
+fun ConeKotlinType.makeConeTypeDefinitelyNotNullOrNotNull(typeContext: ConeTypeContext): ConeKotlinType {
     if (this is ConeIntersectionType) {
         return ConeIntersectionType(intersectedTypes.map { it.makeConeTypeDefinitelyNotNullOrNotNull(typeContext) })
     }
-    return ConeDefinitelyNotNullType.create(this) ?: this.withNullability(ConeNullability.NOT_NULL, typeContext)
+    return ConeDefinitelyNotNullType.create(this, typeContext) ?: this.withNullability(ConeNullability.NOT_NULL, typeContext)
 }
 
-fun <T : ConeKotlinType> T.withArguments(arguments: Array<out ConeTypeProjection>): T {
+fun <T : ConeKotlinType> T.withArguments(arguments: Array<out ConeTypeProjection>, typeSystemContext: ConeTypeContext): T {
     if (this.typeArguments === arguments) {
         return this
     }
@@ -72,12 +81,15 @@ fun <T : ConeKotlinType> T.withArguments(arguments: Array<out ConeTypeProjection
     return when (this) {
         is ConeClassErrorType -> this
         is ConeClassLikeTypeImpl -> ConeClassLikeTypeImpl(lookupTag, arguments, nullability.isNullable) as T
-        is ConeDefinitelyNotNullType -> ConeDefinitelyNotNullType.create(original.withArguments(arguments))!! as T
+        is ConeDefinitelyNotNullType -> ConeDefinitelyNotNullType.create(
+            original.withArguments(arguments, typeSystemContext),
+            typeSystemContext
+        )!! as T
         else -> error("Not supported: $this: ${this.render()}")
     }
 }
 
-fun <T : ConeKotlinType> T.withAttributes(attributes: ConeAttributes): T {
+fun <T : ConeKotlinType> T.withAttributes(attributes: ConeAttributes, typeSystemContext: ConeTypeContext): T {
     if (this.attributes == attributes) {
         return this
     }
@@ -86,16 +98,22 @@ fun <T : ConeKotlinType> T.withAttributes(attributes: ConeAttributes): T {
     return when (this) {
         is ConeClassErrorType -> this
         is ConeClassLikeTypeImpl -> ConeClassLikeTypeImpl(lookupTag, typeArguments, nullability.isNullable, attributes)
-        is ConeDefinitelyNotNullType -> ConeDefinitelyNotNullType.create(original.withAttributes(attributes))!!
+        is ConeDefinitelyNotNullType -> ConeDefinitelyNotNullType.create(
+            original.withAttributes(attributes, typeSystemContext),
+            typeSystemContext
+        )!!
         is ConeTypeParameterTypeImpl -> ConeTypeParameterTypeImpl(lookupTag, nullability.isNullable, attributes)
-        is ConeFlexibleType -> ConeFlexibleType(lowerBound.withAttributes(attributes), upperBound.withAttributes(attributes))
+        is ConeFlexibleType -> ConeFlexibleType(
+            lowerBound.withAttributes(attributes, typeSystemContext),
+            upperBound.withAttributes(attributes, typeSystemContext)
+        )
         else -> error("Not supported: $this: ${this.render()}")
     } as T
 }
 
 fun <T : ConeKotlinType> T.withNullability(
     nullability: ConeNullability,
-    typeContext: ConeInferenceContext,
+    typeContext: ConeTypeContext,
     attributes: ConeAttributes = this.attributes,
 ): T {
     if (this.nullability == nullability && this.attributes == attributes) {
@@ -140,7 +158,7 @@ fun <T : ConeKotlinType> T.withNullability(
 }
 
 fun coneFlexibleOrSimpleType(
-    typeContext: ConeInferenceContext,
+    typeContext: ConeTypeContext,
     lowerBound: ConeKotlinType,
     upperBound: ConeKotlinType,
 ): ConeKotlinType {
@@ -197,7 +215,7 @@ fun FirTypeRef.isUnsafeVarianceType(session: FirSession): Boolean {
 fun FirTypeRef.hasEnhancedNullability(): Boolean =
     coneTypeSafe<ConeKotlinType>()?.hasEnhancedNullability == true
 
-fun FirTypeRef.withoutEnhancedNullability(): FirTypeRef {
+fun FirTypeRef.withoutEnhancedNullability(typeSystemContext: ConeTypeContext): FirTypeRef {
     require(this is FirResolvedTypeRef)
     if (!hasEnhancedNullability()) return this
     return buildResolvedTypeRef {
@@ -205,7 +223,8 @@ fun FirTypeRef.withoutEnhancedNullability(): FirTypeRef {
         type = this@withoutEnhancedNullability.type.withAttributes(
             ConeAttributes.create(
                 this@withoutEnhancedNullability.type.attributes.filter { it != CompilerConeAttributes.EnhancedNullability }
-            )
+            ),
+            typeSystemContext,
         )
         annotations += this@withoutEnhancedNullability.annotations
     }
@@ -263,13 +282,14 @@ fun FirTypeRef.approximated(
 fun FirTypeRef.approximatedIfNeededOrSelf(
     approximator: ConeTypeApproximator,
     containingCallableVisibility: Visibility?,
-    isInlineFunction: Boolean = false
+    typeSystemContext: ConeTypeContext,
+    isInlineFunction: Boolean = false,
 ): FirTypeRef {
     val approximated = if (containingCallableVisibility == Visibilities.Public || containingCallableVisibility == Visibilities.Protected)
         approximatedForPublicPosition(approximator)
     else
         this
-    return approximated.hideLocalTypeIfNeeded(containingCallableVisibility, isInlineFunction).withoutEnhancedNullability()
+    return approximated.hideLocalTypeIfNeeded(containingCallableVisibility, isInlineFunction).withoutEnhancedNullability(typeSystemContext)
 }
 
 fun FirTypeRef.approximatedForPublicPosition(approximator: ConeTypeApproximator): FirTypeRef =
@@ -334,11 +354,11 @@ fun ConeTypeContext.captureFromArgumentsInternal(type: ConeKotlinType, status: C
     val capturedArguments = captureArguments(type, status) ?: return null
     return if (type is ConeFlexibleType) {
         ConeFlexibleType(
-            type.lowerBound.withArguments(capturedArguments),
+            type.lowerBound.withArguments(capturedArguments, this),
-            type.upperBound.withArguments(capturedArguments),
+            type.upperBound.withArguments(capturedArguments, this),
         )
     } else {
-        type.withArguments(capturedArguments)
+        type.withArguments(capturedArguments, this)
     }
 }
 
@@ -418,12 +438,12 @@ fun ConeTypeContext.captureFromExpressionInternal(type: ConeKotlinType): ConeKot
             typeToReplace.intersectedTypes.map { componentType ->
                 val capturedArguments = findCorrespondingCapturedArgumentsForType(componentType)
                     ?: return@map componentType
-                componentType.withArguments(capturedArguments)
+                componentType.withArguments(capturedArguments, this)
             }
         } else {
             val capturedArguments = findCorrespondingCapturedArgumentsForType(typeToReplace)
                 ?: return listOf(typeToReplace)
-            listOf(typeToReplace.withArguments(capturedArguments))
+            listOf(typeToReplace.withArguments(capturedArguments, this))
         }
     }
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirTypeUtils.kt

@@ -10,9 +10,11 @@ import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirConstExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.render
-import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlin.types.AbstractNullabilityChecker
+import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 import org.jetbrains.kotlin.types.ConstantValueKind
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
@@ -137,15 +139,37 @@ fun FirTypeProjection.toConeTypeProjection(): ConeTypeProjection =
         else -> error("!")
     }
 
-fun makesSenseToBeDefinitelyNotNull(type: ConeKotlinType): Boolean =
+fun AbstractTypeCheckerContext.makesSenseToBeDefinitelyNotNull(
-    type.canHaveUndefinedNullability() // TODO: also check nullability
+    type: ConeKotlinType,
+    useCorrectedNullabilityForFlexibleTypeParameters: Boolean
+): Boolean {
+    if (!type.canHaveUndefinedNullability()) return false
+
+    // Replacing `useCorrectedNullabilityForFlexibleTypeParameters` with true for all call-sites seems to be correct
+    // But it seems that it should be a new feature: KT-28785 would be automatically fixed then
+    // (see the tests org.jetbrains.kotlin.spec.checkers.DiagnosticsTestSpecGenerated.NotLinked.Dfa.Pos.test12/13)
+    // So it should be a language feature, but it's hard correctly identify language version settings for all call sites
+    // Thus, we have non-trivial value at org.jetbrains.kotlin.load.java.typeEnhancement.JavaTypeEnhancement.notNullTypeParameter
+    // that run under related language-feature only
+    if (useCorrectedNullabilityForFlexibleTypeParameters && type is ConeTypeParameterType) {
+        // Effectively checks if the type is flexible or has nullable bound
+        return with(typeSystemContext) {
+            type.isNullableType()
+        }
+    }
+
+    // Actually, this code should work for type parameters as well, but it breaks some cases
+    // See KT-40114
+
+    return !AbstractNullabilityChecker.isSubtypeOfAny(this, type)
+}
 
 fun ConeKotlinType.canHaveUndefinedNullability(): Boolean {
     return when (this) {
         is ConeTypeVariableType,
-        is ConeCapturedType
+        is ConeCapturedType,
+        is ConeTypeParameterType
         -> true
-        is ConeTypeParameterType -> type.isMarkedNullable || !hasNotNullUpperBound()
         else -> false
     }
 }

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

@@ -26,6 +26,7 @@ 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
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclaration
@@ -79,11 +80,8 @@ internal fun KtNamedClassOrObjectSymbol.typeForClassSymbol(psiElement: PsiElemen
 
 private class AnonymousTypesSubstitutor(
     private val session: FirSession,
-    private val state: FirModuleResolveState
+    private val state: FirModuleResolveState,
-) : AbstractConeSubstitutor() {
+) : AbstractConeSubstitutor(session.typeContext) {
-    override val typeInferenceContext: ConeInferenceContext
-        get() = session.inferenceComponents.ctx
-
     override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
 
         if (type !is ConeClassLikeType) return null