[LL FIR] optimize FirCacheValue, do not create ConcurrentHashMap to store a single value

^KTIJ-24640
^KTIJ-24641

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.fir.caches
 import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
 import java.util.concurrent.ConcurrentHashMap
+import org.jetbrains.kotlin.fir.caches.FirLazyValue
 
 object FirThreadSafeCachesFactory : FirCachesFactory() {
     override fun <KEY : Any, VALUE, CONTEXT> createCache(createValue: (KEY, CONTEXT) -> VALUE): FirCache<KEY, VALUE, CONTEXT> =
@@ -29,4 +30,7 @@ object FirThreadSafeCachesFactory : FirCachesFactory() {
         postCompute: (KEY, VALUE, DATA) -> Unit
     ): FirCache<KEY, VALUE, CONTEXT> =
         FirThreadSafeCacheWithPostCompute(createValue, postCompute)
+
+    override fun <V> createLazyValue(createValue: () -> V): FirLazyValue<V> =
+        FirThreadSafeValue(createValue)
 }

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

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2021 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.low.level.api.fir.fir.caches
+
+import org.jetbrains.kotlin.fir.caches.FirLazyValue
+
+internal class FirThreadSafeValue<V>(createValue: () -> V) : FirLazyValue<V>() {
+    private val lazyValue by lazy(LazyThreadSafetyMode.SYNCHRONIZED, createValue)
+    override fun getValue(): V = lazyValue
+}

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

@@ -25,7 +25,7 @@ internal class LLFirIdeRegisteredPluginAnnotations(
     override val annotations: Set<AnnotationFqn>
         get() = allAnnotationsCache.getValue()
 
-    private val allAnnotationsCache: FirLazyValue<Set<AnnotationFqn>, Nothing?> = session.firCachesFactory.createLazyValue {
+    private val allAnnotationsCache: FirLazyValue<Set<AnnotationFqn>> = session.firCachesFactory.createLazyValue {
         // at this point, both metaAnnotations and annotationsFromPlugins should be collected
         val result = metaAnnotations.flatMapTo(mutableSetOf()) { getAnnotationsWithMetaAnnotation(it) }
 

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt

@@ -85,11 +85,11 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
         }
 
     @FirExtensionApiInternals
-    val topLevelClassIdsCache: FirLazyValue<Set<ClassId>, Nothing?> =
+    val topLevelClassIdsCache: FirLazyValue<Set<ClassId>> =
         session.firCachesFactory.createLazyValue { getTopLevelClassIds() }
 
     @FirExtensionApiInternals
-    val topLevelCallableIdsCache: FirLazyValue<Set<CallableId>, Nothing?> =
+    val topLevelCallableIdsCache: FirLazyValue<Set<CallableId>> =
         session.firCachesFactory.createLazyValue { getTopLevelCallableIds() }
 
 }

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirExtensionDeclarationsSymbolProvider.kt

@@ -53,7 +53,7 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
         hasPackage(packageFqName)
     }
 
-    private val callableNamesInPackageCache: FirLazyValue<Map<FqName, Set<Name>>, Nothing?> =
+    private val callableNamesInPackageCache: FirLazyValue<Map<FqName, Set<Name>>> =
         cachesFactory.createLazyValue {
             computeNamesGroupedByPackage(
                 FirDeclarationGenerationExtension::getTopLevelCallableIds,
@@ -61,7 +61,7 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
             )
         }
 
-    private val classNamesInPackageCache: FirLazyValue<Map<FqName, Set<String>>, Nothing?> =
+    private val classNamesInPackageCache: FirLazyValue<Map<FqName, Set<String>>> =
         cachesFactory.createLazyValue {
             computeNamesGroupedByPackage(
                 FirDeclarationGenerationExtension::getTopLevelClassIds,
@@ -82,12 +82,12 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
             }
         }
 
-    private val extensionsByTopLevelClassId: FirLazyValue<Map<ClassId, List<FirDeclarationGenerationExtension>>, Nothing?> =
+    private val extensionsByTopLevelClassId: FirLazyValue<Map<ClassId, List<FirDeclarationGenerationExtension>>> =
         session.firCachesFactory.createLazyValue {
             extensions.flatGroupBy { it.topLevelClassIdsCache.getValue() }
         }
 
