Optimize declaration providers for case of non-existing packages

Avoid creating memoized function nodes when the value is obviously null
Otherwise, it may lead to about 25M retained by empty concurrent hashmap
nodes

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/ResolveSession.java

@@ -63,7 +63,7 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
     private final BindingTrace trace;
     private final DeclarationProviderFactory declarationProviderFactory;
 
-    private final MemoizedFunctionToNullable<FqName, LazyPackageDescriptor> packages;
+    private final CacheWithNotNullValues<FqName, LazyPackageDescriptor> packages;
     private final PackageFragmentProvider packageFragmentProvider;
 
     private final MemoizedFunctionToNotNull<KtFile, LazyAnnotations> fileAnnotations;
@@ -169,7 +169,7 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
         this.trace = lockBasedLazyResolveStorageManager.createSafeTrace(delegationTrace);
         this.module = rootDescriptor;
 
-        this.packages = storageManager.createMemoizedFunctionWithNullableValues(this::createPackage);
+        this.packages = storageManager.createCacheWithNotNullValues();
 
         this.declarationProviderFactory = declarationProviderFactory;
 
@@ -216,7 +216,15 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
     @Override
     @Nullable
     public LazyPackageDescriptor getPackageFragment(@NotNull FqName fqName) {
-        return packages.invoke(fqName);
+        PackageMemberDeclarationProvider provider = declarationProviderFactory.getPackageMemberDeclarationProvider(fqName);
+        if (provider == null) {
+            return null;
+        }
+
+        return packages.computeIfAbsent(
+                fqName,
+                () -> new LazyPackageDescriptor(module, fqName, this, provider)
+        );
     }
 
 
@@ -231,15 +239,6 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
         return packageDescriptor;
     }
 
-    @Nullable
-    private LazyPackageDescriptor createPackage(FqName fqName) {
-        PackageMemberDeclarationProvider provider = declarationProviderFactory.getPackageMemberDeclarationProvider(fqName);
-        if (provider == null) {
-            return null;
-        }
-        return new LazyPackageDescriptor(module, fqName, this, provider);
-    }
-
     @NotNull
     @Override
     public ModuleDescriptor getModuleDescriptor() {

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/declarations/AbstractDeclarationProviderFactory.java

@@ -34,8 +34,11 @@ public abstract class AbstractDeclarationProviderFactory implements DeclarationP
     @Nullable
     protected abstract PackageMemberDeclarationProvider createPackageMemberDeclarationProvider(@NotNull FqName name);
 
+    public abstract boolean packageExists(@NotNull FqName fqName);
+
     @Override
     public PackageMemberDeclarationProvider getPackageMemberDeclarationProvider(@NotNull FqName packageFqName) {
+        if (!packageExists(packageFqName)) return null;
         return packageDeclarationProviders.invoke(packageFqName);
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/declarations/FileBasedDeclarationProviderFactory.java

@@ -65,7 +65,8 @@ public class FileBasedDeclarationProviderFactory extends AbstractDeclarationProv
         }
     }
 
-    private boolean isPackageDeclaredExplicitly(@NotNull FqName packageFqName) {
+    @Override
+    public boolean packageExists(@NotNull FqName packageFqName) {
         return index.invoke().declaredPackages.contains(packageFqName);
     }
 
@@ -76,7 +77,7 @@ public class FileBasedDeclarationProviderFactory extends AbstractDeclarationProv
     @Nullable
     @Override
     protected PackageMemberDeclarationProvider createPackageMemberDeclarationProvider(@NotNull FqName packageFqName) {
-        if (isPackageDeclaredExplicitly(packageFqName)) {
+        if (packageExists(packageFqName)) {
             return new FileBasedPackageMemberDeclarationProvider(
                     storageManager, packageFqName, this, index.invoke().filesByPackage.get(packageFqName));
         }

core/descriptors.jvm/src/org/jetbrains/kotlin/load/java/lazy/LazyJavaPackageFragmentProvider.kt

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.descriptors.PackageFragmentProvider
 import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaPackageFragment
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.storage.MemoizedFunctionToNullable
+import org.jetbrains.kotlin.storage.CacheWithNotNullValues
 
 class LazyJavaPackageFragmentProvider(
         components: JavaResolverComponents
@@ -28,17 +28,16 @@ class LazyJavaPackageFragmentProvider(
 
     private val c = LazyJavaResolverContext(components, TypeParameterResolver.EMPTY, lazyOf(null))
 
-    private val packageFragments: MemoizedFunctionToNullable<FqName, LazyJavaPackageFragment> =
-            c.storageManager.createMemoizedFunctionWithNullableValues {
-                fqName ->
-                val jPackage = c.components.finder.findPackage(fqName)
-                if (jPackage != null) {
-                    LazyJavaPackageFragment(c, jPackage)
-                }
-                else null
-            }
+    private val packageFragments: CacheWithNotNullValues<FqName, LazyJavaPackageFragment> =
+        c.storageManager.createCacheWithNotNullValues()
 
-    private fun getPackageFragment(fqName: FqName) = packageFragments(fqName)
+    private fun getPackageFragment(fqName: FqName): LazyJavaPackageFragment? {
+        val jPackage = c.components.finder.findPackage(fqName) ?: return null
+
+        return packageFragments.computeIfAbsent(fqName) {
+            LazyJavaPackageFragment(c, jPackage)
+        }
+    }
 
     override fun getPackageFragments(fqName: FqName) = listOfNotNull(getPackageFragment(fqName))
 

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

@@ -57,7 +57,10 @@ class PluginDeclarationProviderFactory(
         }
     }
 
-    private fun packageExists(name: FqName): Boolean {
+    override fun packageExists(fqName: FqName) =
+        fileBasedDeclarationProviderFactory.packageExists(fqName) || stubBasedPackageExists(fqName)
+
+    private fun stubBasedPackageExists(name: FqName): Boolean {
         return if (moduleInfo is ModuleSourceInfo)
             project.service<PerModulePackageCacheService>().packageExists(name, moduleInfo)
         else
@@ -65,7 +68,7 @@ class PluginDeclarationProviderFactory(
     }
 
     private fun getStubBasedPackageMemberDeclarationProvider(name: FqName): PackageMemberDeclarationProvider? {
-        if (!packageExists(name)) return null
+        if (!stubBasedPackageExists(name)) return null
 
         return StubBasedPackageMemberDeclarationProvider(name, project, indexedFilesScope)
     }