CliLightClassGenerationSupport: correctly find facade files

This fixes a problem where JvmPackageName annotation could force file
    facades to be in the package different to declared kotlin package

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/CliLightClassGenerationSupport.kt

@@ -27,7 +27,10 @@ import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.asJava.LightClassBuilder
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
-import org.jetbrains.kotlin.asJava.builder.*
+import org.jetbrains.kotlin.asJava.builder.InvalidLightClassDataHolder
+import org.jetbrains.kotlin.asJava.builder.LightClassConstructionContext
+import org.jetbrains.kotlin.asJava.builder.LightClassDataHolder
+import org.jetbrains.kotlin.asJava.builder.LightClassDataHolderImpl
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForScript
@@ -35,8 +38,8 @@ import org.jetbrains.kotlin.asJava.classes.KtLightClassForSourceDeclaration
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.PackageViewDescriptor
-import org.jetbrains.kotlin.fileClasses.JvmFileClassUtil
-import org.jetbrains.kotlin.load.kotlin.PackagePartClassUtils
+import org.jetbrains.kotlin.fileClasses.javaFileFacadeFqName
+import org.jetbrains.kotlin.load.java.components.FilesByFacadeFqNameIndexer
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.*
@@ -190,9 +193,9 @@ class CliLightClassGenerationSupport(project: Project) : LightClassGenerationSup
     override fun findFilesForFacade(facadeFqName: FqName, scope: GlobalSearchScope): Collection<KtFile> {
         if (facadeFqName.isRoot) return emptyList()
 
-        return PackagePartClassUtils.getFilesWithCallables(findFilesForPackage(facadeFqName.parent(), scope)).filter {
-            JvmFileClassUtil.getFileClassInfoNoResolve(it).facadeClassFqName == facadeFqName
-        }
+        return bindingContext.get(FilesByFacadeFqNameIndexer.FACADE_FILES_BY_FQ_NAME, facadeFqName)?.filter {
+            PsiSearchScopeUtil.isInScope(scope, it)
+        } ?: emptyList()
     }
 
     override fun createTrace(): BindingTraceContext {
@@ -244,14 +247,20 @@ class CliLightClassGenerationSupport(project: Project) : LightClassGenerationSup
     }
 
     override fun getFacadeClassesInPackage(packageFqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> {
-        return PackagePartClassUtils.getFilesWithCallables(findFilesForPackage(packageFqName, scope)).groupBy {
-            JvmFileClassUtil.getFileClassInfoNoResolve(it).facadeClassFqName
-        }.mapNotNull { KtLightClassForFacade.createForFacade(psiManager, it.key, scope, it.value) }
+        return findFacadeFilesInPackage(packageFqName, scope)
+            .groupBy { it.javaFileFacadeFqName }
+            .mapNotNull { KtLightClassForFacade.createForFacade(psiManager, it.key, scope, it.value) }
     }
 
     override fun getFacadeNames(packageFqName: FqName, scope: GlobalSearchScope): Collection<String> {
-        return PackagePartClassUtils.getFilesWithCallables(findFilesForPackage(packageFqName, scope)).map {
-            JvmFileClassUtil.getFileClassInfoNoResolve(it).facadeClassFqName.shortName().asString()
-        }
+        return findFacadeFilesInPackage(packageFqName, scope)
+            .map { it.javaFileFacadeFqName.shortName().asString() }
     }
+
+    private fun findFacadeFilesInPackage(
+        packageFqName: FqName,
+        scope: GlobalSearchScope
+    ) = bindingContext.get(FilesByFacadeFqNameIndexer.FACADE_FILES_BY_PACKAGE_NAME, packageFqName)
+        ?.filter { PsiSearchScopeUtil.isInScope(scope, it) }
+        .orEmpty()
 }

compiler/frontend.java/src/org/jetbrains/kotlin/frontend/java/di/injection.kt