-    private val extensionsByTopLevelCallableId: FirLazyValue<Map<CallableId, List<FirDeclarationGenerationExtension>>, Nothing?> =
+    private val extensionsByTopLevelCallableId: FirLazyValue<Map<CallableId, List<FirDeclarationGenerationExtension>>> =
         session.firCachesFactory.createLazyValue {
             extensions.flatGroupBy { it.topLevelCallableIdsCache.getValue() }
         }
@@ -171,7 +171,7 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
 
     override fun computePackageSetWithTopLevelCallables(): Set<String> =
         extensions.flatMapTo(mutableSetOf()) { extension ->
-            extension.topLevelCallableIdsCache.getValue(null).map { it.packageName.asString() }
+            extension.topLevelCallableIdsCache.getValue().map { it.packageName.asString() }
         }
 
     override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String> =

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirGeneratedClassDeclaredMemberScope.kt

@@ -74,7 +74,7 @@ class FirGeneratedClassDeclaredMemberScope private constructor(
         generateMemberProperties(callableId)
     }
 
-    private val constructorCache: FirLazyValue<List<FirConstructorSymbol>, Nothing?> = firCachesFactory.createLazyValue {
+    private val constructorCache: FirLazyValue<List<FirConstructorSymbol>> = firCachesFactory.createLazyValue {
         generateConstructors()
     }
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/caches/FirCacheWithPostCompute.kt

@@ -20,17 +20,10 @@ operator fun <K : Any, V> FirCache<K, V, Nothing>.contains(key: K): Boolean {
     return getValueIfComputed(key) != null
 }
 
-class FirLazyValue<out V, in CONTEXT>(private val cache: FirCache<Unit, V, CONTEXT>) {
-    fun getValue(context: CONTEXT): V {
-        return cache.getValue(Unit, context)
-    }
+abstract class FirLazyValue<out V> {
+    abstract fun getValue(): V
 }
 
-operator fun <V> FirLazyValue<V, Nothing?>.getValue(thisRef: Any?, property: KProperty<*>): V {
-    return getValue(context = null)
-}
-
-@Suppress("NOTHING_TO_INLINE")
-inline fun <V> FirLazyValue<V, Nothing?>.getValue(): V {
-    return getValue(null)
-}
+operator fun <V> FirLazyValue<V>.getValue(thisRef: Any?, property: KProperty<*>): V {
+    return getValue()
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/caches/FirCachesFactory.kt

@@ -59,9 +59,7 @@ abstract class FirCachesFactory : FirSessionComponent {
         postCompute: (K, V, DATA) -> Unit
     ): FirCache<K, V, CONTEXT>
 
-    fun <V, CONTEXT> createLazyValue(createValue: (CONTEXT) -> V): FirLazyValue<V, CONTEXT> {
-        return FirLazyValue(createCache { _, context -> createValue(context) })
-    }
+    abstract fun <V> createLazyValue(createValue: () -> V): FirLazyValue<V>
 }
 
 val FirSession.firCachesFactory: FirCachesFactory by FirSession.sessionComponentAccessor()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/caches/FirThreadUnsafeCachesFactory.kt

@@ -24,6 +24,9 @@ object FirThreadUnsafeCachesFactory : FirCachesFactory() {
         postCompute: (K, V, DATA) -> Unit
     ): FirCache<K, V, CONTEXT> =
         FirThreadUnsafeCacheWithPostCompute(createValue, postCompute)
+
+    override fun <V> createLazyValue(createValue: () -> V): FirLazyValue<V> =
+        FirThreadUnsafeValue(createValue)
 }
 
 @Suppress("UNCHECKED_CAST")
@@ -63,3 +66,8 @@ private class FirThreadUnsafeCacheWithPostCompute<K : Any, V, CONTEXT, DATA>(
     override fun getValueIfComputed(key: K): V? =
         map.getOrElse(key) { null as V }
 }
+
+private class FirThreadUnsafeValue<V>(createValue: () -> V) : FirLazyValue<V>() {
+    private val lazyValue by lazy(LazyThreadSafetyMode.NONE, createValue)
+    override fun getValue(): V = lazyValue
+}