[FIR] add possibility to specify initialCapacity and loadFactor in FirCachesFactory.createCache

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

@@ -9,9 +9,9 @@ import org.jetbrains.kotlin.fir.caches.FirCache
 import java.util.concurrent.ConcurrentHashMap
 
 internal class FirThreadSafeCache<K : Any, V, CONTEXT>(
+    private val map: ConcurrentHashMap<K, Any> = ConcurrentHashMap<K, Any>(),
     private val createValue: (K, CONTEXT) -> V
 ) : FirCache<K, V, CONTEXT>() {
-    private val map = ConcurrentHashMap<K, Any>()
 
     override fun getValue(key: K, context: CONTEXT): V =
         map.getOrPutWithNullableValue(key) { createValue(it, context) }

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

@@ -5,11 +5,24 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.fir.caches
 
-import org.jetbrains.kotlin.fir.caches.*
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.FirCachesFactory
+import java.util.concurrent.ConcurrentHashMap
 
 object FirThreadSafeCachesFactory : FirCachesFactory() {
     override fun <KEY : Any, VALUE, CONTEXT> createCache(createValue: (KEY, CONTEXT) -> VALUE): FirCache<KEY, VALUE, CONTEXT> =
-        FirThreadSafeCache(createValue)
+        FirThreadSafeCache(createValue = createValue)
+
+    override fun <K : Any, V, CONTEXT> createCache(
+        initialCapacity: Int,
+        loadFactor: Float,
+        createValue: (K, CONTEXT) -> V
+    ): FirCache<K, V, CONTEXT> =
+        FirThreadSafeCache(
+            ConcurrentHashMap<K, Any>(initialCapacity, loadFactor),
+            createValue
+        )
+
 
     override fun <KEY : Any, VALUE, CONTEXT, DATA> createCacheWithPostCompute(
         createValue: (KEY, CONTEXT) -> Pair<VALUE, DATA>,

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

@@ -20,6 +20,24 @@ abstract class FirCachesFactory : FirSessionComponent {
      */
     abstract fun <K : Any, V, CONTEXT> createCache(createValue: (K, CONTEXT) -> V): FirCache<K, V, CONTEXT>
 
+    /**
+     * Creates a cache with returns a value by key on demand if it is computed
+     * Otherwise computes the value in [createValue] and caches it for future invocations
+     *
+     * [FirCache.getValue] should not be called inside [createValue]
+     *
+     * Where:
+     * [CONTEXT] -- type of value which be used to create value by [createValue]
+     *
+     * @param initialCapacity initial capacity for the underlying cache map
+     * @param loadFactor loadFactor for the underlying cache map
+     */
+    abstract fun <K : Any, V, CONTEXT> createCache(
+        initialCapacity: Int,
+        loadFactor: Float,
+        createValue: (K, CONTEXT) -> V
+    ): FirCache<K, V, CONTEXT>
+
     /**
      * Creates a cache with returns a caches value on demand if it is computed
      * Otherwise computes the value in two phases:

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

@@ -7,7 +7,17 @@ package org.jetbrains.kotlin.fir.caches
 
 object FirThreadUnsafeCachesFactory : FirCachesFactory() {
     override fun <K : Any, V, CONTEXT> createCache(createValue: (K, CONTEXT) -> V): FirCache<K, V, CONTEXT> =
-        FirThreadUnsafeCache(createValue)
+        FirThreadUnsafeCache(createValue = createValue)
+
+    override fun <K : Any, V, CONTEXT> createCache(
+        initialCapacity: Int,
+        loadFactor: Float,
+        createValue: (K, CONTEXT) -> V
+    ): FirCache<K, V, CONTEXT> =
+        FirThreadUnsafeCache(
+            NullableMap(HashMap(initialCapacity, loadFactor)),
+            createValue
+        )
 
     override fun <K : Any, V, CONTEXT, DATA> createCacheWithPostCompute(
         createValue: (K, CONTEXT) -> Pair<V, DATA>,
@@ -18,9 +28,9 @@ object FirThreadUnsafeCachesFactory : FirCachesFactory() {
 
 @Suppress("UNCHECKED_CAST")
 private class FirThreadUnsafeCache<K : Any, V, CONTEXT>(
+    private val map: NullableMap<K, V> = NullableMap<K, V>(),
     private val createValue: (K, CONTEXT) -> V
 ) : FirCache<K, V, CONTEXT>() {
-    private val map = NullableMap<K, V>()
 
     override fun getValue(key: K, context: CONTEXT): V =
         map.getOrElse(key) {