FIR: use withNullability with type context whenever possible

It could simplify flexible type cases.

compiler/fir/dump/src/org/jetbrains/kotlin/fir/dump/HtmlFirDump.kt

@@ -676,7 +676,7 @@ class HtmlFirDump internal constructor(private var linkResolver: FirLinkResolver
             AbstractStrictEqualityTypeChecker.strictEqualTypes(
                 session.typeContext,
                 flexibleType.lowerBound,
-                flexibleType.upperBound.withNullability(ConeNullability.NOT_NULL)
+                flexibleType.upperBound.withNullability(ConeNullability.NOT_NULL, session.typeContext)
             )
         ) {
             generate(flexibleType.lowerBound)

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

@@ -582,7 +582,7 @@ class Fir2IrVisitor(
                 IrGetValueImpl(startOffset, endOffset, irLhsVariable.type, irLhsVariable.symbol)
 
             val originalType = firLhsVariable.returnTypeRef.coneType
-            val notNullType = originalType.withNullability(ConeNullability.NOT_NULL)
+            val notNullType = originalType.withNullability(ConeNullability.NOT_NULL, session.typeContext)
             val irBranches = listOf(
                 IrBranchImpl(
                     startOffset, endOffset, primitiveOp2(

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.declarations.expandedConeType
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.utils.WeakPair
@@ -58,14 +59,17 @@ fun ConeClassLikeType.directExpansionType(
     val typeAliasSymbol = lookupTag.toSymbol(useSiteSession) as? FirTypeAliasSymbol ?: return null
     val typeAlias = typeAliasSymbol.fir
 
-    val resultType = expandedConeType(typeAlias)?.applyNullabilityFrom(this) ?: return null
+    val resultType = expandedConeType(typeAlias)?.applyNullabilityFrom(useSiteSession, this) ?: return null
 
     if (resultType.typeArguments.isEmpty()) return resultType
     return mapTypeAliasArguments(typeAlias, this, resultType) as? ConeClassLikeType
 }
 
-private fun ConeClassLikeType.applyNullabilityFrom(abbreviation: ConeClassLikeType): ConeClassLikeType {
+private fun ConeClassLikeType.applyNullabilityFrom(
-    if (abbreviation.isMarkedNullable) return withNullability(ConeNullability.NULLABLE)
+    session: FirSession,
+    abbreviation: ConeClassLikeType
+): ConeClassLikeType {
+    if (abbreviation.isMarkedNullable) return withNullability(ConeNullability.NULLABLE, session.typeContext)
     return this
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CreateFreshTypeVariableSubstitutorStage.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.resolve.inference.model.ConeDeclaredUpperBoundCo
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
+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.FirTypePlaceholderProjection
@@ -68,7 +69,7 @@ internal object CreateFreshTypeVariableSubstitutorStage : ResolutionStage() {
                     getTypePreservingFlexibilityWrtTypeVariable(
                         typeArgument.typeRef.coneType,
                         typeParameter,
-                        context.session.inferenceComponents.ctx
+                        context.session
                     ).fullyExpandedType(context.session),
                     SimpleConstraintSystemConstraintPosition // TODO
                 )
@@ -88,11 +89,11 @@ internal object CreateFreshTypeVariableSubstitutorStage : ResolutionStage() {
     private fun getTypePreservingFlexibilityWrtTypeVariable(
         type: ConeKotlinType,
         typeParameter: FirTypeParameterRef,
-        context: ConeTypeContext
+        session: FirSession,
     ): ConeKotlinType {
-        return if (typeParameter.shouldBeFlexible(context)) {
+        return if (typeParameter.shouldBeFlexible(session.typeContext)) {
-            val notNullType = type.withNullability(ConeNullability.NOT_NULL)
+            val notNullType = type.withNullability(ConeNullability.NOT_NULL, session.typeContext)
-            ConeFlexibleType(notNullType, notNullType.withNullability(ConeNullability.NULLABLE))
+            ConeFlexibleType(notNullType, notNullType.withNullability(ConeNullability.NULLABLE, session.typeContext))
         } else {
             type
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Synthetics.kt

@@ -91,8 +91,8 @@ class FirSyntheticPropertiesScope(
                     // I think details here are worth designing
                     if (!AbstractTypeChecker.isSubtypeOf(
                             session.typeContext,
-                            getterReturnType.withNullability(NOT_NULL),
+                            getterReturnType.withNullability(NOT_NULL, session.typeContext),
-                            parameterType.withNullability(NOT_NULL)
+                            parameterType.withNullability(NOT_NULL, session.typeContext)
                         )
                     ) {
                         return

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -6,8 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.config.LanguageFeature
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.description.ConeBooleanConstantReference
 import org.jetbrains.kotlin.fir.contracts.description.ConeConditionalEffectDeclaration
@@ -15,7 +14,6 @@ import org.jetbrains.kotlin.fir.contracts.description.ConeConstantReference
 import org.jetbrains.kotlin.fir.contracts.description.ConeReturnsEffectDeclaration
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
@@ -33,7 +31,6 @@ import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.visibilityChecker
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -777,7 +774,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
         safeCall.receiver.let { receiver ->
             val type = receiver.coneType.takeIf { it.isMarkedNullable }
-                ?.withNullability(ConeNullability.NOT_NULL)
+                ?.withNullability(ConeNullability.NOT_NULL, components.session.typeContext)
                 ?: return@let
 
             val variable = variableStorage.getOrCreateVariable(flow, receiver)
@@ -998,7 +995,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
         if (isAssignment) {
             if (initializer is FirConstExpression<*> && initializer.kind == ConstantValueKind.Null) {
-                flow.addTypeStatement(propertyVariable typeEq property.returnTypeRef.coneType.withNullability(ConeNullability.NULLABLE))
+                flow.addTypeStatement(
+                    propertyVariable typeEq
+                            property.returnTypeRef.coneType.withNullability(ConeNullability.NULLABLE, components.session.typeContext)
+                )
             } else {
                 flow.addTypeStatement(propertyVariable typeEq initializer.typeRef.coneType)
             }

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.FirWhenExhaustivenessTransformer
 import org.jetbrains.kotlin.fir.resolve.withExpectedType
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.visitors.transformSingle
@@ -217,7 +218,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         elvisExpression.transformAnnotations(transformer, data)
 
         val expectedType = data.expectedType?.coneTypeSafe<ConeKotlinType>()
-        val resolutionModeForLhs = withExpectedType(expectedType?.withNullability(ConeNullability.NULLABLE))
+        val resolutionModeForLhs = withExpectedType(expectedType?.withNullability(ConeNullability.NULLABLE, session.typeContext))
         elvisExpression.transformLhs(transformer, resolutionModeForLhs)
         dataFlowAnalyzer.exitElvisLhs(elvisExpression)
 

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

@@ -113,7 +113,11 @@ fun <T : ConeKotlinType> T.withNullability(
                     return this
                 }
             }
-            coneFlexibleOrSimpleType(typeContext, lowerBound.withNullability(nullability), upperBound.withNullability(nullability))
+            coneFlexibleOrSimpleType(
+                typeContext,
+                lowerBound.withNullability(nullability, typeContext),
+                upperBound.withNullability(nullability, typeContext)
+            )
         }
         is ConeTypeVariableType -> ConeTypeVariableType(nullability, lookupTag)
         is ConeCapturedType -> ConeCapturedType(captureStatus, lowerType, nullability, constructor, attributes)
@@ -127,8 +131,8 @@ fun <T : ConeKotlinType> T.withNullability(
         is ConeStubType -> ConeStubType(variable, nullability)
         is ConeDefinitelyNotNullType -> when (nullability) {
             ConeNullability.NOT_NULL -> this
-            ConeNullability.NULLABLE -> original.withNullability(nullability)
+            ConeNullability.NULLABLE -> original.withNullability(nullability, typeContext)
-            ConeNullability.UNKNOWN -> original.withNullability(nullability)
+            ConeNullability.UNKNOWN -> original.withNullability(nullability, typeContext)
         }
         is ConeIntegerLiteralType -> ConeIntegerLiteralTypeImpl(value, isUnsigned, nullability)
         else -> error("sealed: ${this::class}")

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -412,7 +412,11 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
                     val lowerRendered = lowerBound.tryToRenderConeAsFunctionType()
                     if (lowerBound.nullability == ConeNullability.NOT_NULL && upperBound.nullability == ConeNullability.NULLABLE &&
                         AbstractStrictEqualityTypeChecker
-                            .strictEqualTypes(session.typeContext, lowerBound, upperBound.withNullability(ConeNullability.NOT_NULL))
+                            .strictEqualTypes(
+                                session.typeContext,
+                                lowerBound,
+                                upperBound.withNullability(ConeNullability.NOT_NULL, session.typeContext)
+                            )
                     ) {
                         "$lowerRendered!"
                     } else {