FIR IDE: allow the same entity to be accessed from different threads

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/ValidityOwner.kt

@@ -33,7 +33,6 @@ interface ValidityOwnerByValidityToken : ValidityOwner {
 
 class ReadActionConfinementValidityToken(project: Project) : ValidityOwner {
     private val modificationTracker = KotlinModificationTrackerService.getInstance(project).modificationTracker
-    private val ownerThread = Thread.currentThread()
     private val onCreatedTimeStamp = modificationTracker.modificationCount
 
 
@@ -41,7 +40,6 @@ class ReadActionConfinementValidityToken(project: Project) : ValidityOwner {
         val application = ApplicationManager.getApplication()
         if (application.isDispatchThread) return false
         if (!application.isReadAccessAllowed) return false
-        if (ownerThread != Thread.currentThread()) return false
         return onCreatedTimeStamp == modificationTracker.modificationCount
     }
 
@@ -49,9 +47,6 @@ class ReadActionConfinementValidityToken(project: Project) : ValidityOwner {
         val application = ApplicationManager.getApplication()
         if (application.isDispatchThread) return "Called in EDT thread"
         if (!application.isReadAccessAllowed) return "Called outside read action"
-        if (ownerThread != Thread.currentThread()) {
-            return "Called outside the thread was created in (was created in $ownerThread, but accessed in ${Thread.currentThread()}"
-        }
         if (onCreatedTimeStamp != modificationTracker.modificationCount) return "PSI has changed since creation"
         error("Getting invalidation reason for valid invalidatable")
     }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtSymbolByFirBuilder.kt

@@ -5,9 +5,9 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir
 
+import com.google.common.collect.MapMaker
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
-import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
@@ -41,6 +41,9 @@ import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.idea.stubindex.PackageIndexUtil
 
+/**
+ * Maps FirElement to KtSymbol & ConeType to KtType, thread safe
+ */
 internal class KtSymbolByFirBuilder(
     firProvider: FirSymbolProvider,
     typeCheckerContext: ConeTypeCheckerContext,
@@ -149,8 +152,7 @@ internal class KtSymbolByFirBuilder(
 }
 
 private class BuilderCache<From, To> {
-    // the capacity & load factor values here are default ones which are used ing java.util.Map
-    private val cache = ContainerUtil.createWeakMap<From, To>(16, .75f, ContainerUtil.identityStrategy())
+    private val cache = MapMaker().weakKeys().makeMap<From, To>()
 
     inline fun <reified S : To> cache(key: From, calculation: () -> S): S =
         cache.getOrPut(key, calculation) as S

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/ValidityAwareCachedValue.kt

@@ -11,24 +11,20 @@ import org.jetbrains.kotlin.idea.frontend.api.assertIsValid
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KProperty
 
-// all access are considered to be from a single thread
-// which should be checked in invalidatable.assertIsValid()
+/**
+ * Lazy value that guaranties safe publication and checks validity on every access
+ */
 class ValidityAwareCachedValue<T>(
     private val validityOwner: ValidityOwner,
-    private val init: () -> T
+    init: () -> T
 ) : ReadOnlyProperty<Any, T> {
-    var value: Any? = UNITITIALIZED
+    private val lazyValue = lazy(LazyThreadSafetyMode.PUBLICATION, init)
 
     @Suppress("UNCHECKED_CAST")
     override fun getValue(thisRef: Any, property: KProperty<*>): T {
         validityOwner.assertIsValid()
-        if (value === UNITITIALIZED) {
-            value = init()
-        }
-        return value as T
+        return lazyValue.value
     }
 }
 
 internal fun <T> ValidityOwnerByValidityToken.cached(init: () -> T) = ValidityAwareCachedValue(token, init)
-
-internal object UNITITIALIZED