FIR IDE: allow PCE between resolve phases, forbid in other cases

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

@@ -83,7 +83,7 @@ internal open class FirModuleResolveStateImpl(
     }
 
     override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
-        elementBuilder.lazyResolveDeclaration(declaration, fileCache, toPhase)
+        elementBuilder.lazyResolveDeclarationWithPCECheck(declaration, fileCache, toPhase)
         return declaration
     }
 

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.createTransformerBasedProcessorByPhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
 import java.util.concurrent.locks.ReentrantLock
 import kotlin.concurrent.withLock
 
@@ -30,6 +31,6 @@ internal class FirPhaseRunner {
 
     private fun runPhaseWithoutLock(firFile: FirFile, phase: FirResolvePhase, scopeSession: ScopeSession) {
         val phaseProcessor = phase.createTransformerBasedProcessorByPhase(firFile.session, scopeSession)
-        phaseProcessor.processFile(firFile)
+        executeWithoutPCE { phaseProcessor.processFile(firFile) }
     }
 }

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

@@ -16,6 +16,6 @@ internal class IdePhasedFirFileResolver(
     private val cache: ModuleFileCache
 ) : PhasedFirFileResolver() {
     override fun resolveFile(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
-        firFileBuilder.runResolve(firFile, cache, fromPhase, toPhase)
+        firFileBuilder.runResolveWithLock(firFile, cache, fromPhase, toPhase)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt

@@ -5,8 +5,6 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 
-import com.intellij.psi.util.parentOfType
-import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirStatement
@@ -15,7 +13,6 @@ import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -28,7 +25,6 @@ import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
 import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
-import org.jetbrains.kotlin.psi.psiUtil.parents
 
 /**
  * Maps [KtElement] to [FirElement]
@@ -56,47 +52,33 @@ internal class FirElementBuilder(
             else -> error("Unsupported: ${container.text}")
         }
 
-        firFileBuilder.runCustomResolve(firFile, moduleFileCache) {
+        runLazyResolveWithPCECheck(containerFir, moduleFileCache, firFile, firFile.session.firIdeProvider, toPhase)
-            runLazyResolveWithoutLock(containerFir, firFile, firFile.session.firIdeProvider, toPhase)
-        }
 
         return psiToFirCache.getFir(element, containerFir, firFile)
     }
 
-    fun lazyResolveDeclaration(
+    fun lazyResolveDeclarationWithPCECheck(
         declaration: FirDeclaration,
         moduleFileCache: ModuleFileCache,
         toPhase: FirResolvePhase
     ) {
         if (declaration.resolvePhase < toPhase) {
             val firFile = declaration.getContainingFile() ?: error("FirFile was not found for\n${declaration.render()}")
-            firFileBuilder.runCustomResolve(firFile, moduleFileCache) {
+            runLazyResolveWithPCECheck(declaration, moduleFileCache, firFile, firFile.session.firIdeProvider, toPhase)
-                runLazyResolveWithoutLock(declaration, firFile, declaration.session.firProvider, toPhase)
-            }
         }
     }
 
-    /**
+    private fun runLazyResolvePhase(
-     * Should be invoked under lock
-     */
-    private fun runLazyResolveWithoutLock(
         firDeclarationToResolve: FirDeclaration,
         containerFirFile: FirFile,
         provider: FirProvider,
         toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>?,
     ) {
-        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
-        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
-            firFileBuilder.runResolveWithoutLock(containerFirFile, fromPhase = firDeclarationToResolve.resolvePhase, toPhase = nonLazyPhase)
-        }
-        if (toPhase <= nonLazyPhase) return
-
         val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider)
 
         val designation = mutableListOf<FirDeclaration>(containerFirFile)
         if (nonLocalDeclarationToResolve !is FirFile) {
-
             val id = when (nonLocalDeclarationToResolve) {
                 is FirCallableDeclaration<*> -> {
                     nonLocalDeclarationToResolve.symbol.callableId.classId
@@ -127,6 +109,51 @@ internal class FirElementBuilder(
         containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
     }
 
+    /**
+     * Should be invoked under lock
+     */
+    private fun runLazyResolveWithoutLock(
+        firDeclarationToResolve: FirDeclaration,
+        containerFirFile: FirFile,
+        provider: FirProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+    ) {
+        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
+        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
+            firFileBuilder.runResolveWithoutLockNoPCECheck(
+                containerFirFile,
+                fromPhase = firDeclarationToResolve.resolvePhase,
+                toPhase = nonLazyPhase
+            )
+        }
+        if (toPhase <= nonLazyPhase) return
+        runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement)
+    }
+
+    fun runLazyResolveWithPCECheck(
+        firDeclarationToResolve: FirDeclaration,
+        cache: ModuleFileCache,
+        containerFirFile: FirFile,
+        provider: FirProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+    ) {
+        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
+        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
+            firFileBuilder.runResolveWithPCECheck(
+                containerFirFile,
+                fromPhase = firDeclarationToResolve.resolvePhase,
+                toPhase = nonLazyPhase,
+                cache = cache,
+            )
+        }
+        if (toPhase <= nonLazyPhase) return
+        firFileBuilder.runCustomResolveWithPCECheck(containerFirFile, cache) {
+            runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement)
+        }
+    }
+
     //TODO for completion only
     fun runLazyResolveForCompletion(
         firFunction: FirFunction<*>,

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

@@ -13,6 +13,8 @@ import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.lockWithPCECheck
 
 /**
  * Responsible for building [FirFile] by [KtFile]
@@ -43,7 +45,7 @@ internal class FirFileBuilder(
         if (toPhase > FirResolvePhase.RAW_FIR) {
             cache.firFileLockProvider.withLock(firFile) {
                 //add lock for implit type resolve phase & super type
-                runResolveWithoutLock(firFile, fromPhase = firFile.resolvePhase, toPhase = toPhase)
+                runResolveWithoutLockNoPCECheck(firFile, fromPhase = firFile.resolvePhase, toPhase = toPhase)
             }
         }
         return firFile
@@ -52,16 +54,35 @@ internal class FirFileBuilder(
     /**
      * Runs [resolve] function (which is considered to do some resolve on [firFile]) under a lock for [firFile]
      */
-    inline fun <R> runCustomResolve(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R =
+    inline fun <R> runCustomResolveUnderLock(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R =
         cache.firFileLockProvider.withLock(firFile) { resolve() }
 
-    fun runResolve(firFile: FirFile, cache: ModuleFileCache, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+    inline fun <R : Any> runCustomResolveWithPCECheck(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R {
+        val lock = cache.firFileLockProvider.getLockFor(firFile)
+        return lock.lockWithPCECheck(LOCKING_INTERVAL_MS) { resolve() }
+    }
+
+    fun runResolveWithLock(firFile: FirFile, cache: ModuleFileCache, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
         cache.firFileLockProvider.withLock(firFile) {
-            runResolveWithoutLock(firFile, fromPhase, toPhase)
+            runResolveWithoutLockNoPCECheck(firFile, fromPhase, toPhase)
         }
     }
 
-    fun runResolveWithoutLock(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+    fun runResolveWithPCECheck(firFile: FirFile, cache: ModuleFileCache, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+        val lock = cache.firFileLockProvider.getLockFor(firFile)
+        lock.lockWithPCECheck(LOCKING_INTERVAL_MS) {
+            val scopeSession = ScopeSession()
+            var currentPhase = fromPhase
+            while (currentPhase < toPhase) {
+                checkCanceled()
+                currentPhase = currentPhase.next
+                firPhaseRunner.runPhase(firFile, currentPhase, scopeSession)
+            }
+        }
+    }
+
+
+    fun runResolveWithoutLockNoPCECheck(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
         assert(fromPhase <= toPhase) {
             "Trying to resolve file ${firFile.name} from $fromPhase to $toPhase"
         }
@@ -72,6 +93,10 @@ internal class FirFileBuilder(
             firPhaseRunner.runPhase(firFile, currentPhase, scopeSession)
         }
     }
+
+    companion object {
+        private const val LOCKING_INTERVAL_MS = 500L
+    }
 }
 
 

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

@@ -6,9 +6,12 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.util
 
 import com.intellij.openapi.progress.ProcessCanceledException
+import com.intellij.openapi.progress.ProgressManager
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.locks.ReentrantLock
 
 
 internal inline fun <T> executeOrReturnDefaultValueOnPCE(defaultValue: T, action: () -> T): T =
@@ -18,6 +21,34 @@ internal inline fun <T> executeOrReturnDefaultValueOnPCE(defaultValue: T, action
         defaultValue
     }
 
+internal inline fun <T : Any> executeWithoutPCE(crossinline action: () -> T): T {
+    var result: T? = null
+    ProgressManager.getInstance().executeNonCancelableSection { result = action() }
+    return result!!
+}
+
+internal inline fun <T : Any> ReentrantLock.lockWithPCECheck(lockingIntervalMs: Long, action: () -> T): T {
+    var needToRun = true
+    var result: T? = null
+    while (needToRun) {
+        checkCanceled()
+        if (tryLock(lockingIntervalMs, TimeUnit.MILLISECONDS)) {
+            try {
+                needToRun = false
+                result = action()
+            } finally {
+                unlock()
+            }
+        }
+    }
+    return result!!
+}
+
+@Suppress("NOTHING_TO_INLINE")
+internal inline fun checkCanceled() {
+    ProgressManager.checkCanceled()
+}
+
 internal val FirElement.isErrorElement
     get() = this is FirDiagnosticHolder