[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

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) }
 }

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))
     }
 }