[FIR] Implement capturing of captured types

This is in preparation of a future commit, where captured types won't be
approximated anymore after completion.

Consider a case like

class Box<T>(val value: T)
interface Foo<T> { fun bar() }

fun test(x: Box<out Foo<*>>) {
    x.value.bar()
}

The type of `x.value` will be `CapturedType(out Foo<*>)`.
Note that capturing only applies to the top level, i.e., nested
projections are not captured.
That's why it becomes necessary to support capturing of captured types,
otherwise the star projection in `CapturedType(out Foo<*>)` is not
properly captured in the receiver of the call `x.value.bar()`.

#KT-62959

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirUpperBoundViolatedHelpers.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
+import org.jetbrains.kotlin.fir.resolve.substitution.wrapProjection
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.resolve.withCombinedAttributesFrom
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol

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

@@ -24,18 +24,48 @@ import org.jetbrains.kotlin.types.model.typeConstructor
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 
-abstract class AbstractConeSubstitutor(protected val typeContext: ConeTypeContext) : ConeSubstitutor() {
-    protected fun wrapProjection(old: ConeTypeProjection, newType: ConeKotlinType): ConeTypeProjection {
-        return when (old) {
-            is ConeStarProjection -> old
-            is ConeKotlinTypeProjectionIn -> ConeKotlinTypeProjectionIn(newType)
-            is ConeKotlinTypeProjectionOut -> ConeKotlinTypeProjectionOut(newType)
-            is ConeKotlinTypeConflictingProjection -> ConeKotlinTypeConflictingProjection(newType)
-            is ConeKotlinType -> newType
-            else -> old
-        }
-    }
+inline fun ConeCapturedType.substitute(f: (ConeKotlinType) -> ConeKotlinType?): ConeCapturedType? {
+    val innerType = this.lowerType ?: this.constructor.projection.type
+    // TODO(KT-64024): This early return looks suspicious.
+    //  In fact, if the inner type wasn't substituted we will ignore potential substitution in
+    //  super types
+    val substitutedInnerType = innerType?.let(f) ?: return null
+    if (substitutedInnerType is ConeCapturedType) return substitutedInnerType
+    val substitutedSuperTypes =
+        this.constructor.supertypes?.map { f(it) ?: it }
 
+    // TODO(KT-64027): Creation of new captured types creates unexpected behavior by breaking substitution consistency.
+    //  E.g:
+    //  ```
+    //   substitution = { A => B }
+    //   substituteOrSelf(C<CapturedType(out A)_0>) -> C<CapturedType(out B)_1>
+    //   substituteOrSelf(C<CapturedType(out A)_0>) -> C<CapturedType(out B)_2>
+    //   C<CapturedType(out B)_1> <!:> C<CapturedType(out B)_2>
+    //  ```
+
+    return ConeCapturedType(
+        captureStatus,
+        constructor = ConeCapturedTypeConstructor(
+            wrapProjection(constructor.projection, substitutedInnerType),
+            substitutedSuperTypes,
+            typeParameterMarker = constructor.typeParameterMarker
+        ),
+        lowerType = if (lowerType != null) substitutedInnerType else null
+    )
+}
+
+fun wrapProjection(old: ConeTypeProjection, newType: ConeKotlinType): ConeTypeProjection {
+    return when (old) {
+        is ConeStarProjection -> old
+        is ConeKotlinTypeProjectionIn -> ConeKotlinTypeProjectionIn(newType)
+        is ConeKotlinTypeProjectionOut -> ConeKotlinTypeProjectionOut(newType)
+        is ConeKotlinTypeConflictingProjection -> ConeKotlinTypeConflictingProjection(newType)
+        is ConeKotlinType -> newType
+        else -> old
+    }
+}
+
+abstract class AbstractConeSubstitutor(protected val typeContext: ConeTypeContext) : ConeSubstitutor() {
     abstract fun substituteType(type: ConeKotlinType): ConeKotlinType?
     open fun substituteArgument(projection: ConeTypeProjection, index: Int): ConeTypeProjection? {
         val type = (projection as? ConeKotlinTypeProjection)?.type ?: return null
@@ -84,7 +114,7 @@ abstract class AbstractConeSubstitutor(protected val typeContext: ConeTypeContex
                 if (it.lowerBound == it.upperBound) it.lowerBound
                 else it
             }
-            is ConeCapturedType -> return substituteCapturedType()
+            is ConeCapturedType -> return substitute(::substituteOrNull)
             is ConeDefinitelyNotNullType -> this.substituteOriginal()
             is ConeIntersectionType -> this.substituteIntersectedTypes()
             is ConeStubType -> return null
@@ -92,36 +122,6 @@ abstract class AbstractConeSubstitutor(protected val typeContext: ConeTypeContex
         }
     }
 
-    private fun ConeCapturedType.substituteCapturedType(): ConeCapturedType? {
-        val innerType = this.lowerType ?: this.constructor.projection.type
-        // TODO(KT-64024): This early return looks suspicious.
-        //  In fact, if the inner type wasn't substituted we will ignore potential substitution in
-        //  super types
-        val substitutedInnerType = substituteOrNull(innerType) ?: return null
-        if (substitutedInnerType is ConeCapturedType) return substitutedInnerType
-        val substitutedSuperTypes =
-            this.constructor.supertypes?.map { substituteOrSelf(it) }
-
-        // TODO(KT-64027): Creation of new captured types creates unexpected behavior by breaking substitution consistency.
-        //  E.g:
-        //  ```
-        //   substitution = { A => B }
-        //   substituteOrSelf(C<CapturedType(out A)_0>) -> C<CapturedType(out B)_1>
-        //   substituteOrSelf(C<CapturedType(out A)_0>) -> C<CapturedType(out B)_2>
-        //   C<CapturedType(out B)_1> <!:> C<CapturedType(out B)_2>
-        //  ```
-
-        return ConeCapturedType(
-            captureStatus,
-            constructor = ConeCapturedTypeConstructor(
-                wrapProjection(constructor.projection, substitutedInnerType),
-                substitutedSuperTypes,
-                typeParameterMarker = constructor.typeParameterMarker
-            ),
-            lowerType = if (lowerType != null) substitutedInnerType else null
-        )
-    }
-
     private fun ConeIntersectionType.substituteIntersectedTypes(): ConeIntersectionType? {
         val substitutedTypes = ArrayList<ConeKotlinType>(intersectedTypes.size)
         var somethingIsSubstituted = false

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

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.modality
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.diagnostics.ConeRecursiveTypeParameterDuringErasureError
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.fir.resolve.substitution.substitute
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
@@ -431,6 +432,14 @@ fun ConeTypeContext.captureArguments(type: ConeKotlinType, status: CaptureStatus
 }
 
 internal fun ConeTypeContext.captureFromExpressionInternal(type: ConeKotlinType): ConeKotlinType? {
+    if (type is ConeCapturedType) {
+        return type.substitute { captureFromExpressionInternal(it) }
+    }
+
+    if (type is ConeDefinitelyNotNullType) {
+        return captureFromExpressionInternal(type.original)?.makeConeTypeDefinitelyNotNullOrNotNull(this)
+    }
+
     if (type !is ConeIntersectionType && type !is ConeFlexibleType) {
         return captureFromArgumentsInternal(type, CaptureStatus.FROM_EXPRESSION)
     }