Do not force loading all kt-files inside LazyPackageDescriptor constructor

Effectively it leads to all stubs being built for any once used package

The idea is that they only were used for FILE_TO_PACKAGE_FRAGMENT
that is mostly read in backend and rarely in the IDE, so we can replace
its usages with searching through related package-fragments

compiler/backend/src/org/jetbrains/kotlin/codegen/MultifileClassCodegen.kt

@@ -41,6 +41,7 @@ import org.jetbrains.kotlin.resolve.jvm.diagnostics.MultifileClass
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.MultifileClassPart
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.OtherOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
+import org.jetbrains.kotlin.resolve.lazy.descriptors.findPackageFragmentForFile
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedCallableMemberDescriptor
@@ -63,7 +64,7 @@ class MultifileClassCodegenImpl(
 ) : MultifileClassCodegen {
     private val facadeClassType = AsmUtil.asmTypeByFqNameWithoutInnerClasses(facadeFqName)
 
-    private val packageFragment = getOnlyPackageFragment(facadeFqName.parent(), files, state.bindingContext)
+    private val packageFragment = getOnlyPackageFragment(facadeFqName.parent(), files, state.module)
 
     private val compiledPackageFragment = getCompiledPackageFragment(facadeFqName, state)
 
@@ -352,10 +353,10 @@ class MultifileClassCodegenImpl(
         private val J_L_OBJECT = AsmTypes.OBJECT_TYPE.internalName
         private val FACADE_CLASS_ATTRIBUTES = Opcodes.ACC_PUBLIC or Opcodes.ACC_FINAL or Opcodes.ACC_SUPER
 
-        private fun getOnlyPackageFragment(packageFqName: FqName, files: Collection<KtFile>, bindingContext: BindingContext): PackageFragmentDescriptor? {
+        private fun getOnlyPackageFragment(packageFqName: FqName, files: Collection<KtFile>, moduleDescriptor: ModuleDescriptor): PackageFragmentDescriptor? {
             val fragments = SmartList<PackageFragmentDescriptor>()
             for (file in files) {
-                val fragment = bindingContext.get(BindingContext.FILE_TO_PACKAGE_FRAGMENT, file)
+                val fragment = moduleDescriptor.findPackageFragmentForFile(file)
                                ?: throw AssertionError("package fragment is null for " + file + "\n" + file.text)
 
                 assert(packageFqName == fragment.fqName) { "expected package fq name: " + packageFqName + ", actual: " + fragment.fqName }

compiler/backend/src/org/jetbrains/kotlin/codegen/PackageCodegenImpl.java

@@ -32,8 +32,8 @@ import org.jetbrains.kotlin.fileClasses.JvmFileClassUtil;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus;
 import org.jetbrains.kotlin.psi.*;
-import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOriginKt;
+import org.jetbrains.kotlin.resolve.lazy.descriptors.PackageDescriptorUtilKt;
 import org.jetbrains.org.objectweb.asm.Type;
 
 import java.util.ArrayList;
@@ -131,7 +131,8 @@ public class PackageCodegenImpl implements PackageCodegen {
     private PackageFragmentDescriptor getOnlyPackageFragment(@NotNull FqName expectedPackageFqName) {
         SmartList<PackageFragmentDescriptor> fragments = new SmartList<>();
         for (KtFile file : files) {
-            PackageFragmentDescriptor fragment = state.getBindingContext().get(BindingContext.FILE_TO_PACKAGE_FRAGMENT, file);
+            PackageFragmentDescriptor fragment =
+                    PackageDescriptorUtilKt.findPackageFragmentForFile(state.getModule(), file);
             assert fragment != null : "package fragment is null for " + file + "\n" + file.getText();
 
             assert expectedPackageFqName.equals(fragment.getFqName()) :

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/DelegatePackageMemberDeclarationProvider.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.cli.jvm.repl
 
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtDestructuringDeclarationEntry
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.lazy.declarations.PackageMemberDeclarationProvider
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 
@@ -28,6 +29,8 @@ open class DelegatePackageMemberDeclarationProvider(var delegate: PackageMemberD
 
     override fun getPackageFiles() = delegate.getPackageFiles()
 
+    override fun containsFile(file: KtFile) = delegate.containsFile(file)
+
     override fun getDeclarations(kindFilter: DescriptorKindFilter,
                                  nameFilter: (Name) -> Boolean) = delegate.getDeclarations(kindFilter, nameFilter)
 

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/ReplCodeAnalyzer.kt

@@ -31,7 +31,10 @@ import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.BindingTraceContext
+import org.jetbrains.kotlin.resolve.LazyTopDownAnalyzer
+import org.jetbrains.kotlin.resolve.TopDownAnalysisContext
+import org.jetbrains.kotlin.resolve.TopDownAnalysisMode
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfoFactory
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import org.jetbrains.kotlin.resolve.lazy.*
@@ -107,10 +110,6 @@ class ReplCodeAnalyzer(environment: KotlinCoreEnvironment) {
 
         val context = topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi))
 
-        if (trace.get(BindingContext.FILE_TO_PACKAGE_FRAGMENT, linePsi) == null) {
-            trace.record(BindingContext.FILE_TO_PACKAGE_FRAGMENT, linePsi, resolveSession.getPackageFragment(FqName.ROOT))
-        }
-
         val diagnostics = trace.bindingContext.diagnostics
         val hasErrors = diagnostics.any { it.severity == Severity.ERROR }
         return if (hasErrors) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -262,7 +262,6 @@ public interface BindingContext {
     WritableSlice<ClassDescriptor, FunctionDescriptor> DATA_CLASS_COPY_FUNCTION = Slices.createSimpleSlice();
 
     WritableSlice<FqNameUnsafe, ClassDescriptor> FQNAME_TO_CLASS_DESCRIPTOR = new BasicWritableSlice<>(DO_NOTHING, true);
-    WritableSlice<KtFile, PackageFragmentDescriptor> FILE_TO_PACKAGE_FRAGMENT = Slices.createSimpleSlice();
     WritableSlice<FqName, Collection<KtFile>> PACKAGE_TO_FILES = Slices.createSimpleSlice();
     WritableSlice<KtBinaryExpressionWithTypeRHS, Boolean> CAST_TYPE_USED_AS_EXPECTED_TYPE = Slices.createSimpleSlice();
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/declarations/CombinedPackageMemberDeclarationProvider.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.resolve.lazy.declarations
 
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtDestructuringDeclarationEntry
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 
 class CombinedPackageMemberDeclarationProvider(
@@ -27,6 +28,8 @@ class CombinedPackageMemberDeclarationProvider(
 
     override fun getPackageFiles() = providers.flatMap { it.getPackageFiles() }
 
+    override fun containsFile(file: KtFile) = providers.any { it.containsFile(file) }
+
     override fun getDeclarations(kindFilter: DescriptorKindFilter, nameFilter: (Name) -> Boolean)
             = providers.flatMap { it.getDeclarations(kindFilter, nameFilter) }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/declarations/FileBasedPackageMemberDeclarationProvider.kt

@@ -16,9 +16,9 @@
 
 package org.jetbrains.kotlin.resolve.lazy.declarations
 
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.storage.StorageManager
 
 class FileBasedPackageMemberDeclarationProvider(
@@ -45,6 +45,8 @@ class FileBasedPackageMemberDeclarationProvider(
 
     override fun getPackageFiles() = packageFiles
 
+    override fun containsFile(file: KtFile) = file in packageFiles
+
     override fun toString() = "Declarations for package $fqName with files ${packageFiles.map { it.name }} " +
                               "with declarations inside ${packageFiles.flatMap { it.declarations }.map { it.name ?: "???"}}"
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/declarations/PackageMemberDeclarationProvider.kt

@@ -16,13 +16,14 @@
 
 package org.jetbrains.kotlin.resolve.lazy.declarations
 
-import org.jetbrains.annotations.ReadOnly
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtFile
 
 interface PackageMemberDeclarationProvider : DeclarationProvider {
     fun getAllDeclaredSubPackages(nameFilter: (Name) -> Boolean): Collection<FqName>
 
     fun getPackageFiles(): Collection<KtFile>
+
+    fun containsFile(file: KtFile): Boolean
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/LazyPackageDescriptor.java

@@ -20,8 +20,6 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.PackageFragmentDescriptorImpl;
 import org.jetbrains.kotlin.name.FqName;
-import org.jetbrains.kotlin.psi.KtFile;
-import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil;
 import org.jetbrains.kotlin.resolve.lazy.LazyEntity;
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession;
@@ -42,10 +40,6 @@ public class LazyPackageDescriptor extends PackageFragmentDescriptorImpl impleme
         this.declarationProvider = declarationProvider;
 
         this.memberScope = new LazyPackageMemberScope(resolveSession, declarationProvider, this);
-
-        for (KtFile file : declarationProvider.getPackageFiles()) {
-            resolveSession.getTrace().record(BindingContext.FILE_TO_PACKAGE_FRAGMENT, file, this);
-        }
     }
 
     @NotNull

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/packageDescriptorUtil.kt

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2010-2017 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.resolve.lazy.descriptors
+
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
+import org.jetbrains.kotlin.psi.KtFile
+
+fun ModuleDescriptor.findPackageFragmentForFile(ktFile: KtFile): PackageFragmentDescriptor? =
+        getPackage(ktFile.packageFqName).fragments
+                .filterIsInstance<LazyPackageDescriptor>()
+                .firstOrNull { it.declarationProvider.containsFile(ktFile) }

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ModuleGenerator.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrModuleFragmentImpl
 import org.jetbrains.kotlin.ir.util.ExternalDependenciesGenerator
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.lazy.descriptors.findPackageFragmentForFile
 
 class ModuleGenerator(override val context: GeneratorContext) : Generator {
     fun generateModuleFragment(ktFiles: Collection<KtFile>): IrModuleFragment =
@@ -63,7 +64,7 @@ class ModuleGenerator(override val context: GeneratorContext) : Generator {
 
     private fun createEmptyIrFile(ktFile: KtFile): IrFileImpl {
         val fileEntry = context.sourceManager.getOrCreateFileEntry(ktFile)
-        val packageFragmentDescriptor = getOrFail(BindingContext.FILE_TO_PACKAGE_FRAGMENT, ktFile)
+        val packageFragmentDescriptor = context.moduleDescriptor.findPackageFragmentForFile(ktFile)!!
         val irFile = IrFileImpl(fileEntry, packageFragmentDescriptor)
         context.sourceManager.putFileEntry(irFile, fileEntry)
         return irFile

compiler/tests-common/tests/org/jetbrains/kotlin/resolve/annotation/AbstractAnnotationDescriptorResolveTest.java

@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.renderer.DescriptorRenderer;
 import org.jetbrains.kotlin.renderer.DescriptorRendererModifier;
 import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
+import org.jetbrains.kotlin.resolve.lazy.descriptors.PackageDescriptorUtilKt;
 import org.jetbrains.kotlin.resolve.scopes.MemberScope;
 import org.jetbrains.kotlin.test.KotlinTestUtils;
 import org.jetbrains.kotlin.test.KotlinTestWithEnvironment;
@@ -66,7 +67,11 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
 
     private static final FqName PACKAGE = new FqName("test");
 
-    protected BindingContext context;
+    protected AnalysisResult analysisResult;
+
+    protected BindingContext getContext() {
+        return analysisResult.getBindingContext();
+    }
 
     @Override
     protected KotlinCoreEnvironment createEnvironment() {
@@ -122,7 +127,7 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
 
     private void checkAnnotationsOnFile(String expectedAnnotation, KtFile file) {
         String actualAnnotation = StringUtil.join(file.getAnnotationEntries(), annotationEntry -> {
-            AnnotationDescriptor annotationDescriptor = context.get(BindingContext.ANNOTATION, annotationEntry);
+            AnnotationDescriptor annotationDescriptor = getContext().get(BindingContext.ANNOTATION, annotationEntry);
             assertNotNull(annotationDescriptor);
 
             KtAnnotationUseSiteTarget target = annotationEntry.getUseSiteTarget();
@@ -144,7 +149,7 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
         checkDescriptor(expectedAnnotation, getLocalClassDescriptor("LocalClass"));
         checkDescriptor(expectedAnnotation, getLocalObjectDescriptor("LocalObject"));
         checkDescriptor(expectedAnnotation, getLocalFunDescriptor("localFun"));
-        checkDescriptor(expectedAnnotation, getLocalVarDescriptor(context, "localVar"));
+        checkDescriptor(expectedAnnotation, getLocalVarDescriptor(getContext(), "localVar"));
     }
 
     private static void checkAnnotationsOnProperty(String expectedAnnotation, PropertyDescriptor prop) {
@@ -236,7 +241,7 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
 
     @NotNull
     private ClassDescriptor getLocalClassDescriptor(@NotNull String name) {
-        for (ClassDescriptor descriptor : context.getSliceContents(BindingContext.CLASS).values()) {
+        for (ClassDescriptor descriptor : getContext().getSliceContents(BindingContext.CLASS).values()) {
             if (descriptor.getName().asString().equals(name)) {
                 return descriptor;
             }
@@ -259,7 +264,7 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
 
     @NotNull
     private SimpleFunctionDescriptor getLocalFunDescriptor(@NotNull String name) {
-        for (SimpleFunctionDescriptor descriptor : context.getSliceContents(BindingContext.FUNCTION).values()) {
+        for (SimpleFunctionDescriptor descriptor : getContext().getSliceContents(BindingContext.FUNCTION).values()) {
             if (descriptor.getName().asString().equals(name)) {
                 return descriptor;
             }
@@ -283,7 +288,7 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
 
     @NotNull
     private SimpleFunctionDescriptor getAnonymousFunDescriptor() {
-        for (SimpleFunctionDescriptor descriptor : context.getSliceContents(BindingContext.FUNCTION).values()) {
+        for (SimpleFunctionDescriptor descriptor : getContext().getSliceContents(BindingContext.FUNCTION).values()) {
             if (descriptor instanceof AnonymousFunctionDescriptor) {
                 return descriptor;
             }
@@ -333,15 +338,15 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
     @NotNull
     protected KtFile getFile(@NotNull String content) {
         KtFile ktFile = KotlinTestUtils.createFile("dummy.kt", content, getProject());
-        AnalysisResult analysisResult = KotlinTestUtils.analyzeFile(ktFile, getEnvironment());
+        analysisResult = KotlinTestUtils.analyzeFile(ktFile, getEnvironment());
-        context = analysisResult.getBindingContext();
 
         return ktFile;
     }
 
     @NotNull
     protected PackageFragmentDescriptor getPackage(@NotNull KtFile ktFile) {
-        PackageFragmentDescriptor packageFragment = context.get(BindingContext.FILE_TO_PACKAGE_FRAGMENT, ktFile);
+        PackageFragmentDescriptor packageFragment =
+                PackageDescriptorUtilKt.findPackageFragmentForFile(analysisResult.getModuleDescriptor(), ktFile);
         assertNotNull("Failed to find package: " + PACKAGE, packageFragment);
         return packageFragment;
     }
@@ -384,7 +389,7 @@ public abstract class AbstractAnnotationDescriptorResolveTest extends KotlinTest
 
     @Override
     public void tearDown() throws Exception {
-        context = null;
+        analysisResult = null;
         super.tearDown();
     }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/stubindex/resolve/StubBasedPackageMemberDeclarationProvider.kt

@@ -109,6 +109,10 @@ class StubBasedPackageMemberDeclarationProvider(
         return PackageIndexUtil.findFilesWithExactPackage(fqName, searchScope, project)
     }
 
+    override fun containsFile(file: KtFile): Boolean {
+        return searchScope.contains(file.virtualFile ?: return false)
+    }
+
     override fun getTypeAliasDeclarations(name: Name): Collection<KtTypeAlias> {
         return KotlinTopLevelTypeAliasFqNameIndex.getInstance().get(childName(name), project, searchScope)
     }

idea/src/org/jetbrains/kotlin/idea/quickfix/MakeVisibleFactory.kt

@@ -23,23 +23,19 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithVisibility
 import org.jetbrains.kotlin.descriptors.Visibilities.*
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory3
-import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.caches.resolve.findModuleDescriptor
 import org.jetbrains.kotlin.idea.search.usagesSearch.descriptor
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtModifierListOwner
 import org.jetbrains.kotlin.psi.psiUtil.containingClass
-import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.getAllSuperclassesWithoutAny
-import org.jetbrains.kotlin.resolve.descriptorUtil.module
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
 object MakeVisibleFactory : KotlinIntentionActionsFactory() {
     override fun doCreateActions(diagnostic: Diagnostic): List<IntentionAction> {
         val element = diagnostic.psiElement as? KtElement ?: return emptyList()
-        val context = element.analyze(BodyResolveMode.PARTIAL)
+        val usageModule = element.findModuleDescriptor()
-        val usageModule = context.get(BindingContext.FILE_TO_PACKAGE_FRAGMENT, element.containingKtFile)?.module
 
         @Suppress("UNCHECKED_CAST")
         val factory = diagnostic.factory as DiagnosticFactory3<*, DeclarationDescriptor, *, DeclarationDescriptor>

idea/src/org/jetbrains/kotlin/idea/refactoring/move/moveDeclarations/moveConflictUtils.kt

@@ -65,6 +65,7 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.getImportableDescriptor
 import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatform
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.resolve.lazy.descriptors.findPackageFragmentForFile
 import org.jetbrains.kotlin.resolve.source.KotlinSourceElement
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.utils.SmartSet
@@ -103,7 +104,11 @@ class MoveConflictChecker(
                     is KtNamedDeclaration -> resolutionFacade.resolveToDescriptor(targetElement)
 
                     is KtFile -> {
-                        val packageFragment = resolutionFacade.analyze(targetElement)[BindingContext.FILE_TO_PACKAGE_FRAGMENT, targetElement]
+                        val packageFragment =
+                                targetElement
+                                        .findModuleDescriptor()
+                                        .findPackageFragmentForFile(targetElement)
+
                         packageFragment?.withSource(targetElement)
                     }