From 9cf77e19be836e729921e4247ad7023fc04be92c Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 12 Dec 2022 19:00:03 +0100
Subject: [PATCH] [LL FIR] fix memory leak from LLFirSessionProvider for the
 module

LLFirSessionProvider stores all transitive module dependencies,
even if they are not needed or already invalidated and garbage collected

Now, the dependencies are stored on the weak references, and so then nobody uses them, they will be garbage collected
---
 .../api/fir/sessions/LLFirSessionProvider.kt  | 45 ++++++++++++++++---
 .../sessions/LLFirSessionProviderStorage.kt   |  6 +--
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProvider.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProvider.kt
index d5ceb1488d2..fc0e3a3a961 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProvider.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProvider.kt
@@ -6,20 +6,20 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.Immutable
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirModuleData
-import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirDependentModuleProviders
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.FirSessionProvider
-import org.jetbrains.kotlin.fir.resolve.providers.dependenciesSymbolProvider
+import java.lang.ref.WeakReference
 
 @Immutable
 class LLFirSessionProvider internal constructor(
     val project: Project,
     internal val rootModuleSession: LLFirResolvableModuleSession,
-    private val ktModuleToSession: Map<KtModule, LLFirSession>
+    private val ktModuleToSession: KtModuleToSessionMapping
 ) : FirSessionProvider() {
     override fun getSession(moduleData: FirModuleData): LLFirSession {
         requireIsInstance<LLFirModuleData>(moduleData)
@@ -27,11 +27,44 @@ class LLFirSessionProvider internal constructor(
     }
 
     fun getSession(module: KtModule): LLFirSession =
-        ktModuleToSession.getValue(module)
+        ktModuleToSession.getSession(module)
 
     fun getResolvableSession(module: KtModule): LLFirResolvableModuleSession =
-        ktModuleToSession.getValue(module) as LLFirResolvableModuleSession
+        ktModuleToSession.getSession(module) as LLFirResolvableModuleSession
 
+    @get:TestOnly
     val allSessions: Collection<LLFirSession>
-        get() = ktModuleToSession.values
+        get() = ktModuleToSession.getAllSessions()
 }
+
+internal abstract class KtModuleToSessionMapping {
+    abstract fun getSession(module: KtModule): LLFirSession
+
+    @TestOnly
+    abstract fun getAllSessions(): Collection<LLFirSession>
+}
+
+internal class KtModuleToSessionMappingByMapImpl(
+    private val map: Map<KtModule, LLFirSession>
+) : KtModuleToSessionMapping() {
+    override fun getSession(module: KtModule): LLFirSession =
+        map.getValue(module)
+
+    override fun getAllSessions(): Collection<LLFirSession> =
+        map.values
+}
+
+internal class KtModuleToSessionMappingByWeakValueMapImpl(
+    initialMap: Map<KtModule, LLFirSession>
+) : KtModuleToSessionMapping() {
+    private val softValuesMap = initialMap.entries.associate { (key, value) -> key to WeakReference(value) }
+
+    override fun getSession(module: KtModule): LLFirSession {
+        val softReference = softValuesMap.getValue(module)
+        return softReference.get()
+            ?: error("soft reference for $module was invalidated")
+    }
+
+    override fun getAllSessions(): Collection<LLFirSession> =
+        softValuesMap.keys.map { getSession(it) }
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProviderStorage.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProviderStorage.kt
index f97db309bf6..3cc6465206e 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProviderStorage.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProviderStorage.kt
@@ -48,7 +48,7 @@ class LLFirSessionProviderStorage(val project: Project) {
             is KtNotUnderContentRootModule -> {
                 val session = LLFirNonUnderContentRootSessionFactory.getInstance(project)
                     .getNonUnderContentRootSession(useSiteKtModule)
-                LLFirSessionProvider(project, session, mapOf(useSiteKtModule to session))
+                LLFirSessionProvider(project, session, KtModuleToSessionMappingByMapImpl(mapOf(useSiteKtModule to session)))
             }
 
             else -> error("Unexpected ${useSiteKtModule::class.simpleName}")
@@ -72,7 +72,7 @@ class LLFirSessionProviderStorage(val project: Project) {
             )
             sessions to session
         }
-        return LLFirSessionProvider(project, session, sessions)
+        return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
     }
 
 
@@ -93,7 +93,7 @@ class LLFirSessionProviderStorage(val project: Project) {
             )
             sessions to session
         }
-        return LLFirSessionProvider(project, session, sessions)
+        return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
     }
 }