FIR IDE: refactor lock provider

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/FirModuleResolveState.kt

@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.FirTransformerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
-import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.withReadLock
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.ktDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/FirFileBuilder.kt

@@ -68,10 +68,8 @@ internal class FirFileBuilder(
     inline fun <R> runCustomResolveUnderLock(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R =
         cache.firFileLockProvider.withWriteLock(firFile) { resolve() }
 
-    inline fun <R : Any> runCustomResolveWithPCECheck(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R {
+    inline fun <R : Any> runCustomResolveWithPCECheck(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R =
-        val lock = cache.firFileLockProvider.getLockFor(firFile)
+        cache.firFileLockProvider.withWriteLockPCECheck(firFile, LOCKING_INTERVAL_MS, resolve)
-        return lock.writeLock().lockWithPCECheck(LOCKING_INTERVAL_MS) { resolve() }
-    }
 
     fun runResolveWithoutLock(
         firFile: FirFile,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/LockProvider.kt

@@ -6,28 +6,33 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder
 
 import com.google.common.collect.MapMaker
+import com.intellij.openapi.diagnostic.logger
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.lockWithPCECheck
 import java.util.concurrent.ConcurrentMap
 import java.util.concurrent.locks.ReadWriteLock
-import java.util.concurrent.locks.ReentrantLock
+import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.withLock
 
-internal class LockProvider<KEY, out LOCK>(private val createLock: () -> LOCK) {
+internal class LockProvider<KEY> {
-    private val locks: ConcurrentMap<KEY, LOCK> = MapMaker().weakKeys().makeMap()
+    private val locks: ConcurrentMap<KEY, ReadWriteLock> = MapMaker().weakKeys().makeMap()
-    fun getLockFor(key: KEY) = locks.getOrPut(key) { createLock() }
-}
 
-internal inline fun <KEY, R> LockProvider<KEY, ReadWriteLock>.withReadLock(key: KEY, action: () -> R): R {
+    @Suppress("NOTHING_TO_INLINE")
-    val readLock = getLockFor(key).readLock()
+    private inline fun getLockFor(key: KEY) = locks.getOrPut(key) { ReentrantReadWriteLock() }
-    return readLock.withLock { action() }
-}
 
-internal inline fun <KEY, R> LockProvider<KEY, ReadWriteLock>.withWriteLock(key: KEY, action: () -> R): R {
+    inline fun <R> withReadLock(key: KEY, action: () -> R): R {
-    val writeLock = getLockFor(key).writeLock()
+        val readLock = getLockFor(key).readLock()
-    return writeLock.withLock { action() }
+        return readLock.withLock { action() }
-}
+    }
 
 
-internal inline fun <KEY, R> LockProvider<KEY, ReentrantLock>.withLock(key: KEY, action: () -> R): R {
+    inline fun <R> withWriteLock(key: KEY, action: () -> R): R {
-    val lock = getLockFor(key)
+        val writeLock = getLockFor(key).writeLock()
-    return lock.withLock { action() }
+        return writeLock.withLock { action() }
+    }
+
+    inline fun <R> withWriteLockPCECheck(key: KEY, lockingIntervalMs: Long, action: () -> R): R {
+        val writeLock = getLockFor(key).writeLock()
+        return writeLock.lockWithPCECheck(lockingIntervalMs, action)
+    }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/ModuleFileCache.kt

@@ -51,7 +51,7 @@ internal abstract class ModuleFileCache {
 
     abstract fun getCachedFirFile(ktFile: KtFile): FirFile?
 
-    abstract val firFileLockProvider: LockProvider<FirFile, ReentrantReadWriteLock>
+    abstract val firFileLockProvider: LockProvider<FirFile>
 
     inline fun <D : FirDeclaration, R> withReadLockOn(declaration: D, action: (D) -> R): R {
         val file = getContainerFirFile(declaration)
@@ -76,5 +76,5 @@ internal class ModuleFileCacheImpl(override val session: FirSession) : ModuleFil
         return getCachedFirFile(ktFile)
     }
 
-    override val firFileLockProvider: LockProvider<FirFile, ReentrantReadWriteLock> = LockProvider { ReentrantReadWriteLock() }
+    override val firFileLockProvider: LockProvider<FirFile> = LockProvider()
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt

@@ -12,8 +12,6 @@ import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
-import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.withReadLock
-import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.withWriteLock
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/utils.kt

@@ -32,7 +32,7 @@ internal inline fun <T : Any> executeWithoutPCE(crossinline action: () -> T): T
     return result!!
 }
 
-internal inline fun <T : Any> Lock.lockWithPCECheck(lockingIntervalMs: Long, action: () -> T): T {
+internal inline fun <T> Lock.lockWithPCECheck(lockingIntervalMs: Long, action: () -> T): T {
     var needToRun = true
     var result: T? = null
     while (needToRun) {