Combine Java package scope and Kotlin scopes in JvmPackageScope

Use it in LazyJavaPackageFragment#scope instead of a weird lazy chained scope.
Also move lookup tracker records to it, to prevent them from being written in
each package scope individually (this was hurting performance)

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/JvmPackageScope.kt

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.load.java.lazy.descriptors
+
+import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.incremental.components.LookupLocation
+import org.jetbrains.kotlin.incremental.record
+import org.jetbrains.kotlin.load.java.lazy.LazyJavaResolverContext
+import org.jetbrains.kotlin.load.java.structure.JavaPackage
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
+import org.jetbrains.kotlin.resolve.scopes.MemberScope
+import org.jetbrains.kotlin.storage.getValue
+import org.jetbrains.kotlin.util.collectionUtils.getFirstMatch
+import org.jetbrains.kotlin.util.collectionUtils.getFromAllScopes
+import org.jetbrains.kotlin.utils.Printer
+import org.jetbrains.kotlin.utils.toReadOnlyList
+
+class JvmPackageScope(
+        private val c: LazyJavaResolverContext,
+        private val jPackage: JavaPackage,
+        private val packageFragment: LazyJavaPackageFragment
+): MemberScope {
+    internal val javaScope = LazyJavaPackageScope(c, jPackage, packageFragment)
+
+    private val kotlinScopes by c.storageManager.createLazyValue {
+        packageFragment.binaryClasses.values.mapNotNull { partClass ->
+            c.components.deserializedDescriptorResolver.createKotlinPackagePartScope(packageFragment, partClass)
+        }.toReadOnlyList()
+    }
+
+    override fun getContributedClassifier(name: Name, location: LookupLocation): ClassifierDescriptor? {
+        recordLookup(location, name)
+
+        val javaClassifier = javaScope.getContributedClassifier(name, location)
+        if (javaClassifier != null) return javaClassifier
+
+        return getFirstMatch(kotlinScopes) { it.getContributedClassifier(name, location) }
+    }
+
+    override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
+        recordLookup(location, name)
+        return getFromAllScopes(javaScope, kotlinScopes) { it.getContributedVariables(name, location) }
+    }
+
+    override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> {
+        recordLookup(location, name)
+        return getFromAllScopes(javaScope, kotlinScopes) { it.getContributedFunctions(name, location) }
+    }
+
+    override fun getContributedDescriptors(
+            kindFilter: DescriptorKindFilter, nameFilter: (Name) -> Boolean
+    ): Collection<DeclarationDescriptor> =
+            getFromAllScopes(javaScope, kotlinScopes) { it.getContributedDescriptors(kindFilter, nameFilter) }
+
+    override fun printScopeStructure(p: Printer) {
+        p.println(javaClass.simpleName, " {")
+        p.pushIndent()
+
+        p.println("containingDeclaration: $packageFragment")
+        javaScope.printScopeStructure(p)
+
+        for (kotlinScope in kotlinScopes) {
+            kotlinScope.printScopeStructure(p)
+        }
+
+        p.popIndent()
+        p.println("}")
+    }
+
+    private fun recordLookup(location: LookupLocation, name: Name) {
+        c.components.lookupTracker.record(location, packageFragment, name)
+    }
+}

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaClassMemberScope.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.impl.EnumEntrySyntheticClassDescriptor
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.incremental.record
 import org.jetbrains.kotlin.load.java.*
 import org.jetbrains.kotlin.load.java.BuiltinMethodsWithDifferentJvmName.isRemoveAtByIndex
 import org.jetbrains.kotlin.load.java.BuiltinMethodsWithDifferentJvmName.sameAsRenamedInJvmBuiltin
@@ -659,6 +660,16 @@ class LazyJavaClassMemberScope(
         return nestedClasses(name)
     }
 
+    override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> {
+        recordLookup(name, location)
+        return super.getContributedFunctions(name, location)
+    }
+
+    override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
+        recordLookup(name, location)
+        return super.getContributedVariables(name, location)
+    }
+
     override fun getClassNames(kindFilter: DescriptorKindFilter, nameFilter: (Name) -> Boolean): Collection<Name>
             = nestedClassIndex().keys + enumEntryIndex().keys
 
@@ -673,5 +684,9 @@ class LazyJavaClassMemberScope(
         }
     }
 
+    private fun recordLookup(name: Name, from: LookupLocation) {
+        c.components.lookupTracker.record(from, ownerDescriptor, name)
+    }
+
     override fun toString() = "Lazy Java member scope for " + jClass.fqName
 }

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaPackageFragment.kt

@@ -27,8 +27,6 @@ import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.scopes.ChainedMemberScope
-import org.jetbrains.kotlin.resolve.scopes.LazyScopeAdapter
 import org.jetbrains.kotlin.storage.getValue
 
 class LazyJavaPackageFragment(
@@ -42,17 +40,7 @@ class LazyJavaPackageFragment(
         }.toMap()
     }
 
