[Analysis API FIR] fix invalid usages of CachedValue

it should not cache map as values are stored on soft references

analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/CachingKtAnalysisSessionProvider.kt

@@ -7,21 +7,18 @@ package org.jetbrains.kotlin.analysis.api.impl.base
 
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootModificationTracker
-import com.intellij.psi.util.CachedValueProvider
-import com.intellij.psi.util.CachedValuesManager
 import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
-import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
-import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
-import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenFactory
+import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
+import org.jetbrains.kotlin.analysis.utils.caches.SoftCachedMap
 import org.jetbrains.kotlin.psi.KtElement
-import java.util.concurrent.ConcurrentHashMap
 import kotlin.reflect.KClass
 
 @KtAnalysisApiInternals
@@ -47,10 +44,6 @@ abstract class CachingKtAnalysisSessionProvider<State : Any>(project: Project) :
         }
     }
 
-    private fun getCachedAnalysisSession(firResolveSession: State, token: KtLifetimeToken): KtAnalysisSession? {
-        return cache.getCachedAnalysisSession(firResolveSession to token::class)
-    }
-
     @TestOnly
     final override fun clearCaches() {
         cache.clear()
@@ -58,23 +51,24 @@ abstract class CachingKtAnalysisSessionProvider<State : Any>(project: Project) :
 }
 
 private class KtAnalysisSessionCache(project: Project) {
-    private val cache = CachedValuesManager.getManager(project).createCachedValue {
-        CachedValueProvider.Result(
-            ConcurrentHashMap<Pair<KtModule, KClass<out KtLifetimeToken>>, KtAnalysisSession>(),
+    private val cache = SoftCachedMap.create<Pair<KtModule, KClass<out KtLifetimeToken>>, KtAnalysisSession>(
+        project,
+        SoftCachedMap.Kind.STRONG_KEYS_SOFT_VALUES,
+        listOf(
             PsiModificationTracker.MODIFICATION_COUNT,
             ProjectRootModificationTracker.getInstance(project),
             project.createProjectWideOutOfBlockModificationTracker()
         )
-    }
+    )
 
     @TestOnly
     fun clear() {
-        cache.value.clear()
+        cache.clear()
     }
 
-    inline fun getAnalysisSession(key: Pair<KtModule, KClass<out KtLifetimeToken>>, create: () -> KtAnalysisSession): KtAnalysisSession =
-        cache.value.getOrPut(key) { create() }
-
-    fun getCachedAnalysisSession(key: KEY): KtAnalysisSession? =
-        cache.value[key]
+    fun getAnalysisSession(
+        key: Pair<KtModule, KClass<out KtLifetimeToken>>,
+        create: () -> KtAnalysisSession
+    ): KtAnalysisSession =
+        cache.getOrPut(key) { create() }
 }

analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/caches/SoftCachedMap.kt

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.utils.caches
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.util.CachedValue
+import com.intellij.psi.util.CachedValueProvider
+import com.intellij.psi.util.CachedValuesManager
+import com.intellij.util.containers.ContainerUtil
+import org.jetbrains.annotations.TestOnly
+import java.util.concurrent.ConcurrentHashMap
+
+public abstract class SoftCachedMap<K : Any, V : Any> {
+    public abstract fun getOrPut(key: K, create: () -> V): V
+
+    @TestOnly
+    public abstract fun clear()
+
+    public companion object {
+        public fun <K : Any, V : Any> create(
+            project: Project,
+            kind: Kind,
+            trackers: List<Any>
+        ): SoftCachedMap<K, V> = when {
+            trackers.isEmpty() -> SoftCachedMapWithoutTrackers(kind)
+            else -> SoftCachedMapWithTrackers(project, kind, trackers.toTypedArray())
+        }
+    }
+
+    public enum class Kind {
+        SOFT_KEYS_SOFT_VALUES,
+        STRONG_KEYS_SOFT_VALUES
+    }
+}
+
+private class SoftCachedMapWithTrackers<K : Any, V : Any>(
+    private val project: Project,
+    kind: Kind,
+    private val trackers: Array<Any>
+) : SoftCachedMap<K, V>() {
+    private val cache = when (kind) {
+        Kind.SOFT_KEYS_SOFT_VALUES -> ContainerUtil.createConcurrentSoftMap<K, CachedValue<V>>()
+        Kind.STRONG_KEYS_SOFT_VALUES -> ConcurrentHashMap<K, CachedValue<V>>()
+    }
+
+    override fun clear() {
+        cache.clear()
+    }
+
+    override fun getOrPut(key: K, create: () -> V): V {
+        return cache.getOrPut(key) {
+            CachedValuesManager.getManager(project).createCachedValue {
+                CachedValueProvider.Result(create(), *trackers)
+            }
+        }.value
+    }
+}
+
+private class SoftCachedMapWithoutTrackers<K : Any, V : Any>(kind: Kind) : SoftCachedMap<K, V>() {
+    private val cache = when (kind) {
+        Kind.SOFT_KEYS_SOFT_VALUES -> ContainerUtil.createConcurrentSoftKeySoftValueMap<K, V>()
+        Kind.STRONG_KEYS_SOFT_VALUES -> ContainerUtil.createSoftValueMap<K, V>()
+    }
+
+    override fun clear() {
+        cache.clear()
+    }
+
+    override fun getOrPut(key: K, create: () -> V): V {
+        return cache.getOrPut(key, create)
+    }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirResolveSessionDepended.kt

@@ -41,13 +41,14 @@ internal class LLFirResolveSessionDepended(
     override val useSiteKtModule: KtModule get() = originalFirResolveSession.useSiteKtModule
     override val useSiteFirSession get() = originalFirResolveSession.useSiteFirSession
 
-    private val scopeSessionProviderCache by softCachedValue(
+    private val scopeSessionProviderCache = SoftCachedMap.create<FirSession, LLFirScopeSessionProvider>(
         project,
-        PsiModificationTracker.MODIFICATION_COUNT,
-        ProjectRootModificationTracker.getInstance(project),
-    ) {
-        ConcurrentHashMap<FirSession, LLFirScopeSessionProvider>()
-    }
+        SoftCachedMap.Kind.SOFT_KEYS_SOFT_VALUES,
+        listOf(
+            PsiModificationTracker.MODIFICATION_COUNT,
+            ProjectRootModificationTracker.getInstance(project)
+        )
+    )
 
     override fun getScopeSessionFor(firSession: FirSession): ScopeSession {
         return scopeSessionProviderCache

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirResolveStateService.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootModificationTracker
+import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
@@ -19,33 +20,23 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveS
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
-import org.jetbrains.kotlin.analysis.utils.caches.getValue
-import org.jetbrains.kotlin.analysis.utils.caches.softCachedValue
-import java.util.concurrent.locks.ReentrantReadWriteLock
-import kotlin.concurrent.withLock
+import org.jetbrains.kotlin.analysis.utils.caches.SoftCachedMap
 
 internal class LLFirResolveSessionService(project: Project) {
     private val sessionProviderStorage = LLFirSessionProviderStorage(project)
 
-    private val stateCache by softCachedValue(
+    private val cache = SoftCachedMap.create<KtModule, LLFirResolvableResolveSession>(
         project,
-        project.createProjectWideOutOfBlockModificationTracker(),
-        ProjectRootModificationTracker.getInstance(project),
-    ) {
-        mutableMapOf<KtModule, LLFirResolvableResolveSession>()
-    }
-
-    private val cacheLock = ReentrantReadWriteLock()
+        SoftCachedMap.Kind.STRONG_KEYS_SOFT_VALUES,
+        listOf(
+            ProjectRootModificationTracker.getInstance(project),
+            project.createProjectWideOutOfBlockModificationTracker(),
+        )
+    )
 
     fun getFirResolveSession(module: KtModule): LLFirResolvableResolveSession {
-        cacheLock.readLock().withLock {
-            stateCache[module]?.let { return it }
-        }
-        cacheLock.writeLock().withLock {
-            stateCache[module]?.let { return it }
-            val session = createFirResolveSessionFor(module, sessionProviderStorage)
-            stateCache[module] = session
-            return session
+        return cache.getOrPut(module) {
+            createFirResolveSessionFor(module, sessionProviderStorage)
         }
     }
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdePredicateBasedProvider.kt

@@ -9,6 +9,7 @@ import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSourcesSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
 import org.jetbrains.kotlin.analysis.providers.KotlinAnnotationsResolver
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
@@ -40,7 +41,7 @@ import org.jetbrains.kotlin.psi.*
  * PSI index based implementation of [FirPredicateBasedProvider].
  */
 internal class LLFirIdePredicateBasedProvider(
-    private val session: FirSession,
+    private val session: LLFirSourcesSession,
     private val annotationsResolver: KotlinAnnotationsResolver,
     private val declarationProvider: KotlinDeclarationProvider,
 ) : FirPredicateBasedProvider() {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/LLFirScopeSessionProvider.kt

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.roots.ProjectRootModificationTracker
-import com.intellij.psi.util.PsiModificationTracker
-import org.jetbrains.kotlin.analysis.utils.caches.getValue
-import org.jetbrains.kotlin.analysis.utils.caches.softCachedValue
+import org.jetbrains.kotlin.analysis.utils.caches.SoftCachedMap
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import java.util.concurrent.ConcurrentHashMap
 
@@ -28,12 +25,11 @@ abstract class LLFirScopeSessionProvider {
 private class LLFirInvalidatableScopeSessionProvider(project: Project, invalidationTrackers: List<Any>) : LLFirScopeSessionProvider() {
     // ScopeSession is thread-local, so we use Thread id as a key
     // We cannot use thread locals here as it may lead to memory leaks
-    private val cache by softCachedValue(
+    private val cache = SoftCachedMap.create<Long, ScopeSession>(
         project,
-        *invalidationTrackers.toTypedArray(),
-    ) {
-        ConcurrentHashMap<Long, ScopeSession>()
-    }
+        SoftCachedMap.Kind.STRONG_KEYS_SOFT_VALUES,
+        invalidationTrackers
+    )
 
     override fun getScopeSession(): ScopeSession {
         return cache.getOrPut(Thread.currentThread().id) { ScopeSession() }