@@ -120,6 +120,7 @@ fun createContainerForLazyResolveWithJava(
     targetEnvironment.configure(this)
 
     useImpl<ContractDeserializerImpl>()
+    useImpl<FilesByFacadeFqNameIndexer>()
 }.apply {
     get<AbstractJavaClassFinder>().initialize(bindingTrace, get<KotlinCodeAnalyzer>())
 }

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/components/FilesByFacadeFqNameIndexer.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.load.java.components
+
+import org.jetbrains.kotlin.fileClasses.javaFileFacadeFqName
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.FilePreprocessorExtension
+import org.jetbrains.kotlin.resolve.addElementToSlice
+import org.jetbrains.kotlin.util.slicedMap.Slices
+import org.jetbrains.kotlin.util.slicedMap.WritableSlice
+
+// TODO: this component is actually only needed by CLI, see CliLightClassGenerationSupport
+class FilesByFacadeFqNameIndexer(private val trace: BindingTrace) : FilePreprocessorExtension {
+    override fun preprocessFile(file: KtFile) {
+        if (!file.hasTopLevelCallables()) return
+
+        trace.addElementToSlice(FACADE_FILES_BY_FQ_NAME, file.javaFileFacadeFqName, file)
+        trace.addElementToSlice(FACADE_FILES_BY_PACKAGE_NAME, file.javaFileFacadeFqName.parent(), file)
+    }
+
+    companion object {
+        val FACADE_FILES_BY_FQ_NAME: WritableSlice<FqName, MutableCollection<KtFile>> = Slices.createSimpleSlice()
+        val FACADE_FILES_BY_PACKAGE_NAME: WritableSlice<FqName, MutableCollection<KtFile>> = Slices.createSimpleSlice()
+    }
+}

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/JvmBindingContextSlices.kt

@@ -17,9 +17,10 @@
 package org.jetbrains.kotlin.resolve.jvm
 
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
-import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.util.slicedMap.BasicWritableSlice
 import org.jetbrains.kotlin.util.slicedMap.RewritePolicy
 import org.jetbrains.kotlin.util.slicedMap.Slices

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.resolve
 import com.google.common.collect.Sets
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext.PACKAGE_TO_FILES
+import org.jetbrains.kotlin.util.slicedMap.WritableSlice
 
 interface FilePreprocessorExtension {
     fun preprocessFile(file: KtFile)
@@ -28,9 +29,14 @@ class FilePreprocessor(
     private fun registerFileByPackage(file: KtFile) {
         // Register files corresponding to this package
         // The trace currently does not support bi-di multimaps that would handle this task nicer
-        val fqName = file.packageFqName
-        val files = trace.get(PACKAGE_TO_FILES, fqName) ?: Sets.newIdentityHashSet()
-        files.add(file)
-        trace.record(PACKAGE_TO_FILES, fqName, files)
+        trace.addElementToSlice(PACKAGE_TO_FILES, file.packageFqName, file)
     }
-}
+}
+
+fun <K, T> BindingTrace.addElementToSlice(
+    slice: WritableSlice<K, MutableCollection<T>>, key: K, element: T
+) {
+    val elements = get(slice, key) ?: Sets.newIdentityHashSet()
+    elements.add(element)
+    record(slice, key, elements)
+}

compiler/light-classes/src/org/jetbrains/kotlin/asJava/finder/JavaElementFinder.java

@@ -202,6 +202,7 @@ public class JavaElementFinder extends PsiElementFinder implements KotlinFinderM
     @NotNull
     public PsiFile[] getPackageFiles(@NotNull PsiPackage psiPackage, @NotNull GlobalSearchScope scope) {
         FqName packageFQN = new FqName(psiPackage.getQualifiedName());
+        // TODO: this does not take into account JvmPackageName annotation
         Collection<KtFile> result = lightClassGenerationSupport.findFilesForPackage(packageFQN, scope);
         return result.toArray(new PsiFile[result.size()]);
     }