From 41ebf5ed531d2cbf14f727b5f301b5ef71cd802d Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Tue, 8 Nov 2022 12:44:29 +0100
Subject: [PATCH] [AA] KtPsiBasedSymbolPointer: simplify test logic

^KT-54311
---
 .../pointers/KtPsiBasedSymbolPointer.kt       | 21 ++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/pointers/KtPsiBasedSymbolPointer.kt b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/pointers/KtPsiBasedSymbolPointer.kt
index dde9b3803b7..8238eecd358 100644
--- a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/pointers/KtPsiBasedSymbolPointer.kt
+++ b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/pointers/KtPsiBasedSymbolPointer.kt
@@ -36,7 +36,9 @@ public class KtPsiBasedSymbolPointer<S : KtSymbol> private constructor(private v
 
     public companion object {
         public fun <S : KtSymbol> createForSymbolFromSource(symbol: S): KtPsiBasedSymbolPointer<S>? {
-            if (disablePsiPointer || symbol.origin != KtSymbolOrigin.SOURCE) return null
+            ifDisabled { return null }
+
+            if (symbol.origin != KtSymbolOrigin.SOURCE) return null
 
             val psi = when (val psi = symbol.psi) {
                 is KtDeclaration -> psi
@@ -49,20 +51,29 @@ public class KtPsiBasedSymbolPointer<S : KtSymbol> private constructor(private v
         }
 
         public fun <S : KtSymbol> createForSymbolFromPsi(ktElement: KtElement): KtPsiBasedSymbolPointer<S>? {
-            if (disablePsiPointer) return null
+            ifDisabled { return null }
 
             return KtPsiBasedSymbolPointer(ktElement.createSmartPointer())
         }
 
         @TestOnly
-        @Synchronized
         public fun <T> withDisabledPsiBasedPointers(disable: Boolean, action: () -> T): T = try {
-            disablePsiPointer = disable
+            disablePsiPointer = true
+            disablePsiPointerFlag.set(disable)
             action()
         } finally {
-            disablePsiPointer = false
+            disablePsiPointerFlag.remove()
         }
 
+        private inline fun ifDisabled(action: () -> Unit) {
+            if (!disablePsiPointer) return
+            if (disablePsiPointerFlag.get()) {
+                action()
+            }
+        }
+
+        private val disablePsiPointerFlag: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
+
         @Volatile
         private var disablePsiPointer: Boolean = false
     }