Optimize AllUnderImportScope

compiler/frontend/src/org/jetbrains/kotlin/resolve/AllUnderImportScope.kt

@@ -20,28 +20,19 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.scopes.BaseImportingScope
+import org.jetbrains.kotlin.resolve.scopes.*
-import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.MemberScope
-import org.jetbrains.kotlin.resolve.scopes.computeAllNames
 import org.jetbrains.kotlin.util.collectionUtils.flatMapScopes
-import org.jetbrains.kotlin.util.collectionUtils.listOfNonEmptyScopes
+import org.jetbrains.kotlin.util.collectionUtils.forEachScope
 import org.jetbrains.kotlin.utils.Printer
-import org.jetbrains.kotlin.utils.addToStdlib.flatMapToNullable
+import org.jetbrains.kotlin.utils.SmartList
 
-class AllUnderImportScope(
+class AllUnderImportScope private constructor(
     descriptor: DeclarationDescriptor,
-    excludedImportNames: Collection<FqName>
+    excludedImportNames: Collection<FqName>,
+    private val scope1: MemberScope,
+    private val scope2: MemberScope?
 ) : BaseImportingScope(null) {
 
-    private val scopes: Array<MemberScope> = if (descriptor is ClassDescriptor) {
-        listOfNonEmptyScopes(descriptor.staticScope, descriptor.unsubstitutedInnerClassesScope).toTypedArray()
-    } else {
-        assert(descriptor is PackageViewDescriptor) {
-            "Must be class or package view descriptor: $descriptor"
-        }
-        listOfNonEmptyScopes((descriptor as PackageViewDescriptor).memberScope).toTypedArray()
-    }
 
     private val excludedNames: Set<Name> = if (excludedImportNames.isEmpty()) { // optimization
         emptySet<Name>()
@@ -51,10 +42,22 @@ class AllUnderImportScope(
         excludedImportNames.mapNotNull { if (it.parent() == fqName) it.shortName() else null }.toSet()
     }
 
-    override fun computeImportedNames(): Set<Name>? = when (scopes.size) {
+    override fun computeImportedNames(): Set<Name>? {
-        0 -> null
+        val names1 = scope1.computeAllNames()
-        1 -> scopes[0].computeAllNames()
+        return when {
-        else -> scopes.asIterable().flatMapToNullable(hashSetOf(), MemberScope::computeAllNames)
+            scope2 == null -> names1
+            names1 == null -> null
+            else -> {
+                val names2 = scope2.computeAllNames()
+                when {
+                    names2 == null -> null
+                    names1.isEmpty() -> names2
+                    else -> names1.toMutableSet().also {
+                        it.addAll(names2)
+                    }
+                }
+            }
+        }
     }
 
     override fun getContributedDescriptors(
@@ -69,39 +72,61 @@ class AllUnderImportScope(
         }
 
         val noPackagesKindFilter = kindFilter.withoutKinds(DescriptorKindFilter.PACKAGES_MASK)
-        return scopes
+        val result = SmartList<DeclarationDescriptor>()
-            .flatMapScopes { it.getContributedDescriptors(noPackagesKindFilter, nameFilterToUse) }
+        forEachScope(scope1, scope2) { scope ->
-            .filter { it !is PackageViewDescriptor } // subpackages are not imported
+            scope.getContributedDescriptors(noPackagesKindFilter, nameFilterToUse)
+                .filterTo(result) { it !is PackageViewDescriptor }
+        }
+        return result
     }
 
     override fun getContributedClassifier(name: Name, location: LookupLocation): ClassifierDescriptor? {
         if (name in excludedNames) return null
-        var single: ClassifierDescriptor? = null
+        val classifier1 = scope1.getContributedClassifier(name, location)
-        for (scope in scopes) {
+        val classifier2 = scope2?.getContributedClassifier(name, location)
-            val res = scope.getContributedClassifier(name, location) ?: continue
+        return if (classifier1 != null && classifier2 != null) null else (classifier1 ?: classifier2)
-            if (single == null) single = res
-            else return null
-        }
-        return single
     }
 
     override fun getContributedVariables(name: Name, location: LookupLocation): Collection<VariableDescriptor> {
         if (name in excludedNames) return emptyList()
-        return scopes.flatMapScopes { it.getContributedVariables(name, location) }
+        return flatMapScopes(scope1, scope2) { it.getContributedVariables(name, location) }
     }
 
     override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<FunctionDescriptor> {
         if (name in excludedNames) return emptyList()
-        return scopes.flatMapScopes { it.getContributedFunctions(name, location) }
+        return flatMapScopes(scope1, scope2) { it.getContributedFunctions(name, location) }
     }
 
     override fun recordLookup(name: Name, location: LookupLocation) {
-        scopes.forEach { it.recordLookup(name, location) }
+        scope1.recordLookup(name, location)
+        scope2?.recordLookup(name, location)
     }
 
     override fun printStructure(p: Printer) {
         p.println(this::class.java.simpleName)
     }
+
+    companion object {
+        fun create(descriptor: DeclarationDescriptor, excludedImportNames: Collection<FqName>): ImportingScope {
+            val scope1 =
+                if (descriptor is ClassDescriptor) {
+                    descriptor.staticScope
+                } else {
+                    assert(descriptor is PackageViewDescriptor) {
+                        "Must be class or package view descriptor: $descriptor"
+                    }
+                    (descriptor as PackageViewDescriptor).memberScope
+                }
+
+            val scope2 =
+                if (descriptor is ClassDescriptor) {
+                    descriptor.unsubstitutedInnerClassesScope.takeIf { it !== MemberScope.Empty }
+                } else null
+
+            return if (scope1 === MemberScope.Empty) {
+                if (scope2 == null || scope2 === MemberScope.Empty) ImportingScope.Empty
+                else AllUnderImportScope(descriptor, excludedImportNames, scope2, null)
+            } else AllUnderImportScope(descriptor, excludedImportNames, scope1, scope2)
+        }
+    }
 }
-
-

compiler/frontend/src/org/jetbrains/kotlin/resolve/QualifiedExpressionResolver.kt

@@ -249,7 +249,7 @@ class QualifiedExpressionResolver(val languageVersionSettings: LanguageVersionSe
                 return null
             }
 
-            return AllUnderImportScope(packageOrClassDescriptor, excludedImportNames)
+            return AllUnderImportScope.create(packageOrClassDescriptor, excludedImportNames)
         } else {
             return processSingleImport(moduleDescriptor, trace, importDirective, path, lastPart, packageFragmentForCheck)
         }

core/descriptors/src/org/jetbrains/kotlin/util/scopeUtils.kt

@@ -64,12 +64,22 @@ inline fun <Scope, T> getFromAllScopes(firstScope: Scope, restScopes: Array<Scop
     return result ?: emptySet()
 }
 
-inline fun <Scope, R> Array<out Scope>.flatMapScopes(transform: (Scope) -> Collection<R>): Collection<R> =
+inline fun <Scope, R> flatMapScopes(scope1: Scope?, scope2: Scope?, transform: (Scope) -> Collection<R>): Collection<R> {
-    when (size) {
+    val results1 = if (scope1 != null) transform(scope1) else emptyList()
-        0 -> emptyList()
+    if (scope2 == null) return results1
-        1 -> transform(get(0))
+    else {
-        else -> flatMapTo(ArrayList<R>(), transform)
+        val results2 = transform(scope2)
+        if (results1.isEmpty()) return results2
+        else return results1.toMutableList().also {
+            it.addAll(results2)
+        }
     }
+}
+
+inline fun <Scope> forEachScope(scope1: Scope?, scope2: Scope?, action: (Scope) -> Unit) {
+    if (scope1 != null) action(scope1)
+    if (scope2 != null) action(scope2)
+}
 
 fun listOfNonEmptyScopes(vararg scopes: MemberScope?): SmartList<MemberScope> =
     scopes.filterTo(SmartList<MemberScope>()) { it != null && it !== MemberScope.Empty }