From 7f4af06ac732925aa6be9e31bb22a2fe6a380fbc Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Thu, 13 Jul 2023 13:55:12 +0200
Subject: [PATCH] [LL FIR] KT-58257 Turn `LLFirSession.isValid` into a volatile
 variable

- We do not need `AtomicBoolean` here, because setting `_isValid` to a
  constant `false` suffices.
- Also replace the public setter of `isValid` with `markInvalid` to
  prohibit setting `isValid` to `true` again.
---
 .../level/api/fir/sessions/LLFirSession.kt    | 21 +++++++++----------
 .../api/fir/sessions/LLFirSessionCache.kt     |  6 +++---
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
index dfe2a92a7a6..2bfe250491d 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.fir.PrivateSessionConstructor
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import java.lang.ref.WeakReference
-import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicLong
 
 @OptIn(PrivateSessionConstructor::class)
@@ -32,20 +31,20 @@ abstract class LLFirSession(
     val project: Project
         get() = ktModule.project
 
-    private val _isValid = AtomicBoolean(true)
-
     /**
      * Whether the [LLFirSession] is valid. The session should not be used if it is invalid.
-     *
-     * [isValid] should be set to `false` at the same time as the session is removed from [LLFirSessionCache]. Hence, [isValid] should be
+     */
+    @Volatile
+    var isValid: Boolean = true
+        private set
+
+    /**
+     * [markInvalid] should be called at the same time as the session is removed from [LLFirSessionCache]. Hence, session validity should be
      * managed by [LLFirSessionCache].
      */
-    var isValid: Boolean
-        get() = _isValid.get()
-        internal set(value) {
-            check(!value) { "An invalid LL FIR session cannot become valid again." }
-            _isValid.set(value)
-        }
+    internal fun markInvalid() {
+        isValid = false
+    }
 
     fun invalidate() {
         val application = ApplicationManager.getApplication()
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
index 7bf8cab4b50..b49a771cb8d 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
@@ -86,7 +86,7 @@ class LLFirSessionCache(private val project: Project) {
 
     private fun removeSessionFrom(module: KtModule, storage: SessionStorage): Boolean {
         val session = storage.remove(module) ?: return false
-        session.isValid = false
+        session.markInvalid()
         return true
     }
 
@@ -122,7 +122,7 @@ class LLFirSessionCache(private val project: Project) {
     private fun removeAllSessionsFrom(storage: SessionStorage) {
         // Because `removeAllSessionsFrom` is executed in a write action, the order of setting `isValid` and clearing `storage` is not
         // important.
-        storage.values.forEach { it.isValid = false }
+        storage.values.forEach { it.markInvalid() }
         storage.clear()
     }
 
@@ -133,7 +133,7 @@ class LLFirSessionCache(private val project: Project) {
         // "collect and remove" approach also works.
         val scriptEntries = storage.entries.filter { (module, _) -> shouldBeRemoved(module) }
         for ((module, session) in scriptEntries) {
-            session.isValid = false
+            session.markInvalid()
             storage.remove(module)
         }
     }