From c60fd38e05ff9dd0a498de80e4fe205d796bdf4f Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Thu, 16 Feb 2023 14:10:26 +0100
Subject: [PATCH] [LL FIR] allow recursive access to the same
 FirThreadSafeCache

^KT-56721 fixed
---
 .../low/level/api/fir/caches/FirThreadSafeCache.kt     | 10 +---------
 .../analysis/low/level/api/fir/caches/NullValue.kt     |  2 +-
 .../extensions/FirDeclarationGenerationExtension.kt    |  6 ++++++
 .../jetbrains/kotlin/fir/caches/FirCachesFactory.kt    |  3 +++
 4 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/FirThreadSafeCache.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/FirThreadSafeCache.kt
index 94df2ac2c50..c2d7c927922 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/FirThreadSafeCache.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/FirThreadSafeCache.kt
@@ -13,17 +13,9 @@ internal class FirThreadSafeCache<K : Any, V, CONTEXT>(
     private val createValue: (K, CONTEXT) -> V
 ) : FirCache<K, V, CONTEXT>() {
 
-    private val updating: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
-
     override fun getValue(key: K, context: CONTEXT): V =
         map.getOrPutWithNullableValue(key) {
-            if (updating.get()) error("Recursive update")
-            try {
-                updating.set(true)
-                createValue(it, context)
-            } finally {
-                updating.set(false)
-            }
+            createValue(it, context)
         }
 
     override fun getValueIfComputed(key: K): V? =
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/NullValue.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/NullValue.kt
index 930a2921661..fe1bad1614a 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/NullValue.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/NullValue.kt
@@ -19,6 +19,6 @@ internal inline fun <KEY : Any, RESULT> ConcurrentMap<KEY, Any>.getOrPutWithNull
     key: KEY,
     crossinline compute: (KEY) -> Any?
 ): RESULT {
-    val value = computeIfAbsent(key) { compute(key) ?: NullValue }
+    val value = getOrPut(key) { compute(key) ?: NullValue }
     return value.nullValueToNull()
 }
\ No newline at end of file
diff --git a/compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt b/compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt
index c99f4666848..4c50c598310 100644
--- a/compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt
+++ b/compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt
@@ -23,6 +23,12 @@ import kotlin.reflect.KClass
  * TODO:
  *  - check that annotations or meta-annotations is not empty
  */
+
+/**
+ * All `generate*` members have the contract that the computation should be side-effect-free.
+ * That means that all `generate*` function implementations should not modify any state or leak the generated `FirElement` or `FirBasedSymbol` (e.g., by putting it to some cache).
+ * This restriction is imposed by the corresponding IDE cache implementation, which might retry the computation several times.
+ */
 abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExtension(session) {
     companion object {
         val NAME = FirExtensionPointName("ExistingClassModification")
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/caches/FirCachesFactory.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/caches/FirCachesFactory.kt
index 0bdecd9af3d..304196ea7a8 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/caches/FirCachesFactory.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/caches/FirCachesFactory.kt
@@ -15,6 +15,9 @@ abstract class FirCachesFactory : FirSessionComponent {
      *
      * [FirCache.getValue] should not be called inside [createValue]
      *
+     * Note, that [createValue] might be called multiple times for the same value,
+     * but all threads will always get the same value
+     *
      * Where:
      * [CONTEXT] -- type of value which be used to create value by [createValue]
      */