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

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