Fix a memory leak of old resolve session from light classes

The problem is that diagnostics might indirectly retain a reference to an
obsolete ResolverForProject (it becomes outdated after out of block
modification).

At the same time light class instances may survive afer OOM when they
are leaked somewhere (e.g. to Java resolve)

The idea is to store the Diagnostics instance inside a
ResolverForProject component which will be collected once
ResolverForProject is collected itself.

 #KT-23380 In Progress

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/lightClasses/LazyLightClassDataHolder.kt

@@ -5,15 +5,11 @@
 
 package org.jetbrains.kotlin.idea.caches.lightClasses
 
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiMember
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.impl.java.stubs.PsiJavaFileStub
-import com.intellij.psi.util.CachedValueProvider
-import com.intellij.psi.util.CachedValuesManager
-import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.kotlin.asJava.LightClassBuilder
 import org.jetbrains.kotlin.asJava.builder.*
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
@@ -26,46 +22,55 @@ import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.NotNullableUserDataProperty
 import org.jetbrains.kotlin.psi.debugText.getDebugText
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
+import org.jetbrains.kotlin.storage.StorageManager
 import org.jetbrains.kotlin.utils.KotlinExceptionWithAttachments
 
 typealias ExactLightClassContextProvider = () -> LightClassConstructionContext
 typealias DummyLightClassContextProvider = (() -> LightClassConstructionContext?)?
+typealias DiagnosticsHolderProvider = () -> LazyLightClassDataHolder.DiagnosticsHolder
 
 sealed class LazyLightClassDataHolder(
     builder: LightClassBuilder,
-    project: Project,
     exactContextProvider: ExactLightClassContextProvider,
     dummyContextProvider: DummyLightClassContextProvider,
-    isLocal: Boolean = false
+    private val diagnosticsHolderProvider: DiagnosticsHolderProvider
 ) : LightClassDataHolder {
 
-    private data class CachedLightClassBuilderResult(val stub: PsiJavaFileStub, val diagnostics: Diagnostics)
+    class DiagnosticsHolder(private val storageManager: StorageManager) {
+        private val computedLightClassDiagnostics = hashMapOf<LazyLightClassDataHolder, Diagnostics>()
 
-    private val exactResultCachedValue =
-        CachedValuesManager.getManager(project).createCachedValue(
-            {
-                val (stub, _, diagnostics) = builder(exactContextProvider())
-                val cachedResult = CachedLightClassBuilderResult(stub, diagnostics)
+        fun putDiagnostics(lazyLightClassDataHolder: LazyLightClassDataHolder, diagnostics: Diagnostics) {
+            storageManager.compute {
+                computedLightClassDiagnostics[lazyLightClassDataHolder] = diagnostics
+            }
+        }
 
-                CachedValueProvider.Result.create(
-                    cachedResult,
-                    if (isLocal)
-                        PsiModificationTracker.MODIFICATION_COUNT
-                    else
-                        PsiModificationTracker.OUT_OF_CODE_BLOCK_MODIFICATION_COUNT
-                )
-            }, false
-        )
+        fun getComputedDiagnostics(lazyLightClassDataHolder: LazyLightClassDataHolder): Diagnostics? =
+            storageManager.compute { computedLightClassDiagnostics[lazyLightClassDataHolder] }
+    }
+
+    private val exactResultLazyValue = lazyPub {
+        val (stub, _, diagnostics) = builder(exactContextProvider())
+        diagnosticsHolderProvider().putDiagnostics(this, diagnostics)
+        stub
+    }
 
     private val lazyInexactStub by lazyPub {
         dummyContextProvider?.let { provider -> provider()?.let { context -> builder.invoke(context).stub } }
     }
 
     private val inexactStub: PsiJavaFileStub?
-        get() = if (exactResultCachedValue.hasUpToDateValue()) null else lazyInexactStub
+        get() = if (exactResultLazyValue.isInitialized()) null else lazyInexactStub
 
-    override val javaFileStub get() = exactResultCachedValue.value.stub
-    override val extraDiagnostics get() = exactResultCachedValue.value.diagnostics
+    override val javaFileStub by exactResultLazyValue
+
+    override val extraDiagnostics: Diagnostics
+        get() {
+            // run light class builder
+            exactResultLazyValue.value
+
+            return diagnosticsHolderProvider().getComputedDiagnostics(this) ?: Diagnostics.EMPTY
+        }
 
     // for facade or defaultImpls
     override fun findData(findDelegate: (PsiJavaFileStub) -> PsiClass): LightClassData =
@@ -74,9 +79,12 @@ sealed class LazyLightClassDataHolder(
         }
 
     class ForClass(
-        builder: LightClassBuilder, project: Project, isLocal: Boolean,
-        exactContextProvider: ExactLightClassContextProvider, dummyContextProvider: DummyLightClassContextProvider
-    ) : LazyLightClassDataHolder(builder, project, exactContextProvider, dummyContextProvider, isLocal), LightClassDataHolder.ForClass {
+        builder: LightClassBuilder,
+        exactContextProvider: ExactLightClassContextProvider,
+        dummyContextProvider: DummyLightClassContextProvider,
+        diagnosticsHolderProvider: DiagnosticsHolderProvider
+    ) : LazyLightClassDataHolder(builder, exactContextProvider, dummyContextProvider, diagnosticsHolderProvider),
+        LightClassDataHolder.ForClass {
         override fun findDataForClassOrObject(classOrObject: KtClassOrObject): LightClassData =
             LazyLightClassData { stub ->
                 stub.findDelegate(classOrObject)
@@ -84,14 +92,20 @@ sealed class LazyLightClassDataHolder(
     }
 
     class ForFacade(
-        builder: LightClassBuilder, project: Project,
-        exactContextProvider: ExactLightClassContextProvider, dummyContextProvider: DummyLightClassContextProvider
-    ) : LazyLightClassDataHolder(builder, project, exactContextProvider, dummyContextProvider), LightClassDataHolder.ForFacade
+        builder: LightClassBuilder,
+        exactContextProvider: ExactLightClassContextProvider,
+        dummyContextProvider: DummyLightClassContextProvider,
+        diagnosticsHolderProvider: DiagnosticsHolderProvider
+    ) : LazyLightClassDataHolder(builder, exactContextProvider, dummyContextProvider, diagnosticsHolderProvider),
+        LightClassDataHolder.ForFacade
 
     class ForScript(
-        builder: LightClassBuilder, project: Project,
-        exactContextProvider: ExactLightClassContextProvider, dummyContextProvider: DummyLightClassContextProvider
-    ) : LazyLightClassDataHolder(builder, project, exactContextProvider, dummyContextProvider), LightClassDataHolder.ForScript
+        builder: LightClassBuilder,
+        exactContextProvider: ExactLightClassContextProvider,
+        dummyContextProvider: DummyLightClassContextProvider,
+        diagnosticsHolderProvider: DiagnosticsHolderProvider
+    ) : LazyLightClassDataHolder(builder, exactContextProvider, dummyContextProvider, diagnosticsHolderProvider),
+        LightClassDataHolder.ForScript
 
     private inner class LazyLightClassData(
         findDelegate: (PsiJavaFileStub) -> PsiClass

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/IDELightClassGenerationSupport.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.idea.caches.lightClasses.IDELightClassContexts
 import org.jetbrains.kotlin.idea.caches.lightClasses.LazyLightClassDataHolder
+import org.jetbrains.kotlin.idea.resolve.frontendService
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.lazy.NoDescriptorForDeclarationException
@@ -39,17 +40,15 @@ class IDELightClassGenerationSupport(project: Project) : LightClassGenerationSup
             classOrObject.shouldNotBeVisibleAsLightClass() -> InvalidLightClassDataHolder
             classOrObject.isLocal -> LazyLightClassDataHolder.ForClass(
                 builder,
-                classOrObject.project,
                 exactContextProvider = { IDELightClassContexts.contextForLocalClassOrObject(classOrObject) },
                 dummyContextProvider = null,
-                isLocal = true
+                diagnosticsHolderProvider = { classOrObject.getDiagnosticsHolder() }
             )
             else -> LazyLightClassDataHolder.ForClass(
                 builder,
-                classOrObject.project,
                 exactContextProvider = { IDELightClassContexts.contextForNonLocalClassOrObject(classOrObject) },
                 dummyContextProvider = { IDELightClassContexts.lightContextForClassOrObject(classOrObject) },
-                isLocal = false
+                diagnosticsHolderProvider = { classOrObject.getDiagnosticsHolder() }
             )
         }
     }