-    private val javaScope by c.storageManager.createLazyValue {
+    private val scope = JvmPackageScope(c, jPackage, this)
-        LazyJavaPackageScope(c, jPackage, this)
-    }
-
-    private val scope by c.storageManager.createLazyValue {
-        ChainedMemberScope.create("Java + Deserialized Kotlin scope", listOf(javaScope) + LazyScopeAdapter(c.storageManager.createLazyValue {
-            ChainedMemberScope.create("Deserialized Kotlin scope", binaryClasses.values.mapNotNull { partClass ->
-                c.components.deserializedDescriptorResolver.createKotlinPackagePartScope(this, partClass)
-            })
-        }))
-    }
 
     private val subPackages = c.storageManager.createRecursionTolerantLazyValue(
             { jPackage.subPackages.map(JavaPackage::fqName) },
@@ -62,7 +50,7 @@ class LazyJavaPackageFragment(
 
     internal fun getSubPackageFqNames(): List<FqName> = subPackages()
 
-    internal fun findClassifierByJavaClass(jClass: JavaClass): ClassDescriptor? = javaScope.findClassifierByJavaClass(jClass)
+    internal fun findClassifierByJavaClass(jClass: JavaClass): ClassDescriptor? = scope.javaScope.findClassifierByJavaClass(jClass)
 
     private val partToFacade by c.storageManager.createLazyValue {
         val result = hashMapOf<String, String>()

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaPackageScope.kt

@@ -103,13 +103,11 @@ class LazyJavaPackageScope(
         override fun hashCode() = name.hashCode()
     }
 
-    override fun getContributedClassifier(name: Name, location: LookupLocation) = findClassifier(name, null, location)
+    override fun getContributedClassifier(name: Name, location: LookupLocation) = findClassifier(name, null)
 
-    private fun findClassifier(name: Name, javaClass: JavaClass?, location: LookupLocation): ClassDescriptor? {
+    private fun findClassifier(name: Name, javaClass: JavaClass?): ClassDescriptor? {
         if (!SpecialNames.isSafeIdentifier(name)) return null
 
-        recordLookup(name, location)
-
         val knownClassNamesInPackage = knownClassNamesInPackage()
         if (javaClass == null && knownClassNamesInPackage != null && name.asString() !in knownClassNamesInPackage) {
             return null
@@ -118,8 +116,7 @@ class LazyJavaPackageScope(
         return classes(FindClassRequest(name, javaClass))
     }
 
-    internal fun findClassifierByJavaClass(javaClass: JavaClass) =
+    internal fun findClassifierByJavaClass(javaClass: JavaClass) = findClassifier(javaClass.name, javaClass)
-            findClassifier(javaClass.name, javaClass, NoLookupLocation.FROM_JAVA_LOADER)
 
     override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> = emptyList()
 

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaScope.kt

@@ -22,7 +22,6 @@ import org.jetbrains.kotlin.descriptors.impl.PropertyDescriptorImpl
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
-import org.jetbrains.kotlin.incremental.record
 import org.jetbrains.kotlin.load.java.components.TypeUsage
 import org.jetbrains.kotlin.load.java.descriptors.JavaMethodDescriptor
 import org.jetbrains.kotlin.load.java.descriptors.JavaPropertyDescriptor
@@ -214,10 +213,7 @@ abstract class LazyJavaScope(protected val c: LazyJavaResolverContext) : MemberS
         return ResolvedValueParameters(descriptors, synthesizedNames)
     }
 
-    override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> {
+    override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> = functions(name)
-        recordLookup(name, location)
-        return functions(name)
-    }
 
     protected open fun getFunctionNames(kindFilter: DescriptorKindFilter, nameFilter: (Name) -> Boolean): Collection<Name>
             = memberIndex().getMethodNames(nameFilter)
@@ -291,13 +287,9 @@ abstract class LazyJavaScope(protected val c: LazyJavaResolverContext) : MemberS
         return propertyType
     }
 
-    override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
+    override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> = properties(name)
-        recordLookup(name, location)
-        return properties(name)
-    }
 
-    override fun getContributedDescriptors(kindFilter: DescriptorKindFilter,
+    override fun getContributedDescriptors(kindFilter: DescriptorKindFilter, nameFilter: (Name) -> Boolean) = allDescriptors()
-                                           nameFilter: (Name) -> Boolean) = allDescriptors()
 
     protected fun computeDescriptors(
             kindFilter: DescriptorKindFilter,
@@ -336,19 +328,15 @@ abstract class LazyJavaScope(protected val c: LazyJavaResolverContext) : MemberS
 
     protected abstract fun getClassNames(kindFilter: DescriptorKindFilter, nameFilter: (Name) -> Boolean): Collection<Name>
 
-    override fun toString() = "Lazy scope for ${ownerDescriptor}"
+    override fun toString() = "Lazy scope for $ownerDescriptor"
     
     override fun printScopeStructure(p: Printer) {
         p.println(javaClass.simpleName, " {")
         p.pushIndent()
 
-        p.println("containingDeclaration: ${ownerDescriptor}")
+        p.println("containingDeclaration: $ownerDescriptor")
 
         p.popIndent()
         p.println("}")
     }
-
-    protected fun recordLookup(name: Name, from: LookupLocation) {
-        c.components.lookupTracker.record(from, ownerDescriptor, name)
-    }
 }

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

@@ -62,6 +62,14 @@ inline fun <Scope, T> getFromAllScopes(scopes: List<Scope>, callback: (Scope) ->
     return result ?: emptySet()
 }
 
+inline fun <Scope, T> getFromAllScopes(firstScope: Scope, restScopes: List<Scope>, callback: (Scope) -> Collection<T>): Collection<T> {
+    var result: Collection<T>? = callback(firstScope)
+    for (scope in restScopes) {
+        result = result.concat(callback(scope))
+    }
+    return result ?: emptySet()
+}
+
 inline fun <Scope, T : Any> getFirstMatch(scopes: List<Scope>, callback: (Scope) -> T?): T? {
     // NOTE: This is performance-sensitive; please don't replace with map().firstOrNull()
     for (scope in scopes) {

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/descriptors/DeserializedClassDescriptor.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.descriptors.impl.AbstractClassDescriptor
 import org.jetbrains.kotlin.descriptors.impl.EnumEntrySyntheticClassDescriptor
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.incremental.record
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorFactory
 import org.jetbrains.kotlin.resolve.NonReportingOverrideStrategy
@@ -192,8 +193,19 @@ class DeserializedClassDescriptor(
             computeDescriptors(DescriptorKindFilter.ALL, MemberScope.ALL_NAME_FILTER, NoLookupLocation.WHEN_GET_ALL_DESCRIPTORS)
         }
 
-        override fun getContributedDescriptors(kindFilter: DescriptorKindFilter,
+        override fun getContributedDescriptors(
-                                               nameFilter: (Name) -> Boolean): Collection<DeclarationDescriptor> = allDescriptors()
+                kindFilter: DescriptorKindFilter, nameFilter: (Name) -> Boolean
+        ): Collection<DeclarationDescriptor> = allDescriptors()
+
+        override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> {
+            recordLookup(name, location)
+            return super.getContributedFunctions(name, location)
+        }
+
+        override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
+            recordLookup(name, location)
+            return super.getContributedVariables(name, location)
+        }
 
         override fun computeNonDeclaredFunctions(name: Name, functions: MutableCollection<SimpleFunctionDescriptor>) {
             val fromSupertypes = ArrayList<SimpleFunctionDescriptor>()
@@ -257,6 +269,10 @@ class DeserializedClassDescriptor(
         override fun addEnumEntryDescriptors(result: MutableCollection<DeclarationDescriptor>, nameFilter: (Name) -> Boolean) {
             result.addAll(classDescriptor.enumEntries?.all().orEmpty())
         }
+
+        private fun recordLookup(name: Name, from: LookupLocation) {
+            c.components.lookupTracker.record(from, c.containingDeclaration, name)
+        }
     }
 
     private inner class NestedClassDescriptors {

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/descriptors/DeserializedMemberScope.kt

@@ -17,9 +17,11 @@
 package org.jetbrains.kotlin.serialization.deserialization.descriptors
 
 import com.google.protobuf.MessageLite
-import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.incremental.components.LookupLocation
-import org.jetbrains.kotlin.incremental.record
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.MemberScopeImpl
@@ -78,10 +80,7 @@ abstract class DeserializedMemberScope protected constructor(
     protected open fun computeNonDeclaredFunctions(name: Name, functions: MutableCollection<SimpleFunctionDescriptor>) {
     }
 
-    override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> {
+    override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> = functions(name)
-        recordLookup(name, location)
-        return functions(name)
-    }
 
     private fun computeProperties(name: Name): Collection<PropertyDescriptor> {
         val protos = propertyProtos()[ProtoKey(name, isExtension = false)].orEmpty() +
@@ -98,10 +97,7 @@ abstract class DeserializedMemberScope protected constructor(
     protected open fun computeNonDeclaredProperties(name: Name, descriptors: MutableCollection<PropertyDescriptor>) {
     }
 
-    override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
+    override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> = properties(name)
-        recordLookup(name, location)
-        return properties.invoke(name)
-    }
 
     protected abstract fun addClassDescriptors(result: MutableCollection<DeclarationDescriptor>, nameFilter: (Name) -> Boolean)
 
@@ -177,8 +173,4 @@ abstract class DeserializedMemberScope protected constructor(
         p.popIndent()
         p.println("}")
     }
-
-    protected fun recordLookup(name: Name, from: LookupLocation) {
-        c.components.lookupTracker.record(from, c.containingDeclaration, name)
-    }
 }

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/descriptors/DeserializedPackageMemberScope.kt

@@ -51,7 +51,6 @@ open class DeserializedPackageMemberScope(
     override fun getContributedClassifier(name: Name, location: LookupLocation): ClassifierDescriptor? {
         if (SpecialNames.isSafeIdentifier(name) &&
             (name in classNames || c.components.fictitiousClassDescriptorFactory.shouldCreateClass(packageFqName, name))) {
-            recordLookup(name, location)
             return getClassDescriptor(name)
         }
         return null