[Commonizer] Implement ClassOrTypeAliasTypeCommonizer with type substitution

^KT-48288

native/commonizer/src/org/jetbrains/kotlin/commonizer/cir/CirSupertypesResolver.kt

@@ -64,7 +64,26 @@ internal class SimpleCirSupertypesResolver(
     }
 }
 
-private fun CirProvided.ClassType.toCirClassTypeOrNull(): CirClassType? {
+internal fun CirProvidedClassifiers.toCirClassOrTypeAliasTypeOrNull(type: CirProvided.Type): CirClassOrTypeAliasType? {
+    return when (type) {
+        is CirProvided.ClassType -> type.toCirClassTypeOrNull()
+        is CirProvided.TypeParameterType -> null
+        is CirProvided.TypeAliasType -> TODO()
+    }
+}
+
+internal fun CirProvidedClassifiers.toCirTypeAliasTypeOrNull(type: CirProvided.TypeAliasType): CirTypeAliasType? {
+    val typeAlias = this.classifier(type.typeAliasId) as? CirProvided.TypeAlias ?: return null
+    return CirTypeAliasType.createInterned(
+        typeAliasId = type.typeAliasId,
+        isMarkedNullable = type.isMarkedNullable,
+        arguments = type.arguments.map { it.toCirTypeProjection() ?: return null },
+        underlyingType = toCirClassOrTypeAliasTypeOrNull(typeAlias.underlyingType) ?: return null
+    )
+}
+
+// TODO NOW move
+internal fun CirProvided.ClassType.toCirClassTypeOrNull(): CirClassType? {
     return CirClassType.createInterned(
         classId = this.classId,
         outerType = this.outerType?.let { it.toCirClassTypeOrNull() ?: return null },
@@ -74,7 +93,7 @@ private fun CirProvided.ClassType.toCirClassTypeOrNull(): CirClassType? {
     )
 }
 
-private fun CirProvided.TypeProjection.toCirTypeProjection(): CirTypeProjection? {
+internal fun CirProvided.TypeProjection.toCirTypeProjection(): CirTypeProjection? {
     return when (this) {
         is CirProvided.StarTypeProjection -> CirStarTypeProjection
         is CirProvided.RegularTypeProjection -> CirRegularTypeProjection(

native/commonizer/src/org/jetbrains/kotlin/commonizer/core/AbstractFunctionOrPropertyCommonizer.kt

@@ -48,7 +48,7 @@ class FunctionOrPropertyBaseCommonizer(
             kind = values.singleDistinctValueOrNull { it.kind } ?: return null,
             modality = ModalityCommonizer().commonize(values.map { it.modality }) ?: return null,
             visibility = VisibilityCommonizer.lowering().commonize(values) ?: return null,
-            extensionReceiver = extensionReceiverCommonizer(values.map { it.extensionReceiver })?.receiver ?: return null,
+            extensionReceiver = (extensionReceiverCommonizer(values.map { it.extensionReceiver }) ?: return null).receiver,
             returnType = returnTypeCommonizer(values) ?: return null,
             typeParameters = TypeParameterListCommonizer(typeCommonizer).commonize(values.map { it.typeParameters }) ?: return null
         )

native/commonizer/src/org/jetbrains/kotlin/commonizer/core/ClassOrTypeAliasTypeCommonizer.kt

@@ -5,60 +5,105 @@
 
 package org.jetbrains.kotlin.commonizer.core
 
-import org.jetbrains.kotlin.commonizer.cir.CirClassOrTypeAliasType
-import org.jetbrains.kotlin.commonizer.cir.CirClassType
-import org.jetbrains.kotlin.commonizer.cir.CirTypeAliasType
+import org.jetbrains.kotlin.commonizer.cir.*
+import org.jetbrains.kotlin.commonizer.mergedtree.CirKnownClassifiers
+import org.jetbrains.kotlin.commonizer.mergedtree.CirProvided
+import org.jetbrains.kotlin.commonizer.utils.isUnderKotlinNativeSyntheticPackages
+import org.jetbrains.kotlin.commonizer.utils.safeCastValues
+import org.jetbrains.kotlin.commonizer.utils.singleDistinctValueOrNull
+import org.jetbrains.kotlin.descriptors.Visibilities
 
 internal class ClassOrTypeAliasTypeCommonizer(
-    private val typeCommonizer: TypeCommonizer
+    private val typeCommonizer: TypeCommonizer,
+    private val classifiers: CirKnownClassifiers
 ) : NullableSingleInvocationCommonizer<CirClassOrTypeAliasType> {
 
+    private val isMarkedNullableCommonizer = TypeNullabilityCommonizer(typeCommonizer.options)
+
     override fun invoke(values: List<CirClassOrTypeAliasType>): CirClassOrTypeAliasType? {
-        TODO()
+        if (values.isEmpty()) return null
+        val expansions = values.map { it.expandedType() }
+        val isMarkedNullable = isMarkedNullableCommonizer.commonize(expansions.map { it.isMarkedNullable }) ?: return null
+        val arguments = TypeArgumentListCommonizer(typeCommonizer).commonize(values.map { it.arguments }) ?: return null
+        val classifierId = selectClassifierId(values) ?: return null
+
+        val outerTypes = values.safeCastValues<CirClassOrTypeAliasType, CirClassType>()?.map { it.outerType }
+        val outerType = when {
+            outerTypes == null -> null
+            outerTypes.all { it == null } -> null
+            outerTypes.any { it == null } -> return null
+            else -> invoke(outerTypes.map { checkNotNull(it) }) as? CirClassType ?: return null
+        }
+
+        if (classifierId.packageName.isUnderKotlinNativeSyntheticPackages) {
+            return CirClassType.createInterned(
+                classId = classifierId,
+                outerType = outerType,
+                arguments = arguments,
+                visibility = Visibilities.Public,
+                isMarkedNullable = isMarkedNullable
+            )
+        }
+
+        when (val dependencyClassifier = classifiers.commonDependencies.classifier(classifierId)) {
+            is CirProvided.Class -> return CirClassType.createInterned(
+                classId = classifierId,
+                outerType = outerType,
+                arguments = arguments,
+                visibility = Visibilities.Public,
+                isMarkedNullable = isMarkedNullable
+            )
+
+            is CirProvided.TypeAlias -> return CirTypeAliasType.createInterned(
+                typeAliasId = classifierId,
+                arguments = arguments,
+                isMarkedNullable = isMarkedNullable,
+                underlyingType = classifiers.commonDependencies
+                    .toCirClassOrTypeAliasTypeOrNull(dependencyClassifier.underlyingType) ?: return null
+            )
+
+            else -> Unit
+        }
+
+        val commonizedClassifier = classifiers.commonizedNodes.classNode(classifierId)?.commonDeclaration?.invoke()
+            ?: classifiers.commonizedNodes.typeAliasNode(classifierId)?.commonDeclaration?.invoke()
+
+
+        when (commonizedClassifier) {
+            is CirClass -> return CirClassType.createInterned(
+                classId = classifierId,
+                outerType = outerType,
+                arguments = arguments,
+                visibility = Visibilities.Public,
+                isMarkedNullable = isMarkedNullable
+            )
+
+            is CirTypeAlias -> return CirTypeAliasType.createInterned(
+                typeAliasId = classifierId,
+                arguments = arguments,
+                isMarkedNullable = isMarkedNullable,
+                underlyingType = computeSuitableUnderlyingType(
+                    classifiers, typeCommonizer, commonizedClassifier.underlyingType
+                )?.makeNullableIfNecessary(isMarkedNullable) ?: return null
+            )
+
+            else -> Unit
+        }
+
+        return null
     }
 
-    /*
-    override fun commonize(first: CirClassOrTypeAliasType, second: CirClassOrTypeAliasType): CirClassOrTypeAliasType? {
-        if (first is CirClassType && second is CirClassType) {
-            return ClassTypeCommonizer(classifiers, options).commonize(listOf(first, second))
-                ?: if (options.enableOptimisticNumberTypeCommonization) OptimisticNumbersTypeCommonizer.commonize(first, second) else null
-        }
+    private fun selectClassifierId(types: List<CirClassOrTypeAliasType>): CirEntityId? {
+        types.singleDistinctValueOrNull { it.classifierId }?.let { return it }
 
-        if (first is CirTypeAliasType && second is CirTypeAliasType) {
-            /*
-            In case regular type-alias-type commonization fails, we try to expand all type-aliases and
-            try our luck with commonizing those class types
-             */
-            return TypeAliasTypeCommonizer(classifiers, options).commonize(listOf(first, second))
-                ?: ClassOrTypeAliasTypeCommonizer(
-                    classifiers, options.withOptimisticNumberTypeCommonizationEnabled()
-                ).commonize(first.expandedType(), second.expandedType())
-        }
+        if (typeCommonizer.options.enableTypeAliasSubstitution) {
+            val commonId = types.singleDistinctValueOrNull {
+                classifiers.commonClassifierIdResolver.findCommonId(it.classifierId)
+            } ?: return null
 
-        val classType = when {
-            first is CirClassType -> first
-            second is CirClassType -> second
-            else -> return null
-        }
-
-        val typeAliasType = when {
-            first is CirTypeAliasType -> first
-            second is CirTypeAliasType -> second
-            else -> return null
-        }
-
-        /*
-        TypeAliasCommonizer will be able to figure out if the typealias will be represented as expect class in common.
-        If so, re-use this class type, otherwise: try to expand the typeAlias
-         */
-        val typeAliasClassType = TypeAliasTypeCommonizer(classifiers, options).commonize(listOf(typeAliasType))?.expandedType()
-            ?: typeAliasType.expandedType()
-
-        return commonize(classType, typeAliasClassType)
+            return commonId.aliases.maxByOrNull { candidate -> types.count { it.classifierId == candidate } }!!
+        } else return null
     }
-
-     */
-
 }
 
 internal tailrec fun CirClassOrTypeAliasType.expandedType(): CirClassType = when (this) {

native/commonizer/src/org/jetbrains/kotlin/commonizer/core/TypeAliasCommonizer.kt

@@ -13,9 +13,14 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 class TypeAliasCommonizer(
-    private val typeCommonizer: TypeCommonizer,
+    typeCommonizer: TypeCommonizer,
     private val classifiers: CirKnownClassifiers,
 ) : NullableSingleInvocationCommonizer<CirTypeAlias> {
+
+    private val typeCommonizer = typeCommonizer.withOptions {
+        withTypeAliasSubstitutionEnabled(false).withOptimisticNumberTypeCommonizationEnabled(true)
+    }
+
     override fun invoke(values: List<CirTypeAlias>): CirTypeAlias? {
         if (values.isEmpty()) return null
 
@@ -23,8 +28,7 @@ class TypeAliasCommonizer(
 
         val typeParameters = TypeParameterListCommonizer(typeCommonizer).commonize(values.map { it.typeParameters }) ?: return null
 
-        val underlyingType = typeCommonizer.withOptions { withOptimisticNumberTypeCommonizationEnabled() }
-            .invoke(values.map { it.underlyingType }) as? CirClassOrTypeAliasType ?: return null
+        val underlyingType = typeCommonizer(values.map { it.underlyingType }) as? CirClassOrTypeAliasType ?: return null
 
         val visibility = VisibilityCommonizer.lowering().commonize(values) ?: return null
 

native/commonizer/src/org/jetbrains/kotlin/commonizer/core/TypeCommonizer.kt

@@ -14,7 +14,7 @@ class TypeCommonizer(
     val options: Options = Options.default
 ) : NullableSingleInvocationCommonizer<CirType> {
 
-    private val classOrTypeAliasTypeCommonizer = ClassOrTypeAliasTypeCommonizer(this)
+    private val classOrTypeAliasTypeCommonizer = ClassOrTypeAliasTypeCommonizer(this, classifiers)
     private val flexibleTypeCommonizer = FlexibleTypeAssociativeCommonizer(this)
 
     override fun invoke(values: List<CirType>): CirType? {
@@ -35,7 +35,8 @@ class TypeCommonizer(
 
     data class Options(
         val enableOptimisticNumberTypeCommonization: Boolean = false,
-        val enableCovariantNullabilityCommonization: Boolean = false
+        val enableCovariantNullabilityCommonization: Boolean = false,
+        val enableTypeAliasSubstitution: Boolean = true
     ) {
 
         fun withOptimisticNumberTypeCommonizationEnabled(enabled: Boolean = true): Options {
@@ -48,6 +49,11 @@ class TypeCommonizer(
             else copy(enableCovariantNullabilityCommonization = enabled)
         }
 
+        fun withTypeAliasSubstitutionEnabled(enabled: Boolean = true): Options {
+            return if (enableTypeAliasSubstitution == enabled) this
+            else copy(enableTypeAliasSubstitution = enabled)
+        }
+
         companion object {
             val default = Options()
         }

native/commonizer/src/org/jetbrains/kotlin/commonizer/facade.kt

@@ -56,6 +56,7 @@ internal fun commonizeTarget(
             targetDependencies = availableTrees.mapValue(CirTreeRoot::dependencies),
             commonizedNodes = CirCommonizedClassifierNodes.default(),
             commonDependencies = parameters.dependencyClassifiers(output)
+
         )
 
         val mergedTree = mergeCirTree(parameters.storageManager, classifiers, availableTrees)

native/commonizer/src/org/jetbrains/kotlin/commonizer/mergedtree/CirClassifierIndex.kt

@@ -28,29 +28,6 @@ interface CirClassifierIndex {
     fun findTypeAliasesWithUnderlyingType(underlyingClassifier: CirEntityId): List<CirTreeTypeAlias>
 }
 
-/**
- * Resolves all typealias that will either declare the given [underlyingClassifier] as their underlying type directly, *or transitively*.
- * The returned List will be ordered BFS:
- * Elements at the beginning of the returned List will have a 'shorter path' to the underlying classifier
- */
-internal fun CirClassifierIndex.findAllTypeAliasesWithUnderlyingType(underlyingClassifier: CirEntityId): List<CirTreeTypeAlias> {
-    /* Fast Path (no aliases) */
-    val firstAliases = findTypeAliasesWithUnderlyingType(underlyingClassifier)
-    if (firstAliases.isEmpty()) return emptyList()
-
-    val resolved = ArrayList<CirTreeTypeAlias>(firstAliases)
-    val enqueued = ArrayDeque<CirEntityId>(firstAliases.size * 2)
-    firstAliases.forEach { alias -> enqueued.add(alias.id) }
-
-    while (enqueued.isNotEmpty()) {
-        val next = enqueued.removeFirst()
-        val aliases = findTypeAliasesWithUnderlyingType(next)
-        resolved.addAll(aliases)
-        aliases.forEach { alias -> enqueued.add(alias.id) }
-    }
-
-    return resolved
-}
 
 internal fun CirClassifierIndex.findClass(id: CirEntityId): CirClass? = findClassifier(id) as? CirClass
 

native/commonizer/src/org/jetbrains/kotlin/commonizer/mergedtree/CirCommonClassifierId.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.commonizer.mergedtree
 import org.jetbrains.kotlin.commonizer.cir.CirEntityId
 
 @JvmInline
-internal value class CirCommonClassifierId(val aliases: LinkedHashSet<CirEntityId>) {
+value class CirCommonClassifierId(val aliases: LinkedHashSet<CirEntityId>) {
     override fun toString(): String {
         return aliases.joinToString(prefix = "(", postfix = ")")
     }

native/commonizer/src/org/jetbrains/kotlin/commonizer/mergedtree/CirCommonClassifierIdResolver.kt

@@ -11,16 +11,20 @@ import org.jetbrains.kotlin.commonizer.TargetDependent
 import org.jetbrains.kotlin.commonizer.cir.CirEntityId
 import org.jetbrains.kotlin.commonizer.cir.CirTypeAlias
 
-internal fun CirCommonClassifierIdResolver(classifierIndices: TargetDependent<CirClassifierIndex>): CirCommonClassifierIdResolver {
-    return CirCommonClassifierIdResolverImpl(classifierIndices)
+internal fun CirCommonClassifierIdResolver(
+    classifierIndices: TargetDependent<CirClassifierIndex>,
+    dependencies: CirProvidedClassifiers = CirProvidedClassifiers.EMPTY
+): CirCommonClassifierIdResolver {
+    return CirCommonClassifierIdResolverImpl(classifierIndices, dependencies)
 }
 
-internal interface CirCommonClassifierIdResolver {
+interface CirCommonClassifierIdResolver {
     fun findCommonId(id: CirEntityId): CirCommonClassifierId?
 }
 
 private class CirCommonClassifierIdResolverImpl(
-    private val classifierIndices: TargetDependent<CirClassifierIndex>
+    private val classifierIndices: TargetDependent<CirClassifierIndex>,
+    private val dependencies: CirProvidedClassifiers
 ) : CirCommonClassifierIdResolver {
 
     private val cachedResults = HashMap<CirEntityId, CirCommonClassifierId>()
@@ -46,7 +50,9 @@ private class CirCommonClassifierIdResolverImpl(
 
         while (queue.isNotEmpty()) {
             val nextClassifierId = queue.removeFirst()
-            val foundClassifiers = classifierIndices.associateWith { index -> index.findClassifier(nextClassifierId) }
+            val foundClassifiers = classifierIndices.associateWith { index ->
+                index.findClassifier(nextClassifierId) ?: dependencies.classifier(nextClassifierId)
+            }
 
             /* Classifier is available for all targets */
             if (foundClassifiers.all { (_, classifier) -> classifier != null }) {
@@ -69,6 +75,17 @@ private class CirCommonClassifierIdResolverImpl(
                         queue.add(classifier.underlyingType.classifierId)
                     }
                 }
+
+                if (classifier is CirProvided.TypeAlias) {
+                    val underlyingTypeId = when (val underlyingType = classifier.underlyingType) {
+                        is CirProvided.ClassType -> underlyingType.classId
+                        is CirProvided.TypeAliasType -> underlyingType.typeAliasId
+                        is CirProvided.TypeParameterType -> null
+                    }
+                    if (underlyingTypeId != null && visited.add(underlyingTypeId)) {
+                        queue.add(underlyingTypeId)
+                    }
+                }
             }
         }
 

native/commonizer/src/org/jetbrains/kotlin/commonizer/mergedtree/CirKnownClassifiers.kt

@@ -13,7 +13,8 @@ class CirKnownClassifiers(
     val classifierIndices: TargetDependent<CirClassifierIndex>,
     val targetDependencies: TargetDependent<CirProvidedClassifiers>,
     val commonizedNodes: CirCommonizedClassifierNodes,
-    val commonDependencies: CirProvidedClassifiers
+    val commonDependencies: CirProvidedClassifiers,
+    val commonClassifierIdResolver: CirCommonClassifierIdResolver = CirCommonClassifierIdResolver(classifierIndices),
 )
 
 /** A set of all CIR nodes built for commonized classes and type aliases. */

native/commonizer/src/org/jetbrains/kotlin/commonizer/transformer/TypeSubstitutionCirNodeTransformer.kt

@@ -14,9 +14,9 @@ internal class TypeSubstitutionCirNodeTransformer(
     private val storageManager: StorageManager,
     private val classifiers: CirKnownClassifiers,
     private val typeSubstitutor: CirTypeSubstitutor,
-    private val commonClassifierIdResolver: CirCommonClassifierIdResolver = CirCommonClassifierIdResolver(classifiers.classifierIndices)
 ) : CirNodeTransformer {
 
+    private val commonClassifierIdResolver = classifiers.commonClassifierIdResolver
 
     override fun invoke(root: CirRootNode) {
         for (index in 0 until root.targetDeclarations.size) {

native/commonizer/src/org/jetbrains/kotlin/commonizer/tree/mergeCirTree.kt

@@ -17,7 +17,6 @@ internal data class TargetBuildingContext(
     val storageManager: StorageManager,
     val classifiers: CirKnownClassifiers,
     val memberContext: CirMemberContext = CirMemberContext.empty,
-    val commonClassifierIdResolver: CirCommonClassifierIdResolver = CirCommonClassifierIdResolver(classifiers.classifierIndices),
     val targets: Int, val targetIndex: Int
 ) {
     fun withMemberContextOf(clazz: CirClass) = copy(memberContext = memberContext.withContextOf(clazz))
@@ -83,7 +82,7 @@ internal fun CirNodeWithMembers<*, *>.buildFunction(
     context: TargetBuildingContext, function: CirFunction, parent: CirNode<*, *>? = null
 ) {
     val functionNode = functions.getOrPut(
-        FunctionApproximationKey.create(context.memberContext, context.commonClassifierIdResolver, function)
+        FunctionApproximationKey.create(context.memberContext, context.classifiers.commonClassifierIdResolver, function)
     ) {
         buildFunctionNode(context.storageManager, context.targets, context.classifiers, ParentNode(parent))
     }
@@ -94,7 +93,7 @@ internal fun CirNodeWithMembers<*, *>.buildProperty(
     context: TargetBuildingContext, property: CirProperty, parent: CirNode<*, *>? = null
 ) {
     val propertyNode = properties.getOrPut(
-        PropertyApproximationKey.create(context.memberContext, context.commonClassifierIdResolver, property)
+        PropertyApproximationKey.create(context.memberContext, context.classifiers.commonClassifierIdResolver, property)
     ) {
         buildPropertyNode(context.storageManager, context.targets, context.classifiers, ParentNode(parent))
     }
@@ -105,7 +104,7 @@ internal fun CirClassNode.buildConstructor(
     context: TargetBuildingContext, constructor: CirClassConstructor, parent: CirNode<*, *>
 ) {
     val constructorNode = constructors.getOrPut(
-        ConstructorApproximationKey.create(context.memberContext, context.commonClassifierIdResolver, constructor)
+        ConstructorApproximationKey.create(context.memberContext, context.classifiers.commonClassifierIdResolver, constructor)
     ) {
         buildClassConstructorNode(context.storageManager, context.targets, context.classifiers, ParentNode(parent))
     }