From b75024e3e003b4a77d843b8205cab97c56a8ed98 Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Thu, 26 Apr 2018 13:04:24 +0300
Subject: [PATCH] Fix memory leak related to diagnostics computation within
 light classes

Before this change, when somebody has requested the exact light class
we also store LazyJvmDiagnostic for it in a map inside DiagnosticHolder.

The problem is that LazyJvmDiagnostic retains a lot of memory:
GenerationState, bindingContext, etc.

At the same time, it commonly happens that light classes are being computed from Java resolve and diagnostics are obviously
unnecessary there.

The solution has two parts:
1. We don't retain diagnostics until somebody explicitly requested them
2. We force computation of LazyJvmDiagnostic to avoid retaining
backend's parts

The solution has a kind of drawback: for an opened in editor Kotlin file,
we'll run back-end twice for contained classes (analysis parts should be reused, though)

 #KT-24048 Fixed
---
 .../lightClasses/LazyLightClassDataHolder.kt  | 34 +++++++------------
 1 file changed, 12 insertions(+), 22 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/lightClasses/LazyLightClassDataHolder.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/lightClasses/LazyLightClassDataHolder.kt
index 2a894ee9970..12fd6ac64ab 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/lightClasses/LazyLightClassDataHolder.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/lightClasses/LazyLightClassDataHolder.kt
@@ -30,30 +30,20 @@ typealias DummyLightClassContextProvider = (() -> LightClassConstructionContext?
 typealias DiagnosticsHolderProvider = () -> LazyLightClassDataHolder.DiagnosticsHolder
 
 sealed class LazyLightClassDataHolder(
-    builder: LightClassBuilder,
-    exactContextProvider: ExactLightClassContextProvider,
+    private val builder: LightClassBuilder,
+    private val exactContextProvider: ExactLightClassContextProvider,
     dummyContextProvider: DummyLightClassContextProvider,
     private val diagnosticsHolderProvider: DiagnosticsHolderProvider
 ) : LightClassDataHolder {
 
-    class DiagnosticsHolder(private val storageManager: StorageManager) {
-        private val computedLightClassDiagnostics = hashMapOf<LazyLightClassDataHolder, Diagnostics>()
+    class DiagnosticsHolder(storageManager: StorageManager) {
+        private val cache = storageManager.createCacheWithNotNullValues<LazyLightClassDataHolder, Diagnostics>()
 
-        fun putDiagnostics(lazyLightClassDataHolder: LazyLightClassDataHolder, diagnostics: Diagnostics) {
-            storageManager.compute {
-                computedLightClassDiagnostics[lazyLightClassDataHolder] = diagnostics
-            }
-        }
-
-        fun getComputedDiagnostics(lazyLightClassDataHolder: LazyLightClassDataHolder): Diagnostics? =
-            storageManager.compute { computedLightClassDiagnostics[lazyLightClassDataHolder] }
+        fun getOrCompute(lazyLightClassDataHolder: LazyLightClassDataHolder, diagnostics: () -> Diagnostics) =
+            cache.computeIfAbsent(lazyLightClassDataHolder, diagnostics)
     }
 
-    private val exactResultLazyValue = lazyPub {
-        val (stub, _, diagnostics) = builder(exactContextProvider())
-        diagnosticsHolderProvider().putDiagnostics(this, diagnostics)
-        stub
-    }
+    private val exactResultLazyValue = lazyPub { builder(exactContextProvider()).stub }
 
     private val lazyInexactStub by lazyPub {
         dummyContextProvider?.let { provider -> provider()?.let { context -> builder.invoke(context).stub } }
@@ -65,11 +55,11 @@ sealed class LazyLightClassDataHolder(
     override val javaFileStub by exactResultLazyValue
 
     override val extraDiagnostics: Diagnostics
-        get() {
-            // run light class builder
-            exactResultLazyValue.value
-
-            return diagnosticsHolderProvider().getComputedDiagnostics(this) ?: Diagnostics.EMPTY
+        get() = diagnosticsHolderProvider().getOrCompute(this) {
+            builder(exactContextProvider()).diagnostics
+                // Force lazy diagnostics computation because otherwise a lot of memory is retained by computation.
+                // NB: Laziness here is not crucial anyway since somebody already has requested diagnostics and we hope one will use them
+                .takeUnless { it.isEmpty() } ?: Diagnostics.EMPTY
         }
 
     // for facade or defaultImpls