[KLIB] Reduce amount of descriptors loaded during desc-idSig resolution

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/DescriptorByIdSignatureFinder.kt

@@ -7,9 +7,13 @@ package org.jetbrains.kotlin.backend.common.serialization
 
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.ir.util.KotlinMangler
-import org.jetbrains.kotlin.resolve.scopes.getDescriptorsFiltered
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.scopes.MemberScope
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 class DescriptorByIdSignatureFinder(
     private val moduleDescriptor: ModuleDescriptor,
@@ -35,11 +39,12 @@ class DescriptorByIdSignatureFinder(
         MODULE_ONLY
     }
 
-    fun findDescriptorBySignature(signature: IdSignature): DeclarationDescriptor? = when (signature) {
+    fun findDescriptorBySignature(signature: IdSignature): DeclarationDescriptor? =
-        is IdSignature.AccessorSignature -> findDescriptorForAccessorSignature(signature)
+        when (signature) {
-        is IdSignature.PublicSignature -> findDescriptorForPublicSignature(signature)
+            is IdSignature.AccessorSignature -> findDescriptorForAccessorSignature(signature)
-        else -> error("only PublicSignature or AccessorSignature should reach this point, got $signature")
+            is IdSignature.PublicSignature -> findDescriptorForPublicSignature(signature)
-    }
+            else -> error("only PublicSignature or AccessorSignature should reach this point, got $signature")
+        }
 
     private fun findDescriptorForAccessorSignature(signature: IdSignature.AccessorSignature): DeclarationDescriptor? {
         val propertyDescriptor = findDescriptorBySignature(signature.propertySignature) as? PropertyDescriptor
@@ -48,37 +53,77 @@ class DescriptorByIdSignatureFinder(
         return propertyDescriptor.accessors.singleOrNull { it.name.asString() == shortName }
     }
 
-    private fun performLookup(signature: IdSignature.PublicSignature): Collection<DeclarationDescriptor> {
+    private fun isConstructorName(n: Name) = n.isSpecial && n.asString() == "<init>"
-        val declarationName = signature.firstNameSegment
+
+    private fun MemberScope.loadDescriptors(name: String, isLeaf: Boolean): Collection<DeclarationDescriptor> {
+        val descriptorName = Name.guessByFirstCharacter(name)
+        val classifier = getContributedClassifier(descriptorName, NoLookupLocation.FROM_BACKEND)
+        if (!isLeaf) {
+            return listOfNotNull(classifier)
+        }
+
+        val result = mutableListOf<DeclarationDescriptor>()
+        classifier?.let { result.add(it) }
+
+        result.addAll(getContributedFunctions(descriptorName, NoLookupLocation.FROM_BACKEND))
+        result.addAll(getContributedVariables(descriptorName, NoLookupLocation.FROM_BACKEND))
+
+        return result
+    }
+
+    private fun performLookup(nameSegments: List<String>, packageFqName: FqName): Collection<DeclarationDescriptor> {
+        val declarationName = nameSegments[0]
+        val isLeaf = nameSegments.size == 1
         return when (lookupMode) {
             LookupMode.MODULE_WITH_DEPENDENCIES -> {
                 moduleDescriptor
-                    .getPackage(signature.packageFqName())
+                    .getPackage(packageFqName)
-                    .memberScope
+                    .memberScope.loadDescriptors(declarationName, isLeaf)
-                    .getContributedDescriptors { name -> name.asString() == declarationName }
             }
             LookupMode.MODULE_ONLY -> {
                 (moduleDescriptor as ModuleDescriptorImpl)
                     .packageFragmentProviderForModuleContentWithoutDependencies
-                    .getPackageFragments(signature.packageFqName())
+                    .getPackageFragments(packageFqName)
-                    .map { it.getMemberScope() }
+                    .flatMap { it.getMemberScope().loadDescriptors(declarationName, isLeaf) }
-                    .flatMap { it.getContributedDescriptors { name -> name.asString() == declarationName } }
             }
         }
     }
 
     private fun findDescriptorForPublicSignature(signature: IdSignature.PublicSignature): DeclarationDescriptor? {
-        val toplevelDescriptors = performLookup(signature)
+        val nameSegments = signature.nameSegments
+        val toplevelDescriptors = performLookup(nameSegments, signature.packageFqName())
             .ifEmpty { return null }
-        val candidates = signature.nameSegments.drop(1).fold(toplevelDescriptors) { acc, current ->
+
-            acc.flatMap { container ->
+        var acc = toplevelDescriptors
+        val lastIndex = nameSegments.lastIndex
+
+        // The code bellow could look tricky because of it is bottle neck so here is put some attempt including
+        // 1. Minimize amount of descriptors is loaded on each step
+        // 2. Reduce memory pollution
+        for (i in 1 until nameSegments.size) {
+            val current = Name.guessByFirstCharacter(nameSegments[i])
+            acc = acc.flatMap { container ->
                 val classDescriptor = container as? ClassDescriptor ?: return@flatMap emptyList<DeclarationDescriptor>()
-                classDescriptor.constructors.filter { it.name.asString() == current } +
+                val isLeaf = i == lastIndex
-                        classDescriptor.unsubstitutedMemberScope.getDescriptorsFiltered { name -> name.asString() == current } +
+                val memberScope = classDescriptor.unsubstitutedMemberScope
-                        // Static scope is required only for Enum.values() and Enum.valueOf().
+
-                        classDescriptor.staticScope.getDescriptorsFiltered { name -> name.asString() == current }
+                val classifier = memberScope.getContributedClassifier(current, NoLookupLocation.FROM_BACKEND)
+                if (!isLeaf) {
+                    classifier?.let { listOf(it) } ?: emptyList()
+                } else {
+                    mutableListOf<DeclarationDescriptor>().apply {
+                        addIfNotNull(classifier)
+                        if (signature.id != null) {
+                            if (isConstructorName(current)) addAll(classDescriptor.constructors)
+                            addAll(memberScope.getContributedFunctions(current, NoLookupLocation.FROM_BACKEND))
+                            addAll(memberScope.getContributedVariables(current, NoLookupLocation.FROM_BACKEND))
+                            addAll(classDescriptor.staticScope.getContributedDescriptors { it == current })
+                        }
+                    }
+                }
             }
         }
+        val candidates = acc
 
         return when (candidates.size) {
             1 -> candidates.first()