FIR: introduce caches factory

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.fir.session
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.CheckersComponent
+import org.jetbrains.kotlin.fir.caches.FirCachesFactory
+import org.jetbrains.kotlin.fir.caches.FirThreadUnsafeCachesFactory
 import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
@@ -40,6 +42,12 @@ fun FirSession.registerCommonComponents(languageVersionSettings: LanguageVersion
     register(InferenceComponents::class, InferenceComponents(this))
 }
 
+@OptIn(SessionConfiguration::class)
+fun FirSession.registerThreadUnsafeCaches() {
+    register(FirCachesFactory::class, FirThreadUnsafeCachesFactory)
+}
+
+
 // -------------------------- Resolve components --------------------------
 
 /*

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirSessionFactory.kt

@@ -19,6 +19,8 @@ import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationChecker
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.checkersComponent
 import org.jetbrains.kotlin.fir.analysis.extensions.additionalCheckers
+import org.jetbrains.kotlin.fir.caches.FirCachesFactory
+import org.jetbrains.kotlin.fir.caches.FirThreadUnsafeCachesFactory
 import org.jetbrains.kotlin.fir.checkers.registerCommonCheckers
 import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
@@ -68,6 +70,7 @@ object FirSessionFactory {
         init: FirSessionConfigurator.() -> Unit = {}
     ): FirJavaModuleBasedSession {
         return FirJavaModuleBasedSession(moduleInfo, sessionProvider).apply {
+            registerThreadUnsafeCaches()
             registerCommonComponents(languageVersionSettings)
             registerResolveComponents()
             registerJavaSpecificResolveComponents()
@@ -113,6 +116,7 @@ object FirSessionFactory {
 
         val kotlinClassFinder = VirtualFileFinderFactory.getInstance(project).create(scope)
         return FirLibrarySession(moduleInfo, sessionProvider).apply {
+            registerThreadUnsafeCaches()
             registerCommonComponents(languageVersionSettings)
 
             val javaSymbolProvider = JavaSymbolProvider(this, project, scope)

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

@@ -0,0 +1,15 @@
+/*
+ * 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.fir.caches
+
+abstract class FirCache<in KEY : Any, out VALUE, in CONTEXT> {
+    abstract fun getValue(key: KEY, context: CONTEXT): VALUE
+    abstract fun getValueIfComputed(key: KEY): VALUE?
+}
+
+@Suppress("NOTHING_TO_INLINE")
+inline fun <KEY : Any, VALUE> FirCache<KEY, VALUE, Nothing?>.getValue(key: KEY): VALUE =
+    getValue(key, null)

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

@@ -0,0 +1,74 @@
+/*
+ * 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.fir.caches
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+
+abstract class FirCachesFactory : FirSessionComponent {
+    /**
+     * 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]
+     */
+    abstract fun <KEY : Any, VALUE, CONTEXT> createCache(createValue: (KEY, CONTEXT) -> VALUE): FirCache<KEY, VALUE, CONTEXT>
+
+    /**
+     * Creates a cache with returns a caches value on demand if it is computed
+     * Otherwise computes the value in two phases:
+     *  - [createValue] -- creates values and stores [VALUE] to cache and passes [VALUE] & [DATA] to [postCompute]
+     *  - [postCompute] -- performs some operations on computed value after it placed into map
+     *
+     * [FirCache.getValue] can be safely called in postCompute from the same thread and correct value computed by [createValue] will be returned
+     * [FirCache.getValue] should not be called inside [createValue]
+     *
+     * Where:
+     *  [CONTEXT] -- type of value which be used to create value by [createValue]
+     *  [DATA] -- type of additional data which will be passed from [createValue] to [postCompute]
+     */
+    abstract fun <KEY : Any, VALUE, CONTEXT, DATA> createCacheWithPostCompute(
+        createValue: (KEY, CONTEXT) -> Pair<VALUE, DATA>,
+        postCompute: (KEY, VALUE, DATA) -> Unit
+    ): FirCache<KEY, VALUE, CONTEXT>
+}
+
+val FirSession.firCachesFactory: FirCachesFactory by FirSession.sessionComponentAccessor()
+
+inline fun <KEY : Any, VALUE> FirCachesFactory.createCache(
+    crossinline createValue: (KEY) -> VALUE,
+): FirCache<KEY, VALUE, Nothing?> = createCache(
+    createValue = { key, _ -> createValue(key) },
+)
+
+inline fun <KEY : Any, VALUE, CONTEXT> FirCachesFactory.createCacheWithPostCompute(
+    crossinline createValue: (KEY, CONTEXT) -> VALUE,
+    crossinline postCompute: (KEY, VALUE) -> Unit
+): FirCache<KEY, VALUE, CONTEXT> = createCacheWithPostCompute(
+    createValue = { key, context -> createValue(key, context) to null },
+    postCompute = { key, value, _ -> postCompute(key, value) }
+)
+
+inline fun <KEY : Any, VALUE> FirCachesFactory.createCacheWithPostCompute(
+    crossinline createValue: (KEY) -> VALUE,
+    crossinline postCompute: (KEY, VALUE) -> Unit
+): FirCache<KEY, VALUE, Nothing?> = createCacheWithPostCompute(
+    createValue = { key, _ -> createValue(key) to null },
+    postCompute = { key, value, _ -> postCompute(key, value) }
+)
+
+inline fun <KEY : Any, VALUE, DATA> FirCachesFactory.createCacheWithPostCompute(
+    crossinline createValue: (KEY) -> Pair<VALUE, DATA>,
+    crossinline postCompute: (KEY, VALUE, DATA) -> Unit
+): FirCache<KEY, VALUE, Nothing?> = createCacheWithPostCompute(
+    createValue = { key, _ -> createValue(key) },
+    postCompute = { key, value, data -> postCompute(key, value, data) }
+)
+
+

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

