[Commonizer] Lookup descriptors that doesn't participate in commonization

^KT-40119

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/TargetProvider.kt

@@ -31,5 +31,5 @@ interface ModulesProvider {
     class ModuleInfo(val name: String, val originalLocation: File)
 
     fun loadModuleInfos(): Map<String, ModuleInfo>
-    fun loadModules(): Collection<ModuleDescriptor>
+    fun loadModules(): Map<String, ModuleDescriptor>
 }

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/builder/builderUtils.kt

@@ -9,9 +9,7 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.descriptors.commonizer.cir.*
 import org.jetbrains.kotlin.descriptors.commonizer.utils.isUnderStandardKotlinPackages
-import org.jetbrains.kotlin.descriptors.commonizer.utils.resolveClassOrTypeAlias
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.resolve.DescriptorFactory
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.types.*
@@ -104,11 +102,11 @@ internal fun CirSimpleType.buildType(
 ): SimpleType {
     val classifier: ClassifierDescriptor = when (val classifierId = classifierId) {
         is CirClassifierId.Class -> {
-            findClassOrTypeAlias(targetComponents, classifierId.classId).checkClassifierType<ClassDescriptor>()
+            targetComponents.findClassOrTypeAlias(classifierId.classId).checkClassifierType<ClassDescriptor>()
         }
         is CirClassifierId.TypeAlias -> {
             val classId = classifierId.classId
-            val classOrTypeAlias: ClassifierDescriptorWithTypeParameters = findClassOrTypeAlias(targetComponents, classId)
+            val classOrTypeAlias: ClassifierDescriptorWithTypeParameters = targetComponents.findClassOrTypeAlias(classId)
 
             if (classId.packageFqName.isUnderStandardKotlinPackages || !targetComponents.isCommon) {
                 // classifier type could be only type alias
@@ -138,26 +136,6 @@ internal fun CirSimpleType.buildType(
         simpleType
 }
 
-internal fun findClassOrTypeAlias(
-    targetComponents: TargetDeclarationsBuilderComponents,
-    classId: ClassId
-): ClassifierDescriptorWithTypeParameters = when {
-    classId.packageFqName.isUnderStandardKotlinPackages -> {
-        // look up for classifier in built-ins module:
-        val builtInsModule = targetComponents.builtIns.builtInsModule
-
-        // TODO: this works fine for Native as far as built-ins module contains full Native stdlib, but this is not enough for JVM and JS
-        builtInsModule.resolveClassOrTypeAlias(classId)
-            ?: error("Classifier ${classId.asString()} not found in built-ins module $builtInsModule for ${targetComponents.target}")
-    }
-
-    else -> {
-        // otherwise, find the appropriate user classifier:
-        targetComponents.findAppropriateClassOrTypeAlias(classId)
-            ?: error("Classifier ${classId.asString()} not found in created descriptors cache for ${targetComponents.target}")
-    }
-}
-
 private inline fun <reified T : ClassifierDescriptorWithTypeParameters> ClassifierDescriptorWithTypeParameters.checkClassifierType(): T {
     check(this is T) { "Mismatched classifier kinds.\nFound: ${this::class.java}, $this\nShould be: ${T::class.java}" }
     return this

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/builder/context.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.descriptors.commonizer.builder
 
+import gnu.trove.THashMap
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.commonizer.Parameters
@@ -21,6 +22,7 @@ import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.storage.NotNullLazyValue
 import org.jetbrains.kotlin.storage.StorageManager
 
 /**
@@ -131,26 +133,56 @@ class TargetDeclarationsBuilderComponents(
     val storageManager: StorageManager,
     val target: Target,
     val builtIns: KotlinBuiltIns,
+    val lazyModulesLookupTable: NotNullLazyValue<MutableMap<String, ModuleDescriptor?>>,
     val isCommon: Boolean,
     val index: Int,
     private val cache: DeclarationsBuilderCache
 ) {
     // N.B. this function may create new classifiers for types from Kotlin/Native forward declarations packages
-    fun findAppropriateClassOrTypeAlias(classId: ClassId): ClassifierDescriptorWithTypeParameters? {
+    fun findClassOrTypeAlias(classId: ClassId): ClassifierDescriptorWithTypeParameters {
+        return when {
+            classId.packageFqName.isUnderStandardKotlinPackages -> {
+                // look up for classifier in built-ins module:
+                val builtInsModule = builtIns.builtInsModule
 
-        return if (classId.packageFqName.isUnderKotlinNativeSyntheticPackages) {
+                // TODO: this works fine for Native as far as built-ins module contains full Native stdlib, but this is not enough for JVM and JS
-            // that's a synthetic Kotlin/Native classifier that was exported as forward declaration in one or more modules,
+                builtInsModule.resolveClassOrTypeAlias(classId)
-            // but did not match any existing class or typealias
+                    ?: error("Classifier ${classId.asString()} not found in built-ins module $builtInsModule for $target")
-            cache.getOrPutForwardDeclarationsModule(index) {
+            }
-                // N.B. forward declarations module is created only on demand
+            classId.packageFqName.isUnderKotlinNativeSyntheticPackages -> {
-                createKotlinNativeForwardDeclarationsModule(
+                // that's a synthetic Kotlin/Native classifier that was exported as forward declaration in one or more modules,
-                    storageManager = storageManager,
+                // but did not match any existing class or typealias
-                    builtIns = builtIns
+                cache.getOrPutForwardDeclarationsModule(index) {
-                )
+                    // N.B. forward declarations module is created only on demand
-            }.resolveClassOrTypeAlias(classId)
+                    createKotlinNativeForwardDeclarationsModule(
-        } else {
+                        storageManager = storageManager,
-            // look up in created descriptors cache
+                        builtIns = builtIns
-            cache.getCachedClassifier(classId, index)
+                    )
+                }.resolveClassOrTypeAlias(classId)
+                    ?: error("Classifier ${classId.asString()} not found for $target")
+            }
+            else -> {
+                // look up in created descriptors cache
+                val fromCache = cache.getCachedClassifier(classId, index)
+                fromCache
+                    ?: run {
+                        // attempt to load the original classifier
+                        if (!classId.packageFqName.isRoot) {
+                            // first, guess containing module and look up in it
+                            val classifier = lazyModulesLookupTable()
+                                .guessModuleByPackageFqName(classId.packageFqName)
+                                ?.resolveClassOrTypeAlias(classId)
+
+                            // if failed, then look up though all modules
+                            classifier
+                                ?: lazyModulesLookupTable().values
+                                    .asSequence()
+                                    .mapNotNull { it?.resolveClassOrTypeAlias(classId) }
+                                    .firstOrNull()
+                        } else null
+                    }
+                    ?: error("Classifier ${classId.asString()} not found for $target")
+            }
         }
     }
 }
@@ -163,6 +195,8 @@ fun CirRootNode.createGlobalBuilderComponents(
 
     val targetContexts = (0 until dimension).map { index ->
         val isCommon = index == indexOfCommon
+
+        // do not leak root inside of createLazyValue {} closures!!
         val root = if (isCommon) commonDeclaration()!! else targetDeclarations[index]!!
 
         val builtIns = root.builtInsProvider.loadBuiltIns()
@@ -170,10 +204,16 @@ fun CirRootNode.createGlobalBuilderComponents(
             "Unexpected built-ins class: ${builtIns::class.java}, $builtIns\nExpected: ${root.builtInsClass}"
         }
 
+        val lazyModulesLookupTable = storageManager.createLazyValue {
+            val source = if (isCommon) emptyMap() else parameters.targetProviders[index].modulesProvider.loadModules()
+            THashMap(source)
+        }
+
         TargetDeclarationsBuilderComponents(
             storageManager = storageManager,
             target = root.target,
             builtIns = builtIns,
+            lazyModulesLookupTable = lazyModulesLookupTable,
             isCommon = isCommon,
             index = index,
             cache = cache

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/konan/NativeDistributionModulesProvider.kt

@@ -43,7 +43,7 @@ internal class NativeDistributionModulesProvider(
         }
     }
 
-    override fun loadModules(): Collection<ModuleDescriptor> {
+    override fun loadModules(): Map<String, ModuleDescriptor> {
         val builtIns = loadBuiltIns()
         val stdlib = builtIns.builtInsModule
 
@@ -74,6 +74,6 @@ internal class NativeDistributionModulesProvider(
             module.setDependencies(listOf(module) + dependencies + forwardDeclarations)
         }
 
-        return platformModulesMap.values.toList()
+        return platformModulesMap
     }
 }

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/mergedtree/CirTreeMerger.kt

@@ -71,11 +71,11 @@ class CirTreeMerger(
             targetProvider.builtInsProvider
         )
 
-        val moduleDescriptors: Collection<ModuleDescriptor> = targetProvider.modulesProvider.loadModules()
+        val moduleDescriptors: Map<String, ModuleDescriptor> = targetProvider.modulesProvider.loadModules()
         val modules: MutableMap<Name, CirModuleNode> = rootNode.modules
 
-        moduleDescriptors.forEach { moduleDescriptor ->
+        moduleDescriptors.forEach { (name, moduleDescriptor) ->
-            if (moduleDescriptor.name.asString().removeSurrounding("<", ">") in commonModuleNames)
+            if (name in commonModuleNames)
                 processModule(modules, targetIndex, moduleDescriptor)
         }
     }

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/utils/moduleDescriptor.kt

@@ -62,4 +62,40 @@ internal fun ModuleDescriptor.resolveClassOrTypeAlias(classId: ClassId): Classif
     ) as? ClassifierDescriptorWithTypeParameters
 }
 
+internal fun MutableMap<String, ModuleDescriptor?>.guessModuleByPackageFqName(packageFqName: FqName): ModuleDescriptor? {
+    if (isEmpty()) return null
+
+    val packageFqNameRaw = packageFqName.asString()
+    if (containsKey(packageFqNameRaw)) {
+        return this[packageFqNameRaw] // might return null if this is a previously cached result
+    }
+
+    fun guessByEnding(): ModuleDescriptor? {
+        return entries
+            .firstOrNull { (name, _) -> name.endsWith(packageFqNameRaw, ignoreCase = true) }
+            ?.value
+    }
+
+    fun guessBySmartEnding(): ModuleDescriptor? {
+        val packageFqNameFragments = packageFqNameRaw.split('.')
+        if (packageFqNameFragments.size < 2) return null
+
+        return entries.firstOrNull { (name, _) ->
+            var startIndex = 0
+            for (fragment in packageFqNameFragments) {
+                val index = name.indexOf(fragment, startIndex = startIndex, ignoreCase = true)
+                if (index < startIndex)
+                    return@firstOrNull false
+                else
+                    startIndex = index + fragment.length
+            }
+            true
+        }?.value
+    }
+
+    val candidate = guessByEnding() ?: guessBySmartEnding()
+    this[packageFqNameRaw] = candidate // cache to speed-up the further look-ups
+    return candidate
+}
+
 internal val NativeFactories = KlibMetadataFactories(::KonanBuiltIns, NullFlexibleTypeDeserializer, NativeTypeTransformer())

native/commonizer/tests/org/jetbrains/kotlin/descriptors/commonizer/utils/mocks.kt

@@ -44,6 +44,7 @@ internal fun mockClassType(
         storageManager = LockBasedStorageManager.NO_LOCKS,
         target = InputTarget("Arbitrary target"),
         builtIns = DefaultBuiltIns.Instance,
+        lazyModulesLookupTable = LockBasedStorageManager.NO_LOCKS.createLazyValue { mutableMapOf() },
         isCommon = false,
         index = 0,
         cache = DeclarationsBuilderCache(1)
@@ -145,7 +146,7 @@ internal class MockModulesProvider : ModulesProvider {
     }
 
     override fun loadModuleInfos() = moduleInfos
-    override fun loadModules() = modules.values
+    override fun loadModules() = modules
 
     private fun fakeModuleInfo(name: String) = ModuleInfo(name, File("/tmp/commonizer/mocks/$name"))
 }