@@ -62,21 +61,23 @@ class IDELightClassGenerationSupport(project: Project) : LightClassGenerationSup
 
         return LazyLightClassDataHolder.ForFacade(
             builder,
-            files.first().project,
             exactContextProvider = { IDELightClassContexts.contextForFacade(sortedFiles) },
-            dummyContextProvider = { IDELightClassContexts.lightContextForFacade(sortedFiles) }
+            dummyContextProvider = { IDELightClassContexts.lightContextForFacade(sortedFiles) },
+            diagnosticsHolderProvider = { files.first().getDiagnosticsHolder() }
         )
     }
 
     override fun createDataHolderForScript(script: KtScript, builder: LightClassBuilder): LightClassDataHolder.ForScript {
         return LazyLightClassDataHolder.ForScript(
             builder,
-            script.project,
             exactContextProvider = { IDELightClassContexts.contextForScript(script) },
-            dummyContextProvider = { IDELightClassContexts.lightContextForScript(script) }
+            dummyContextProvider = { IDELightClassContexts.lightContextForScript(script) },
+            diagnosticsHolderProvider = { script.getDiagnosticsHolder() }
         )
     }
 
+    private fun KtElement.getDiagnosticsHolder() =
+        getResolutionFacade().frontendService<LazyLightClassDataHolder.DiagnosticsHolder>()
 
     override fun resolveToDescriptor(declaration: KtDeclaration): DeclarationDescriptor? {
         try {

idea/idea-analysis/src/org/jetbrains/kotlin/idea/project/IdeaEnvironment.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.idea.project
 
 import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.container.useImpl
+import org.jetbrains.kotlin.idea.caches.lightClasses.LazyLightClassDataHolder
 import org.jetbrains.kotlin.resolve.TargetEnvironment
 
 object IdeaEnvironment : TargetEnvironment("Idea") {
@@ -25,5 +26,6 @@ object IdeaEnvironment : TargetEnvironment("Idea") {
         container.useImpl<ResolveElementCache>()
         container.useImpl<IdeaLocalDescriptorResolver>()
         container.useImpl<IdeaAbsentDescriptorHandler>()
+        container.useImpl<LazyLightClassDataHolder.DiagnosticsHolder>()
     }
-}
+}