@@ -0,0 +1,55 @@
+/*
+ * 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.fir.caches
+
+object FirThreadUnsafeCachesFactory : FirCachesFactory() {
+    override fun <KEY : Any, VALUE, CONTEXT> createCache(createValue: (KEY, CONTEXT) -> VALUE): FirCache<KEY, VALUE, CONTEXT> =
+        FirThreadUnsafeCache(createValue)
+
+    override fun <KEY : Any, VALUE, CONTEXT, DATA> createCacheWithPostCompute(
+        createValue: (KEY, CONTEXT) -> Pair<VALUE, DATA>,
+        postCompute: (KEY, VALUE, DATA) -> Unit
+    ): FirCache<KEY, VALUE, CONTEXT> =
+        FirThreadUnsafeCacheWithPostCompute(createValue, postCompute)
+}
+
+@Suppress("UNCHECKED_CAST")
+private class FirThreadUnsafeCache<KEY : Any, VALUE, CONTEXT>(
+    private val createValue: (KEY, CONTEXT) -> VALUE
+) : FirCache<KEY, VALUE, CONTEXT>() {
+    private val map = NullableMap<KEY, VALUE>()
+
+    override fun getValue(key: KEY, context: CONTEXT): VALUE =
+        map.getOrElse(key) {
+            createValue(key, context).also { createdValue ->
+                map[key] = createdValue
+            }
+        }
+
+    override fun getValueIfComputed(key: KEY): VALUE? =
+        map.getOrElse(key) { null as VALUE }
+}
+
+
+private class FirThreadUnsafeCacheWithPostCompute<KEY : Any, VALUE, CONTEXT, DATA>(
+    private val createValue: (KEY, CONTEXT) -> Pair<VALUE, DATA>,
+    private val postCompute: (KEY, VALUE, DATA) -> Unit
+) : FirCache<KEY, VALUE, CONTEXT>() {
+    private val map = NullableMap<KEY, VALUE>()
+
+    override fun getValue(key: KEY, context: CONTEXT): VALUE =
+        map.getOrElse(key) {
+            val (createdValue, data) = createValue(key, context)
+            map[key] = createdValue
+            postCompute(key, createdValue, data)
+            createdValue
+        }
+
+
+    @Suppress("UNCHECKED_CAST")
+    override fun getValueIfComputed(key: KEY): VALUE? =
+        map.getOrElse(key) { null as VALUE }
+}

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

@@ -0,0 +1,36 @@
+/*
+ * 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.fir.caches
+
+import org.jetbrains.kotlin.fir.PrivateForInline
+
+/**
+ * [Map] which allows store null values
+ */
+@OptIn(PrivateForInline::class)
+internal inline class NullableMap<KEY, VALUE>(private val map: MutableMap<KEY, Any> = HashMap()) {
+
+    /**
+     * Get value if it is present in map
+     * Execute [orElse] otherwise and return it result,
+     * [orElse] can modify the map inside
+     */
+    @Suppress("UNCHECKED_CAST")
+    inline fun getOrElse(key: KEY, orElse: () -> VALUE): VALUE =
+        when (val value = map[key]) {
+            null -> orElse()
+            NullValue -> null
+            else -> value
+        } as VALUE
+
+    @Suppress("NOTHING_TO_INLINE")
+    inline operator fun set(key: KEY, value: VALUE) {
+        map[key] = value ?: NullValue
+    }
+}
+
+@PrivateForInline
+internal object NullValue