Search classes by JavaClass through package scope

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

@@ -68,12 +68,12 @@ class LazyJavaPackageFragmentProvider(
                 return outerClassScope?.getContributedClassifier(javaClass.name, NoLookupLocation.FROM_JAVA_LOADER) as? ClassDescriptor
             }
 
-            val kotlinResult = c.resolveKotlinBinaryClass(c.components.kotlinClassFinder.findKotlinClass(javaClass))
-            if (kotlinResult is KotlinClassLookupResult.Found) return kotlinResult.descriptor
-
             if (fqName == null) return null
 
-            return getPackageFragment(fqName.parent())?.resolveTopLevelClass(javaClass)
+            val packageFragment = getPackageFragment(fqName.parent()) ?: return null
+            return packageFragment
+                           .getMemberScope()
+                           .findClassifierByJavaClass(javaClass, NoLookupLocation.FROM_JAVA_LOADER)
         }
     }
 }

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

@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.load.java.lazy.descriptors
 import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.descriptors.impl.PackageFragmentDescriptorImpl
 import org.jetbrains.kotlin.load.java.lazy.LazyJavaResolverContext
-import org.jetbrains.kotlin.load.java.structure.JavaClass
 import org.jetbrains.kotlin.load.java.structure.JavaPackage
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinaryPackageSourceElement
 import org.jetbrains.kotlin.name.ClassId
@@ -34,11 +33,6 @@ class LazyJavaPackageFragment(
         LazyJavaPackageScope(c, jPackage, this)
     }
 
-    private val topLevelClasses = c.storageManager.createMemoizedFunctionWithNullableValues {
-        javaClass: JavaClass ->
-        LazyJavaClassDescriptor(c, this, javaClass.fqName!!, javaClass)
-    }
-
     internal val kotlinBinaryClasses by c.storageManager.createLazyValue {
         c.components.packageMapper.findPackageParts(fqName.asString()).mapNotNull {
             val classId = ClassId(fqName, Name.identifier(it))
@@ -46,8 +40,6 @@ class LazyJavaPackageFragment(
         }
     }
 
-    internal fun resolveTopLevelClass(javaClass: JavaClass) = topLevelClasses(javaClass)
-
     override fun getMemberScope() = scope
 
     override fun toString() = "lazy java package fragment: $fqName"

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

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.load.java.structure.JavaClass
 import org.jetbrains.kotlin.load.java.structure.JavaPackage
 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.name.SpecialNames
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
@@ -72,38 +73,65 @@ class LazyJavaPackageScope(
         c.components.deserializedDescriptorResolver.createKotlinPackageScope(ownerDescriptor, ownerDescriptor.kotlinBinaryClasses)
     }
 
-    private val classes = c.storageManager.createMemoizedFunctionWithNullableValues<Name, ClassDescriptor> { name ->
+    private val classes = c.storageManager.createMemoizedFunctionWithNullableValues<FindClassRequest, ClassDescriptor> { request ->
-        val classId = ClassId(ownerDescriptor.fqName, name)
+        val classId = ClassId(ownerDescriptor.fqName, request.name)
+
+        val kotlinBinaryClass =
+                // These branches should be semantically equal, but the first one could be faster
+                if (request.javaClass != null)
+                    c.components.kotlinClassFinder.findKotlinClass(request.javaClass)
+                else
+                    c.components.kotlinClassFinder.findKotlinClass(classId)
+
+        val kotlinResult = c.resolveKotlinBinaryClass(kotlinBinaryClass)
 
-        val kotlinResult = c.resolveKotlinBinaryClass(c.components.kotlinClassFinder.findKotlinClass(classId))
         when (kotlinResult) {
             is KotlinClassLookupResult.Found -> kotlinResult.descriptor
             is KotlinClassLookupResult.SyntheticClass -> null
             is KotlinClassLookupResult.NotFound -> {
-                c.components.finder.findClass(classId)?.let { javaClass ->
+
-                    c.javaClassResolver.resolveClass(javaClass).apply {
+                val javaClass = request.javaClass ?: c.components.finder.findClass(classId)
-                        assert(this == null || this.containingDeclaration == ownerDescriptor) {
+                javaClass?.let { it ->
-                            "Wrong package fragment for $this, expected $ownerDescriptor"
+                    LazyJavaClassDescriptor(c, ownerDescriptor, it.fqName!!, it)
-                        }
-                    }
                 }
             }
         }
     }
 
-    override fun getContributedClassifier(name: Name, location: LookupLocation): ClassifierDescriptor? {
+    // javaClass here is only for sake of optimizations
+    private class FindClassRequest(val name: Name, val javaClass: JavaClass?) {
+        override fun equals(other: Any?): Boolean{
+            if (this === other) return true
+
+            other as FindClassRequest
+
+            if (name != other.name) return false
+
+            return true
+        }
+
+        override fun hashCode(): Int{
+            return name.hashCode()
+        }
+    }
+
+    override fun getContributedClassifier(name: Name, location: LookupLocation) = findClassifier(name, null, location)
+
+    private fun findClassifier(name: Name, javaClass: JavaClass?, location: LookupLocation): ClassDescriptor? {
         if (!SpecialNames.isSafeIdentifier(name)) return null
 
         recordLookup(name, location)
 
         val knownClassNamesInPackage = knownClassNamesInPackage()
-        if (knownClassNamesInPackage != null && name.asString() !in knownClassNamesInPackage) {
+        if (javaClass == null && knownClassNamesInPackage != null && name.asString() !in knownClassNamesInPackage) {
             return null
         }
 
-        return classes(name)
+        return classes(FindClassRequest(name, javaClass))
     }
 
+    fun findClassifierByJavaClass(javaClass: JavaClass, location: LookupLocation) = findClassifier(javaClass.name, javaClass, location)
+
     override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
         // We should track lookups here because this scope can be used for kotlin packages too (if it doesn't contain toplevel properties nor functions).
         recordLookup(name, location)