[SLC] migrate from AtomicReference to AtomicFieldUpdater

^KT-56046

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/LazyAnnotationsBox.kt

@@ -7,21 +7,22 @@ package org.jetbrains.kotlin.light.classes.symbol.annotations
 
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiModifierList
+import com.intellij.util.concurrency.AtomicFieldUpdater
 import org.jetbrains.kotlin.light.classes.symbol.toArrayIfNotEmptyOrDefault
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.utils.SmartList
-import java.util.concurrent.atomic.AtomicReference
 
 internal class LazyAnnotationsBox(
     private val annotationsProvider: AnnotationsProvider,
     private val additionalAnnotationsProvider: AdditionalAnnotationsProvider = EmptyAdditionalAnnotationsProvider,
     private val annotationFilter: AnnotationFilter = AlwaysAllowedAnnotationFilter,
 ) : AnnotationsBox {
-    private val cachedCollection: AtomicReference<Collection<PsiAnnotation>?> = AtomicReference()
+    @Volatile
+    private var cachedAnnotations: Collection<PsiAnnotation>? = null
 
     private fun getOrComputeCachedAnnotations(owner: PsiModifierList): Collection<PsiAnnotation> {
-        cachedCollection.get()?.let { return it }
+        cachedAnnotations?.let { return it }
 
         val annotations = annotationsProvider.annotationInfos().mapNotNullTo(SmartList<PsiAnnotation>()) { applicationInfo ->
             applicationInfo.classId?.let { _ ->
@@ -33,7 +34,7 @@ internal class LazyAnnotationsBox(
         additionalAnnotationsProvider.addAllAnnotations(annotations, foundQualifiers, owner)
 
         val resultAnnotations = annotationFilter.filtered(annotations)
-        cachedCollection.compareAndSet(null, resultAnnotations)
+        fieldUpdater.compareAndSet(this, null, resultAnnotations)
 
         return getOrComputeCachedAnnotations(owner)
     }
@@ -50,7 +51,7 @@ internal class LazyAnnotationsBox(
     fun findAnnotation(owner: PsiModifierList, qualifiedName: String, withAdditionalAnnotations: Boolean): PsiAnnotation? {
         if (!annotationFilter.isAllowed(qualifiedName)) return null
 
-        cachedCollection.get()?.let { annotations ->
+        cachedAnnotations?.let { annotations ->
             return annotations.find { it.qualifiedName == qualifiedName }
         }
 
@@ -69,7 +70,7 @@ internal class LazyAnnotationsBox(
     override fun hasAnnotation(owner: PsiModifierList, qualifiedName: String): Boolean {
         if (!annotationFilter.isAllowed(qualifiedName)) return false
 
-        cachedCollection.get()?.let { annotations ->
+        cachedAnnotations?.let { annotations ->
             return annotations.any { it.qualifiedName == qualifiedName }
         }
 
@@ -82,6 +83,8 @@ internal class LazyAnnotationsBox(
     }
 
     companion object {
+        private val fieldUpdater = AtomicFieldUpdater.forFieldOfType(LazyAnnotationsBox::class.java, Collection::class.java)
+
         /**
          * @see org.jetbrains.kotlin.fir.resolve.transformers.plugin.CompilerRequiredAnnotationsHelper
          */

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/SimpleAnnotationsBox.kt

@@ -7,17 +7,19 @@ package org.jetbrains.kotlin.light.classes.symbol.annotations
 
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiModifierList
+import com.intellij.util.concurrency.AtomicFieldUpdater
 import org.jetbrains.kotlin.light.classes.symbol.toArrayIfNotEmptyOrDefault
-import java.util.concurrent.atomic.AtomicReference
 
 internal class SimpleAnnotationsBox(private val annotationsComputer: (PsiModifierList) -> Collection<PsiAnnotation>) : AnnotationsBox {
-    private val cachedCollection: AtomicReference<Collection<PsiAnnotation>?> = AtomicReference()
+    @Volatile
+    private var cachedAnnotations: Collection<PsiAnnotation>? = null
 
     private fun getOrComputeAnnotations(owner: PsiModifierList): Collection<PsiAnnotation> {
-        cachedCollection.get()?.let { return it }
+        cachedAnnotations?.let { return it }
 
         val nonCachedAnnotations = annotationsComputer(owner)
-        cachedCollection.compareAndSet(null, nonCachedAnnotations)
+        fieldUpdater.compareAndSet(this, null, nonCachedAnnotations)
+
         return getOrComputeAnnotations(owner)
     }
 
@@ -29,4 +31,8 @@ internal class SimpleAnnotationsBox(private val annotationsComputer: (PsiModifie
         owner: PsiModifierList,
         qualifiedName: String,
     ): PsiAnnotation? = getOrComputeAnnotations(owner).find { it.qualifiedName == qualifiedName }
+
+    companion object {
+        private val fieldUpdater = AtomicFieldUpdater.forFieldOfType(SimpleAnnotationsBox::class.java, Collection::class.java)
+    }
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/modifierLists/LazyModifiersBox.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.light.classes.symbol.modifierLists
 
 import com.intellij.psi.PsiModifier
+import com.intellij.util.concurrency.AtomicFieldUpdater
 import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.toPersistentHashMap
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithModality
@@ -16,7 +17,6 @@ import org.jetbrains.kotlin.light.classes.symbol.computeSimpleModality
 import org.jetbrains.kotlin.light.classes.symbol.toPsiVisibilityForMember
 import org.jetbrains.kotlin.light.classes.symbol.withSymbol
 import org.jetbrains.kotlin.utils.keysToMap
-import java.util.concurrent.atomic.AtomicReference
 
 internal typealias LazyModifiersComputer = (modifier: String) -> Map<String, Boolean>?
 
@@ -24,19 +24,26 @@ internal class LazyModifiersBox(
     initialValue: Map<String, Boolean> = emptyMap(),
     private val computer: LazyModifiersComputer,
 ) : ModifiersBox {
-    private val modifiersMapReference: AtomicReference<PersistentMap<String, Boolean>> = AtomicReference(initialValue.toPersistentHashMap())
+    @Volatile
+    private var modifiersMapReference: PersistentMap<String, Boolean> = initialValue.toPersistentHashMap()
 
     override fun hasModifier(modifier: String): Boolean {
-        modifiersMapReference.get()[modifier]?.let { return it }
+        modifiersMapReference[modifier]?.let { return it }
+
         val newValues = computer(modifier) ?: mapOf(modifier to false)
-        modifiersMapReference.updateAndGet {
-            it.putAll(newValues)
-        }
+        do {
+            val currentMap = modifiersMapReference
+            currentMap[modifier]?.let { return it }
+
+            val newMap = currentMap.putAll(newValues)
+        } while (fieldUpdater.compareAndSet(/* owner = */ this, /* expected = */ currentMap, /* newValue = */ newMap))
 
         return newValues[modifier] ?: error("Inconsistent state: $modifier")
     }
 
     companion object {
+        private val fieldUpdater = AtomicFieldUpdater.forFieldOfType(LazyModifiersBox::class.java, PersistentMap::class.java)
+
         internal val VISIBILITY_MODIFIERS = setOf(PsiModifier.PUBLIC, PsiModifier.PACKAGE_LOCAL, PsiModifier.PROTECTED, PsiModifier.PRIVATE)
         internal val VISIBILITY_MODIFIERS_MAP: PersistentMap<String, Boolean> =
             VISIBILITY_MODIFIERS.keysToMap {