FIR: Rewrite Java type mapping

Previoisly, there were two places where mapping had happened: - toConeKotlinTypeWithNullability - enhancePossiblyFlexible The first one was used for supertypes and bounds and the second one was used for other signature parts The main idea is to perform type mapping once to a flexible type, and then use it as it's needed (it's lower bound, or for the further ehnancement) Also, this commit fixes flexibility for type arguments, see the tests
2020-02-18 19:53:40 +03:00
parent 748a326104
commit 2ad8488e6a
42 changed files with 452 additions and 358 deletions
@@ -92,7 +92,12 @@ class JavaSymbolProvider(
        stack: JavaTypeParameterStack,
    ) {
        for (upperBound in javaTypeParameter.upperBounds) {
-            bounds += upperBound.toFirResolvedTypeRef(this@JavaSymbolProvider.session, stack, nullability = ConeNullability.UNKNOWN)
+            bounds += upperBound.toFirResolvedTypeRef(
+                this@JavaSymbolProvider.session,
+                stack,
+                isForSupertypes = false,
+                forTypeParameterBounds = true
+            )
        }
        addDefaultBoundIfNecessary()
    }
@@ -260,7 +265,7 @@ class JavaSymbolProvider(
                firJavaClass.replaceSuperTypeRefs(
                    javaClass.supertypes.map { supertype ->
                        supertype.toFirResolvedTypeRef(
-                            this@JavaSymbolProvider.session, javaTypeParameterStack, typeParametersNullability = ConeNullability.UNKNOWN
+                            this@JavaSymbolProvider.session, javaTypeParameterStack, isForSupertypes = true, forTypeParameterBounds = false
                        )
                    }
                )
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
@@ -24,8 +25,11 @@ import org.jetbrains.kotlin.fir.java.declarations.buildJavaValueParameter
 import org.jetbrains.kotlin.fir.java.enhancement.readOnlyToMutable
 import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.getClassDeclaredCallableSymbols
+import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
@@ -38,6 +42,7 @@ import org.jetbrains.kotlin.fir.types.jvm.FirJavaTypeRef
 import org.jetbrains.kotlin.fir.types.jvm.buildJavaTypeRef
 import org.jetbrains.kotlin.load.java.structure.*
 import org.jetbrains.kotlin.load.java.structure.impl.JavaElementImpl
+import org.jetbrains.kotlin.load.java.typeEnhancement.TypeComponentPosition
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance.*
@@ -65,24 +70,17 @@ internal fun ClassId.toConeKotlinType(
    return ConeClassLikeTypeImpl(lookupTag, typeArguments, isNullable)
 }

-internal fun FirTypeRef.toNotNullConeKotlinType(
+internal fun FirTypeRef.toConeKotlinTypeProbablyFlexible(
    session: FirSession, javaTypeParameterStack: JavaTypeParameterStack
 ): ConeKotlinType =
    when (this) {
        is FirResolvedTypeRef -> type
        is FirJavaTypeRef -> {
-            val javaType = type
-            javaType.toNotNullConeKotlinType(session, javaTypeParameterStack)
+            type.toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack)
        }
        else -> ConeKotlinErrorType("Unexpected type reference in JavaClassUseSiteMemberScope: ${this::class.java}")
    }

-internal fun JavaType?.toNotNullConeKotlinType(
-    session: FirSession, javaTypeParameterStack: JavaTypeParameterStack
-): ConeKotlinType {
-    return toConeKotlinTypeWithNullability(session, javaTypeParameterStack, ConeNullability.NOT_NULL)
-}
-
 internal fun JavaType.toFirJavaTypeRef(session: FirSession, javaTypeParameterStack: JavaTypeParameterStack): FirJavaTypeRef {
    val annotations = (this as? JavaClassifierType)?.annotations.orEmpty()
    return buildJavaTypeRef {
@@ -94,25 +92,29 @@ internal fun JavaType.toFirJavaTypeRef(session: FirSession, javaTypeParameterSta
 internal fun JavaClassifierType.toFirResolvedTypeRef(
    session: FirSession,
    javaTypeParameterStack: JavaTypeParameterStack,
-    nullability: ConeNullability = ConeNullability.NOT_NULL,
-    typeParametersNullability: ConeNullability = ConeNullability.NOT_NULL
+    isForSupertypes: Boolean,
+    forTypeParameterBounds: Boolean
 ): FirResolvedTypeRef {
-    val coneType = this.toConeKotlinTypeWithNullability(session, javaTypeParameterStack, nullability, typeParametersNullability)
+    val coneType =
+        this.toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack, forTypeParameterBounds).run {
+            if (isForSupertypes)
+                this.lowerBoundIfFlexible()
+            else
+                this
+        }
    return buildResolvedTypeRef {
        type = coneType
        this@toFirResolvedTypeRef.annotations.mapTo(annotations) { it.toFirAnnotationCall(session, javaTypeParameterStack) }
    }
 }

-internal fun JavaType?.toConeKotlinTypeWithNullability(
+internal fun JavaType?.toConeKotlinTypeWithoutEnhancement(
    session: FirSession,
-    javaTypeParameterStack: JavaTypeParameterStack,
-    nullability: ConeNullability,
-    typeParametersNullability: ConeNullability = ConeNullability.NOT_NULL
+    javaTypeParameterStack: JavaTypeParameterStack
 ): ConeKotlinType {
    return when (this) {
        is JavaClassifierType -> {
-            toConeKotlinTypeWithNullability(session, nullability, typeParametersNullability, javaTypeParameterStack)
+            toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack)
        }
        is JavaPrimitiveType -> {
            val primitiveType = type
@@ -120,77 +122,204 @@ internal fun JavaType?.toConeKotlinTypeWithNullability(
                null -> "Unit"
                else -> javaName.capitalize()
            }
+
            val classId = StandardClassIds.byName(kotlinPrimitiveName)
-            classId.toConeKotlinType(emptyArray(), nullability.isNullable)
+            classId.toConeKotlinType(emptyArray(), isNullable = false)
        }
        is JavaArrayType -> {
            val componentType = componentType
            if (componentType !is JavaPrimitiveType) {
                val classId = StandardClassIds.Array
-                val argumentType = ConeFlexibleType(
-                    componentType.toConeKotlinTypeWithNullability(session, javaTypeParameterStack, ConeNullability.NOT_NULL),
-                    componentType.toConeKotlinTypeWithNullability(session, javaTypeParameterStack, ConeNullability.NULLABLE)
+                val argumentType = componentType.toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack)
+                classId.toConeFlexibleType(
+                    arrayOf(argumentType),
+                    typeArgumentsForUpper = arrayOf(ConeKotlinTypeProjectionOut(argumentType))
                )
-                classId.toConeKotlinType(arrayOf(argumentType), nullability.isNullable)
            } else {
                val javaComponentName = componentType.type?.typeName?.asString()?.capitalize() ?: error("Array of voids")
                val classId = StandardClassIds.byName(javaComponentName + "Array")
-                classId.toConeKotlinType(emptyArray(), nullability.isNullable)
+
+                classId.toConeFlexibleType(emptyArray())
            }
        }
-        is JavaWildcardType -> bound?.toNotNullConeKotlinType(session, javaTypeParameterStack) ?: run {
-            val classId = StandardClassIds.Any
-            classId.toConeKotlinType(emptyArray(), nullability.isNullable)
+        is JavaWildcardType -> bound?.toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack) ?: run {
+            StandardClassIds.Any.toConeFlexibleType(emptyArray())
        }
        null -> {
-            val classId = StandardClassIds.Any
-            classId.toConeKotlinType(emptyArray(), nullability.isNullable)
+            StandardClassIds.Any.toConeFlexibleType(emptyArray())
        }
        else -> error("Strange JavaType: ${this::class.java}")
    }
 }

-internal fun JavaClassifierType.toConeKotlinTypeWithNullability(
+private fun ClassId.toConeFlexibleType(
+    typeArguments: Array<ConeKotlinTypeProjection>,
+    typeArgumentsForUpper: Array<ConeKotlinTypeProjection> = typeArguments
+) = ConeFlexibleType(
+    toConeKotlinType(typeArguments, isNullable = false),
+    toConeKotlinType(typeArgumentsForUpper, isNullable = true)
+)
+
+private fun JavaClassifierType.toConeKotlinTypeWithoutEnhancement(
    session: FirSession,
-    nullability: ConeNullability,
-    typeParametersNullability: ConeNullability,
-    javaTypeParameterStack: JavaTypeParameterStack
+    javaTypeParameterStack: JavaTypeParameterStack,
+    forTypeParameterBounds: Boolean = false
 ): ConeKotlinType {
-    if (nullability == ConeNullability.UNKNOWN) {
-        return ConeFlexibleType(
-            toConeKotlinTypeWithNullability(session, ConeNullability.NOT_NULL, typeParametersNullability, javaTypeParameterStack),
-            toConeKotlinTypeWithNullability(session, ConeNullability.NULLABLE, typeParametersNullability, javaTypeParameterStack)
-        )
+    val lowerBound = toConeKotlinTypeForFlexibleBound(session, javaTypeParameterStack, isLowerBound = true, forTypeParameterBounds)
+    val upperBound = toConeKotlinTypeForFlexibleBound(session, javaTypeParameterStack, isLowerBound = false, forTypeParameterBounds)
+
+    return if (isRaw) ConeRawType(lowerBound, upperBound) else ConeFlexibleType(lowerBound, upperBound)
+}
+
+private fun computeRawProjection(
+    session: FirSession,
+    parameter: FirTypeParameter,
+    attr: TypeComponentPosition,
+    erasedUpperBound: ConeKotlinType = parameter.getErasedUpperBound()
+) = when (attr) {
+    // Raw(List<T>) => (List<Any?>..List<*>)
+    // Raw(Enum<T>) => (Enum<Enum<*>>..Enum<out Enum<*>>)
+    // In the last case upper bound is equal to star projection `Enum<*>`,
+    // but we want to keep matching tree structure of flexible bounds (at least they should have the same size)
+    TypeComponentPosition.FLEXIBLE_LOWER -> {
+        // T : String -> String
+        // in T : String -> String
+        // T : Enum<T> -> Enum<*>
+        erasedUpperBound
    }
+    TypeComponentPosition.FLEXIBLE_UPPER, TypeComponentPosition.INFLEXIBLE -> {
+        if (!parameter.variance.allowsOutPosition)
+        // in T -> Comparable<Nothing>
+            session.builtinTypes.nothingType.type
+        else if (erasedUpperBound is ConeClassLikeType &&
+            erasedUpperBound.lookupTag.toSymbol(session)!!.firUnsafe<FirRegularClass>().typeParameters.isNotEmpty()
+        )
+        // T : Enum<E> -> out Enum<*>
+            ConeKotlinTypeProjectionOut(erasedUpperBound)
+        else
+        // T : String -> *
+            ConeStarProjection
+    }
+}
+
+// Definition:
+// ErasedUpperBound(T : G<t>) = G<*> // UpperBound(T) is a type G<t> with arguments
+// 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(
+    // 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
+    potentiallyRecursiveTypeParameter: FirTypeParameter? = null,
+    defaultValue: (() -> ConeKotlinType) = { ConeKotlinErrorType("Can't compute erased upper bound of type parameter `$this`") }
+): ConeKotlinType {
+    if (this === potentiallyRecursiveTypeParameter) return defaultValue()
+
+    val firstUpperBound = this.bounds.first().coneTypeUnsafe<ConeKotlinType>()
+
+    return getErasedVersionOfFirstUpperBound(firstUpperBound, mutableSetOf(this, potentiallyRecursiveTypeParameter), defaultValue)
+}
+
+private fun getErasedVersionOfFirstUpperBound(
+    firstUpperBound: ConeKotlinType,
+    alreadyVisitedParameters: MutableSet<FirTypeParameter?>,
+    defaultValue: () -> ConeKotlinType
+): ConeKotlinType =
+    when (firstUpperBound) {
+        is ConeClassLikeType ->
+            firstUpperBound.withArguments(firstUpperBound.typeArguments.map { ConeStarProjection }.toTypedArray())
+
+        is ConeFlexibleType -> {
+            val lowerBound =
+                getErasedVersionOfFirstUpperBound(firstUpperBound.lowerBound, alreadyVisitedParameters, defaultValue)
+                    .lowerBoundIfFlexible()
+            if (firstUpperBound.upperBound is ConeTypeParameterType) {
+                // Avoid exponential complexity
+                ConeFlexibleType(
+                    lowerBound,
+                    lowerBound.withNullability(ConeNullability.NULLABLE)
+                )
+            } else {
+                ConeFlexibleType(
+                    lowerBound,
+                    getErasedVersionOfFirstUpperBound(firstUpperBound.upperBound, alreadyVisitedParameters, defaultValue)
+                )
+            }
+        }
+        is ConeTypeParameterType -> {
+            val current = firstUpperBound.lookupTag.typeParameterSymbol.fir
+
+            if (alreadyVisitedParameters.add(current)) {
+                val nextUpperBound = current.bounds.first().coneTypeUnsafe<ConeKotlinType>()
+                getErasedVersionOfFirstUpperBound(nextUpperBound, alreadyVisitedParameters, defaultValue)
+            } else {
+                defaultValue()
+            }
+        }
+        else -> error("Unexpected kind of firstUpperBound: $firstUpperBound [${firstUpperBound::class}]")
+    }
+
+private fun JavaClassifierType.toConeKotlinTypeForFlexibleBound(
+    session: FirSession,
+    javaTypeParameterStack: JavaTypeParameterStack,
+    isLowerBound: Boolean,
+    forTypeParameterBounds: Boolean
+): ConeKotlinType {
    return when (val classifier = classifier) {
        is JavaClass -> {
            //val classId = classifier.classId!!
            var classId = JavaToKotlinClassMap.mapJavaToKotlin(classifier.fqName!!) ?: classifier.classId!!
-            classId = classId.readOnlyToMutable() ?: classId
+
+            if (isLowerBound) {
+                classId = classId.readOnlyToMutable() ?: classId
+            }

            val lookupTag = ConeClassLikeLookupTagImpl(classId)

-            val mappedTypeArguments = when (isRaw) {
-                true -> classifier.typeParameters.map { ConeStarProjection }
-                false -> typeArguments.map { argument ->
-                    argument.toConeProjection(
-                        session, javaTypeParameterStack, boundTypeParameter = null, nullability = typeParametersNullability
+            val mappedTypeArguments = if (isRaw) {
+
+                val defaultArgs = (1..classifier.typeParameters.size).map { ConeStarProjection }
+
+                if (forTypeParameterBounds) {
+                    // This is not fully correct, but it's a simple fix for some time to avoid recursive definition:
+                    // to create a proper raw type arguments, we should take class parameters some time
+                    defaultArgs
+                } else {
+                    val classSymbol = session.firSymbolProvider.getClassLikeSymbolByFqName(classId) as? FirRegularClassSymbol
+                    val position = if (isLowerBound) TypeComponentPosition.FLEXIBLE_LOWER else TypeComponentPosition.FLEXIBLE_UPPER
+
+                    classSymbol?.fir?.createRawArguments(defaultArgs, position) ?: defaultArgs
+                }
+            } else {
+                typeArguments.map { argument ->
+                    argument.toConeProjectionWithoutEnhancement(
+                        session, javaTypeParameterStack, boundTypeParameter = null
                    )
                }
            }

            lookupTag.constructClassType(
-                mappedTypeArguments.toTypedArray(), nullability.isNullable
+                mappedTypeArguments.toTypedArray(), isNullable = !isLowerBound
            )
        }
        is JavaTypeParameter -> {
            val symbol = javaTypeParameterStack[classifier]
-            ConeTypeParameterTypeImpl(symbol.toLookupTag(), nullability.isNullable)
+            ConeTypeParameterTypeImpl(symbol.toLookupTag(), isNullable = !isLowerBound)
        }
        else -> ConeClassErrorType(reason = "Unexpected classifier: $classifier")
    }
 }

+private fun FirRegularClass.createRawArguments(
+    defaultArgs: List<ConeStarProjection>,
+    position: TypeComponentPosition
+) = typeParameters.map { typeParameter ->
+    val erasedUpperBound = typeParameter.getErasedUpperBound {
+        defaultType().withArguments(defaultArgs.toTypedArray())
+    }
+    computeRawProjection(session, typeParameter, position, erasedUpperBound)
+}
+
 internal fun JavaAnnotation.toFirAnnotationCall(
    session: FirSession, javaTypeParameterStack: JavaTypeParameterStack
 ): FirAnnotationCall {
@@ -226,11 +355,10 @@ internal fun JavaValueParameter.toFirValueParameter(
    }
 }

-internal fun JavaType?.toConeProjection(
+private fun JavaType?.toConeProjectionWithoutEnhancement(
    session: FirSession,
    javaTypeParameterStack: JavaTypeParameterStack,
-    boundTypeParameter: FirTypeParameter?,
-    nullability: ConeNullability = ConeNullability.NOT_NULL
+    boundTypeParameter: FirTypeParameter?
 ): ConeKotlinTypeProjection {
    return when (this) {
        null -> ConeStarProjection
@@ -241,7 +369,7 @@ internal fun JavaType?.toConeProjection(
            if (bound == null || parameterVariance != INVARIANT && parameterVariance != argumentVariance) {
                ConeStarProjection
            } else {
-                val boundType = bound.toConeKotlinTypeWithNullability(session, javaTypeParameterStack, nullability)
+                val boundType = bound.toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack)
                if (argumentVariance == OUT_VARIANCE) {
                    ConeKotlinTypeProjectionOut(boundType)
                } else {
@@ -249,7 +377,7 @@ internal fun JavaType?.toConeProjection(
                }
            }
        }
-        is JavaClassifierType -> toConeKotlinTypeWithNullability(session, javaTypeParameterStack, nullability)
+        is JavaClassifierType -> toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack)
        else -> ConeClassErrorType("Unexpected type argument: $this")
    }
 }
@@ -342,7 +470,12 @@ internal fun Any?.createConstant(session: FirSession): FirExpression {
 private fun JavaType.toFirResolvedTypeRef(
    session: FirSession, javaTypeParameterStack: JavaTypeParameterStack
 ): FirResolvedTypeRef {
-    if (this is JavaClassifierType) return toFirResolvedTypeRef(session, javaTypeParameterStack)
+    if (this is JavaClassifierType) return toFirResolvedTypeRef(
+        session,
+        javaTypeParameterStack,
+        isForSupertypes = false,
+        forTypeParameterBounds = false
+    )
    return buildResolvedTypeRef {
        type = ConeClassErrorType("Unexpected JavaType: $this")
    }
@@ -12,9 +12,8 @@ import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.classId
 import org.jetbrains.kotlin.fir.java.JavaTypeParameterStack
-import org.jetbrains.kotlin.fir.java.toConeKotlinTypeWithNullability
+import org.jetbrains.kotlin.fir.java.toConeKotlinTypeWithoutEnhancement
 import org.jetbrains.kotlin.fir.java.toFirJavaTypeRef
-import org.jetbrains.kotlin.fir.java.toNotNullConeKotlinType
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.jvm.FirJavaTypeRef
@@ -58,7 +57,8 @@ internal class EnhancementSignatureParts(
            }
        }

-        val containsFunctionN = current.toNotNullConeKotlinType(session, javaTypeParameterStack).contains {
+        val typeWithoutEnhancement = current.type.toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack)
+        val containsFunctionN = typeWithoutEnhancement.contains {
            if (it is ConeClassErrorType) false
            else {
                val classId = it.lookupTag.classId
@@ -67,7 +67,10 @@ internal class EnhancementSignatureParts(
            }
        }

-        val enhancedCurrent = current.enhance(session, javaTypeParameterStack, qualifiersWithPredefined ?: qualifiers)
+        val enhancedCurrent = current.enhance(
+            session, qualifiersWithPredefined ?: qualifiers,
+            typeWithoutEnhancement
+        )
        return PartEnhancementResult(
            enhancedCurrent, wereChanges = true, containsFunctionN = containsFunctionN
        )
@@ -150,10 +153,10 @@ internal class EnhancementSignatureParts(
                }
            }
            is FirJavaTypeRef -> {
+                val convertedType = type.toConeKotlinTypeWithoutEnhancement(session, javaTypeParameterStack)
                Pair(
-                    // TODO: optimize
-                    type.toConeKotlinTypeWithNullability(session, javaTypeParameterStack, nullability = ConeNullability.NOT_NULL),
-                    type.toConeKotlinTypeWithNullability(session, javaTypeParameterStack, nullability = ConeNullability.NULLABLE)
+                    convertedType.lowerBoundIfFlexible(),
+                    convertedType.upperBoundIfFlexible()
                )
            }
            else -> return JavaTypeQualifiers.NONE
@@ -16,15 +16,11 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildQualifiedAccessExpressi
 import org.jetbrains.kotlin.fir.java.JavaTypeParameterStack
 import org.jetbrains.kotlin.fir.java.declarations.FirJavaClass
 import org.jetbrains.kotlin.fir.java.declarations.FirJavaField
-import org.jetbrains.kotlin.fir.java.toConeProjection
-import org.jetbrains.kotlin.fir.java.toNotNullConeKotlinType
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
-import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
+import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
-import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.typeContext
@@ -36,13 +32,14 @@ import org.jetbrains.kotlin.load.java.JvmAnnotationNames.DEFAULT_VALUE_FQ_NAME
 import org.jetbrains.kotlin.load.java.descriptors.AnnotationDefaultValue
 import org.jetbrains.kotlin.load.java.descriptors.NullDefaultValue
 import org.jetbrains.kotlin.load.java.descriptors.StringDefaultValue
-import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.load.java.structure.JavaClassifierType
+import org.jetbrains.kotlin.load.java.structure.JavaType
 import org.jetbrains.kotlin.load.java.typeEnhancement.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.AbstractStrictEqualityTypeChecker
-import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.types.RawType
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.extractRadix

@@ -56,60 +53,23 @@ internal class IndexedJavaTypeQualifiers(private val data: Array<JavaTypeQualifi

 internal fun FirJavaTypeRef.enhance(
    session: FirSession,
-    javaTypeParameterStack: JavaTypeParameterStack,
-    qualifiers: IndexedJavaTypeQualifiers
+    qualifiers: IndexedJavaTypeQualifiers,
+    typeWithoutEnhancement: ConeKotlinType,
 ): FirResolvedTypeRef {
-    return type.enhancePossiblyFlexible(session, javaTypeParameterStack, annotations, qualifiers, 0)
+    return typeWithoutEnhancement.enhancePossiblyFlexible(session, annotations, qualifiers, 0)
 }

 // The index in the lambda is the position of the type component:
 // Example: for `A<B, C<D, E>>`, indices go as follows: `0 - A<...>, 1 - B, 2 - C<D, E>, 3 - D, 4 - E`,
 // which corresponds to the left-to-right breadth-first walk of the tree representation of the type.
 // For flexible types, both bounds are indexed in the same way: `(A<B>..C<D>)` gives `0 - (A<B>..C<D>), 1 - B and D`.
-private fun JavaType?.enhancePossiblyFlexible(
+private fun ConeKotlinType.enhancePossiblyFlexible(
    session: FirSession,
-    javaTypeParameterStack: JavaTypeParameterStack,
    annotations: List<FirAnnotationCall>,
    qualifiers: IndexedJavaTypeQualifiers,
    index: Int
 ): FirResolvedTypeRef {
-    val type = this
-    val arguments = this?.typeArguments().orEmpty()
-    val enhanced = when (type) {
-        is JavaClassifierType -> {
-            val lowerResult = type.enhanceInflexibleType(
-                session, javaTypeParameterStack, annotations, arguments, TypeComponentPosition.FLEXIBLE_LOWER, qualifiers, index
-            )
-            val upperResult = type.enhanceInflexibleType(
-                session, javaTypeParameterStack, annotations, arguments, TypeComponentPosition.FLEXIBLE_UPPER, qualifiers, index
-            )
-
-            when {
-                type.isRaw -> ConeRawType(lowerResult, upperResult)
-                else -> coneFlexibleOrSimpleType(
-                    session, lowerResult, upperResult, isNotNullTypeParameter = qualifiers(index).isNotNullTypeParameter
-                )
-            }
-        }
-        is JavaArrayType -> {
-            val baseEnhanced = type.toNotNullConeKotlinType(session, javaTypeParameterStack)
-
-            val upperBound = if (baseEnhanced.typeArguments.isNotEmpty()) {
-                val typeArgument = baseEnhanced.typeArguments.first() as ConeKotlinType
-                baseEnhanced.withArguments(arrayOf(ConeKotlinTypeProjectionOut(typeArgument)))
-            } else {
-                baseEnhanced
-            }
-            coneFlexibleOrSimpleType(
-                session, baseEnhanced,
-                upperBound.withNullability(ConeNullability.NULLABLE),
-                isNotNullTypeParameter = false
-            )
-        }
-        else -> {
-            type.toNotNullConeKotlinType(session, javaTypeParameterStack)
-        }
-    }
+    val enhanced = enhanceConeKotlinType(session, qualifiers, index)

    return buildResolvedTypeRef {
        this.type = enhanced
@@ -117,9 +77,34 @@ private fun JavaType?.enhancePossiblyFlexible(
    }
 }

-private fun JavaType?.subtreeSize(): Int {
-    if (this !is JavaClassifierType) return 1
-    return 1 + typeArguments.sumBy { it?.subtreeSize() ?: 0 }
+private fun ConeKotlinType.enhanceConeKotlinType(
+    session: FirSession,
+    qualifiers: IndexedJavaTypeQualifiers,
+    index: Int
+): ConeKotlinType {
+    return when (this) {
+        is ConeFlexibleType -> {
+            val lowerResult = lowerBound.enhanceInflexibleType(
+                session, TypeComponentPosition.FLEXIBLE_LOWER, qualifiers, index
+            )
+            val upperResult = upperBound.enhanceInflexibleType(
+                session, TypeComponentPosition.FLEXIBLE_UPPER, qualifiers, index
+            )
+
+            when (this) {
+                is ConeRawType -> ConeRawType(lowerResult, upperResult)
+                else -> coneFlexibleOrSimpleType(
+                    session, lowerResult, upperResult, isNotNullTypeParameter = qualifiers(index).isNotNullTypeParameter
+                )
+            }
+        }
+        is ConeSimpleKotlinType -> enhanceInflexibleType(session, TypeComponentPosition.INFLEXIBLE, qualifiers, index)
+        else -> this
+    }
+}
+
+private fun ConeKotlinType.subtreeSize(): Int {
+    return 1 + typeArguments.sumBy { ((it as? ConeKotlinType)?.subtreeSize() ?: 0) + 1 }
 }

 private fun coneFlexibleOrSimpleType(
@@ -164,155 +149,43 @@ private fun ClassId.mutableToReadOnly(): ClassId? {
    }
 }

-// Definition:
-// ErasedUpperBound(T : G<t>) = G<*> // UpperBound(T) is a type G<t> with arguments
-// 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(
-    // 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
-    potentiallyRecursiveTypeParameter: FirTypeParameter? = null,
-    defaultValue: (() -> ConeKotlinType) = { ConeKotlinErrorType("Can't compute erased upper bound of type parameter `$this`") }
-): ConeKotlinType {
-    if (this === potentiallyRecursiveTypeParameter) return defaultValue()
-
-    val firstUpperBound = this.bounds.first().coneTypeUnsafe<ConeKotlinType>()
-
-    return getErasedVersionOfFirstUpperBound(firstUpperBound, mutableSetOf(this, potentiallyRecursiveTypeParameter), defaultValue)
-}
-
-private fun getErasedVersionOfFirstUpperBound(
-    firstUpperBound: ConeKotlinType,
-    alreadyVisitedParameters: MutableSet<FirTypeParameter?>,
-    defaultValue: () -> ConeKotlinType
-): ConeKotlinType =
-    when (firstUpperBound) {
-        is ConeClassLikeType ->
-            firstUpperBound.withArguments(firstUpperBound.typeArguments.map { ConeStarProjection }.toTypedArray())
-
-        is ConeFlexibleType -> {
-            val lowerBound =
-                getErasedVersionOfFirstUpperBound(firstUpperBound.lowerBound, alreadyVisitedParameters, defaultValue)
-                    .lowerBoundIfFlexible()
-            if (firstUpperBound.upperBound is ConeTypeParameterType) {
-                // Avoid exponential complexity
-                ConeFlexibleType(
-                    lowerBound,
-                    lowerBound.withNullability(ConeNullability.NULLABLE)
-                )
-            } else {
-                ConeFlexibleType(
-                    lowerBound,
-                    getErasedVersionOfFirstUpperBound(firstUpperBound.upperBound, alreadyVisitedParameters, defaultValue)
-                )
-            }
-        }
-        is ConeTypeParameterType -> {
-            val current = firstUpperBound.lookupTag.typeParameterSymbol.fir
-
-            if (alreadyVisitedParameters.add(current)) {
-                val nextUpperBound = current.bounds.first().coneTypeUnsafe<ConeKotlinType>()
-                getErasedVersionOfFirstUpperBound(nextUpperBound, alreadyVisitedParameters, defaultValue)
-            } else {
-                defaultValue()
-            }
-        }
-        else -> error("Unexpected kind of firstUpperBound: $firstUpperBound [${firstUpperBound::class}]")
-    }
-
-
-fun computeProjection(
+private fun ConeKotlinType.enhanceInflexibleType(
    session: FirSession,
-    parameter: FirTypeParameter,
-    attr: TypeComponentPosition,
-    erasedUpperBound: ConeKotlinType = parameter.getErasedUpperBound()
-) = when (attr) {
-    // Raw(List<T>) => (List<Any?>..List<*>)
-    // Raw(Enum<T>) => (Enum<Enum<*>>..Enum<out Enum<*>>)
-    // In the last case upper bound is equal to star projection `Enum<*>`,
-    // but we want to keep matching tree structure of flexible bounds (at least they should have the same size)
-    TypeComponentPosition.FLEXIBLE_LOWER -> {
-        // T : String -> String
-        // in T : String -> String
-        // T : Enum<T> -> Enum<*>
-        erasedUpperBound
-    }
-    TypeComponentPosition.FLEXIBLE_UPPER, TypeComponentPosition.INFLEXIBLE -> {
-        if (!parameter.variance.allowsOutPosition)
-        // in T -> Comparable<Nothing>
-            session.builtinTypes.nothingType.type
-        else if (erasedUpperBound is ConeClassLikeType &&
-            erasedUpperBound.lookupTag.toSymbol(session)!!.firUnsafe<FirRegularClass>().typeParameters.isNotEmpty()
-        )
-        // T : Enum<E> -> out Enum<*>
-            ConeKotlinTypeProjectionOut(erasedUpperBound)
-        else
-        // T : String -> *
-            ConeStarProjection
-    }
-}
-
-private fun JavaClassifierType.enhanceInflexibleType(
-    session: FirSession,
-    javaTypeParameterStack: JavaTypeParameterStack,
-    annotations: List<FirAnnotationCall>,
-    arguments: List<JavaType?>,
    position: TypeComponentPosition,
    qualifiers: IndexedJavaTypeQualifiers,
    index: Int
 ): ConeKotlinType {
-    val originalTag = when (val classifier = classifier) {
-        is JavaClass -> {
-            val classId = classifier.classId!!
-            var mappedId = JavaToKotlinClassMap.mapJavaToKotlin(classId.asSingleFqName())
-            if (mappedId != null) {
-                if (position == TypeComponentPosition.FLEXIBLE_LOWER) {
-                    mappedId = mappedId.readOnlyToMutable() ?: mappedId
-                }
-            }
-            val kotlinClassId = mappedId ?: classId
-            ConeClassLikeLookupTagImpl(kotlinClassId)
-        }
-        is JavaTypeParameter -> javaTypeParameterStack[classifier].toLookupTag()
-        else -> return toNotNullConeKotlinType(session, javaTypeParameterStack)
+    require(this !is ConeFlexibleType) {
+        "$this should not be flexible"
    }
+    if (this !is ConeLookupTagBasedType) return this
+
+    val originalTag = lookupTag

    val effectiveQualifiers = qualifiers(index)
    val enhancedTag = originalTag.enhanceMutability(effectiveQualifiers, position)

-    val enhancedArguments = if (isRaw) {
-        val firClassifier = originalTag.toSymbol(session)!!.firUnsafe<FirRegularClass>()
-        firClassifier.typeParameters.map {
-            val fir = it
-            val erasedUpperBound = fir.getErasedUpperBound {
-                firClassifier.defaultType().withArguments(firClassifier.typeParameters.map { ConeStarProjection }.toTypedArray())
-            }
-            computeProjection(session, fir, position, erasedUpperBound)
-        }
+    val enhancedArguments = if (this is RawType) {
+        // TODO: Support enhancing for raw types
+        typeArguments
    } else {
        var globalArgIndex = index + 1
-        arguments.mapIndexed { localArgIndex, arg ->
-            if (arg is JavaWildcardType) {
+        typeArguments.mapIndexed { localArgIndex, arg ->
+            if (arg.kind != ProjectionKind.INVARIANT) {
                globalArgIndex++
-                arg.toConeProjection(
-                    session,
-                    javaTypeParameterStack,
-                    ((originalTag as? FirBasedSymbol<*>)?.fir as? FirCallableMemberDeclaration<*>)?.typeParameters?.getOrNull(localArgIndex)
-                )
+                arg
            } else {
-                val argEnhancedTypeRef =
-                    arg.enhancePossiblyFlexible(session, javaTypeParameterStack, annotations, qualifiers, globalArgIndex)
+                require(arg is ConeKotlinType) { "Should be invariant type: $arg" }
                globalArgIndex += arg.subtreeSize()

-                argEnhancedTypeRef.type.type.toTypeProjection(Variance.INVARIANT)
+                arg.enhanceConeKotlinType(session, qualifiers, globalArgIndex)
            }
-        }
+        }.toTypedArray()
    }

    val enhancedNullability = getEnhancedNullability(effectiveQualifiers, position)

-    val enhancedType = enhancedTag.constructType(enhancedArguments.toTypedArray(), enhancedNullability)
+    val enhancedType = enhancedTag.constructType(enhancedArguments, enhancedNullability)

    // TODO: why all of these is needed
 //    val enhancement = if (effectiveQualifiers.isNotNullTypeParameter) NotNullTypeParameter(enhancedType) else enhancedType
@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.modality
 import org.jetbrains.kotlin.fir.java.JavaTypeParameterStack
 import org.jetbrains.kotlin.fir.java.enhancement.readOnlyToMutable
-import org.jetbrains.kotlin.fir.java.toNotNullConeKotlinType
+import org.jetbrains.kotlin.fir.java.toConeKotlinTypeProbablyFlexible
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractOverrideChecker
 import org.jetbrains.kotlin.fir.typeContext
@@ -32,7 +32,7 @@ class JavaOverrideChecker internal constructor(
            return candidateType.lookupTag.classId.let { it.readOnlyToMutable() ?: it } == baseType.lookupTag.classId.let { it.readOnlyToMutable() ?: it }
        }
        if (candidateType is ConeClassLikeType && baseType is ConeTypeParameterType) {
-            val boundType = baseType.lookupTag.typeParameterSymbol.fir.bounds.singleOrNull()?.toNotNullConeKotlinType(
+            val boundType = baseType.lookupTag.typeParameterSymbol.fir.bounds.singleOrNull()?.toConeKotlinTypeProbablyFlexible(
                session, javaTypeParameterStack
            )
            if (boundType != null) {
@@ -49,8 +49,8 @@ class JavaOverrideChecker internal constructor(

    override fun isEqualTypes(candidateTypeRef: FirTypeRef, baseTypeRef: FirTypeRef, substitutor: ConeSubstitutor) =
        isEqualTypes(
-            candidateTypeRef.toNotNullConeKotlinType(session, javaTypeParameterStack),
-            baseTypeRef.toNotNullConeKotlinType(session, javaTypeParameterStack),
+            candidateTypeRef.toConeKotlinTypeProbablyFlexible(session, javaTypeParameterStack),
+            baseTypeRef.toConeKotlinTypeProbablyFlexible(session, javaTypeParameterStack),
            substitutor
        )