Optimized task creation for local functions and variables

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tasks/CallableDescriptorCollectors.kt

@@ -27,6 +27,9 @@ import org.jetbrains.kotlin.resolve.calls.tasks.createSynthesizedInvokes
 import org.jetbrains.kotlin.resolve.calls.util.FakeCallableDescriptorForObject
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasClassObjectType
 import org.jetbrains.kotlin.resolve.scopes.JetScope
+import org.jetbrains.kotlin.resolve.scopes.LexicalScope
+import org.jetbrains.kotlin.resolve.scopes.utils.collectAllFromMeAndParent
+import org.jetbrains.kotlin.resolve.scopes.utils.getLocalVariable
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.types.JetType
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
@@ -34,6 +37,8 @@ import org.jetbrains.kotlin.utils.singletonOrEmptyList
 
 public interface CallableDescriptorCollector<D : CallableDescriptor> {
 
+    public fun getLocalNonExtensionsByName(lexicalScope: LexicalScope, name: Name, location: LookupLocation): Collection<D>
+
     public fun getNonExtensionsByName(scope: JetScope, name: Name, location: LookupLocation): Collection<D>
 
     public fun getMembersByName(receiver: JetType, name: Name, location: LookupLocation): Collection<D>
@@ -71,6 +76,17 @@ public fun <D : CallableDescriptor> CallableDescriptorCollectors<D>.filtered(fil
         CallableDescriptorCollectors(this.collectors.map { it.filtered(filter) })
 
 private object FunctionCollector : CallableDescriptorCollector<FunctionDescriptor> {
+    override fun getLocalNonExtensionsByName(lexicalScope: LexicalScope, name: Name, location: LookupLocation): Collection<FunctionDescriptor> {
+        return lexicalScope.collectAllFromMeAndParent {
+            if (it.ownerDescriptor is FunctionDescriptor) {
+                it.getDeclaredFunctions(name, location).filter { it.extensionReceiverParameter == null } +
+                    getConstructors(it.getDeclaredClassifier(name, location))
+            }
+            else {
+                emptyList()
+            }
+        }
+    }
 
     override fun getNonExtensionsByName(scope: JetScope, name: Name, location: LookupLocation): Collection<FunctionDescriptor> {
         return scope.getFunctions(name, location).filter { it.extensionReceiverParameter == null } + getConstructors(scope, name, location)
@@ -117,6 +133,13 @@ private object FunctionCollector : CallableDescriptorCollector<FunctionDescripto
             filterClassPredicate: (ClassDescriptor) -> Boolean = { true }
     ): Collection<FunctionDescriptor> {
         val classifier = scope.getClassifier(name, location)
+        return getConstructors(classifier, filterClassPredicate)
+    }
+
+    private fun getConstructors(
+            classifier: ClassifierDescriptor?,
+            filterClassPredicate: (ClassDescriptor) -> Boolean = { true }
+    ): Collection<FunctionDescriptor> {
         if (classifier !is ClassDescriptor || ErrorUtils.isError(classifier) || !filterClassPredicate(classifier)
             // Constructors of singletons shouldn't be callable from the code
             || classifier.kind.isSingleton) {
@@ -129,6 +152,9 @@ private object FunctionCollector : CallableDescriptorCollector<FunctionDescripto
 }
 
 private object VariableCollector : CallableDescriptorCollector<VariableDescriptor> {
+    override fun getLocalNonExtensionsByName(lexicalScope: LexicalScope, name: Name, location: LookupLocation): Collection<VariableDescriptor> {
+        return listOfNotNull(lexicalScope.getLocalVariable(name))
+    }
 
     private fun getFakeDescriptorForObject(scope: JetScope, name: Name, location: LookupLocation): VariableDescriptor? {
         val classifier = scope.getClassifier(name, location)
@@ -172,6 +198,10 @@ private object PropertyCollector : CallableDescriptorCollector<VariableDescripto
     private fun filterProperties(variableDescriptors: Collection<VariableDescriptor>) =
             variableDescriptors.filter { it is PropertyDescriptor }
 
+    override fun getLocalNonExtensionsByName(lexicalScope: LexicalScope, name: Name, location: LookupLocation): Collection<VariableDescriptor> {
+        return filterProperties(VARIABLES_COLLECTOR.getLocalNonExtensionsByName(lexicalScope, name, location))
+    }
+
     override fun getNonExtensionsByName(scope: JetScope, name: Name, location: LookupLocation): Collection<VariableDescriptor> {
         return filterProperties(VARIABLES_COLLECTOR.getNonExtensionsByName(scope, name, location))
     }
@@ -194,6 +224,10 @@ private object PropertyCollector : CallableDescriptorCollector<VariableDescripto
 private fun <D : CallableDescriptor> CallableDescriptorCollector<D>.filtered(filter: (D) -> Boolean): CallableDescriptorCollector<D> {
     val delegate = this
     return object : CallableDescriptorCollector<D> {
+        override fun getLocalNonExtensionsByName(lexicalScope: LexicalScope, name: Name, location: LookupLocation): Collection<D> {
+            return delegate.getLocalNonExtensionsByName(lexicalScope, name, location).filter(filter)
+        }
+
         override fun getNonExtensionsByName(scope: JetScope, name: Name, location: LookupLocation): Collection<D> {
             return delegate.getNonExtensionsByName(scope, name, location).filter(filter)
         }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tasks/ResolutionTaskHolder.kt

@@ -37,15 +37,6 @@ public class ResolutionTaskHolder<D : CallableDescriptor, F : D>(
         candidatesList.add(storageManager.createLazyValue { lazyCandidates().toReadOnlyList() })
     }
 
-    public fun addCandidates(candidatesList: List<Collection<ResolutionCandidate<D>>>) {
-        assertNotFinished()
-        for (candidates in candidatesList) {
-            if (candidates.isNotEmpty()) {
-                addCandidates { candidates }
-            }
-        }
-    }
-
     private fun assertNotFinished() {
         assert(internalTasks == null, "Can't add candidates after the resulting tasks were computed.")
     }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tasks/TaskPrioritizer.kt

@@ -56,22 +56,6 @@ public class TaskPrioritizer(
         private val smartCastManager: SmartCastManager
 ) {
 
-    public fun <D : CallableDescriptor> splitLexicallyLocalDescriptors(
-            allDescriptors: Collection<ResolutionCandidate<D>>,
-            containerOfTheCurrentLocality: DeclarationDescriptor,
-            local: MutableCollection<ResolutionCandidate<D>>,
-            nonlocal: MutableCollection<ResolutionCandidate<D>>
-    ) {
-        for (resolvedCall in allDescriptors) {
-            if (ExpressionTypingUtils.isLocal(containerOfTheCurrentLocality, resolvedCall.getDescriptor())) {
-                local.add(resolvedCall)
-            }
-            else {
-                nonlocal.add(resolvedCall)
-            }
-        }
-    }
-
     public fun <D : CallableDescriptor, F : D> computePrioritizedTasks(
             context: BasicCallResolutionContext,
             name: Name,
@@ -264,28 +248,18 @@ public class TaskPrioritizer(
             implicitReceivers: Collection<ReceiverValue>,
             c: TaskPrioritizerContext<D, F>
     ) {
-        val localsList = Lists.newArrayList<Collection<ResolutionCandidate<D>>>()
-        val nonlocalsList = Lists.newArrayList<Collection<ResolutionCandidate<D>>>()
-        for (callableDescriptorCollector in c.callableDescriptorCollectors) {
-
-            val members = convertWithImpliedThisAndNoReceiver(
-                    c.scope,
-                    callableDescriptorCollector.getNonExtensionsByName(c.scope.asJetScope(), c.name, createLookupLocation(c)),
-                    c.context.call
-            )
-
-            if (members.isNotEmpty()) {
-                val nonlocals = Lists.newArrayList<ResolutionCandidate<D>>()
-                val locals = Lists.newArrayList<ResolutionCandidate<D>>()
-                splitLexicallyLocalDescriptors(members, c.scope.ownerDescriptor, locals, nonlocals)
-
-                localsList.add(locals)
-                nonlocalsList.add(nonlocals)
-            }
-        }
+        val lookupLocation = createLookupLocation(c)
 
         //locals
-        c.result.addCandidates(localsList)
+        c.callableDescriptorCollectors.forEach {
+            c.result.addCandidates {
+                convertWithImpliedThisAndNoReceiver(
+                        c.scope,
+                        it.getLocalNonExtensionsByName(c.scope, c.name, lookupLocation),
+                        c.context.call
+                )
+            }
+        }
 
         val implicitReceiversWithTypes = implicitReceivers.map { ReceiverWithTypes(it, c.context) }
 
@@ -295,7 +269,13 @@ public class TaskPrioritizer(
         }
 
         //nonlocals
-        c.result.addCandidates(nonlocalsList)
+        c.callableDescriptorCollectors.forEach {
+            c.result.addCandidates {
+                val descriptors = it.getNonExtensionsByName(c.scope.asJetScope(), c.name, lookupLocation)
+                        .filter { !ExpressionTypingUtils.isLocal(c.scope.ownerDescriptor, it) }
+                convertWithImpliedThisAndNoReceiver(c.scope, descriptors, c.context.call)
+            }
+        }
 
         //static (only for better error reporting)
         for (implicitReceiver in implicitReceiversWithTypes) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/scopes/utils/ScopeUtils.kt

@@ -266,7 +266,7 @@ private inline fun <T: Any> LexicalScope.collectFromMeAndParent(
     return result ?: emptyList()
 }
 
-private inline fun <T: Any> LexicalScope.collectAllFromMeAndParent(
+internal inline fun <T: Any> LexicalScope.collectAllFromMeAndParent(
         collect: (LexicalScope) -> Collection<T>
 ): Collection<T> {
     var result: Collection<T>? = null