FIR IDE: separate lazy phase into 3 phases, add contracts phase

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt

@@ -103,6 +103,7 @@ fun <F : FirClass<F>> F.runContractAndBodiesResolutionForLocalClass(
 fun createReturnTypeCalculatorForIDE(
     session: FirSession,
     scopeSession: ScopeSession,
+    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
     createTransformer: (
         designation: Iterator<FirElement>,
         FirSession,
@@ -112,7 +113,7 @@ fun createReturnTypeCalculatorForIDE(
         BodyResolveContext?
     ) -> FirDesignatedBodyResolveTransformerForReturnTypeCalculator
 ): ReturnTypeCalculator =
-    ReturnTypeCalculatorWithJump(session, scopeSession, ImplicitBodyResolveComputationSession(), createTransformer = createTransformer)
+    ReturnTypeCalculatorWithJump(session, scopeSession, implicitBodyResolveComputationSession, createTransformer = createTransformer)
 
 open class FirImplicitAwareBodyResolveTransformer(
     session: FirSession,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -31,7 +31,7 @@ import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.Name
 
-class FirContractResolveTransformer(
+open class FirContractResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession,
     outerBodyResolveContext: BodyResolveContext? = null

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

@@ -28,6 +28,25 @@ internal class FirPhaseRunner(val transformerProvider: FirTransformerProvider) {
         }
     }
 
+    inline fun runPhaseWithCustomResolve(phase: FirResolvePhase, crossinline resolve: () -> Unit) = when (phase) {
+        FirResolvePhase.SUPER_TYPES -> superTypesBodyResolveLock.withLock {
+            runPhaseWithCustomResolveWithoutLock(resolve)
+        }
+        FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> implicitTypesResolveLock.withLock {
+            runPhaseWithCustomResolveWithoutLock(resolve)
+        }
+        else -> {
+            runPhaseWithCustomResolveWithoutLock(resolve)
+        }
+    }
+
+    private inline fun runPhaseWithCustomResolveWithoutLock(crossinline resolve: () -> Unit) {
+        executeWithoutPCE {
+            resolve()
+        }
+    }
+
+
     private fun runPhaseWithoutLock(firFile: FirFile, phase: FirResolvePhase) {
         val phaseProcessor = transformerProvider.getTransformerForPhase(firFile.session, phase)
         executeWithoutPCE {

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
@@ -26,6 +27,7 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
     returnTypeCalculator = createReturnTypeCalculatorForIDE(
         session,
         ScopeSession(),
+        ImplicitBodyResolveComputationSession(),
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     )
 ) {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -11,11 +11,13 @@ import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirDesignatedBodyResolveTransformerForIDE
+import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedBodyResolveTransformerForIDE
 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.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedContractsResolveTransformerForIDE
+import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedImplicitTypesTransformerForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
@@ -41,9 +43,6 @@ internal class FirLazyDeclarationResolver(
 
         if (checkPCE) {
             firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) {
-                executeWithoutPCE {
-                    calculateLazyBodies(declaration)
-                }
                 runLazyResolveWithoutLock(
                     declaration,
                     moduleFileCache,
@@ -58,7 +57,6 @@ internal class FirLazyDeclarationResolver(
         } else {
             firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
                 executeWithoutPCE {
-                    calculateLazyBodies(declaration)
                     runLazyResolveWithoutLock(
                         declaration,
                         moduleFileCache,
@@ -89,7 +87,7 @@ internal class FirLazyDeclarationResolver(
         checkPCE: Boolean,
     ) {
         if (fromPhase >= toPhase) return
-        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
+        val nonLazyPhase = minOf(toPhase, LAST_NON_LAZY_PHASE)
         if (fromPhase < nonLazyPhase) {
             firFileBuilder.runResolveWithoutLock(
                 containerFirFile,
@@ -99,8 +97,25 @@ internal class FirLazyDeclarationResolver(
             )
         }
         if (toPhase <= nonLazyPhase) return
-        if (checkPCE) checkCanceled()
-        runLazyResolvePhase(firDeclarationToResolve, containerFirFile, moduleFileCache, provider, toPhase, towerDataContextCollector)
+
+        executeWithoutPCE {
+            calculateLazyBodies(firDeclarationToResolve)
+        }
+
+        var currentPhase = nonLazyPhase
+        while (currentPhase < toPhase) {
+            currentPhase = currentPhase.next
+            if (currentPhase.pluginPhase) continue
+            if (checkPCE) checkCanceled()
+            runLazyResolvePhase(
+                firDeclarationToResolve,
+                containerFirFile,
+                moduleFileCache,
+                provider,
+                currentPhase,
+                towerDataContextCollector
+            )
+        }
     }
 
     private fun runLazyResolvePhase(
@@ -108,7 +123,7 @@ internal class FirLazyDeclarationResolver(
         containerFirFile: FirFile,
         moduleFileCache: ModuleFileCache,
         provider: FirProvider,
-        toPhase: FirResolvePhase,
+        phase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector?,
     ) {
         val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
@@ -132,21 +147,36 @@ internal class FirLazyDeclarationResolver(
                 designation += nonLocalDeclarationToResolve
             }
         }
-        if (designation.all { it.resolvePhase >= toPhase }) {
+        if (designation.all { it.resolvePhase >= phase }) {
             return
         }
         val scopeSession = firFileBuilder.firPhaseRunner.transformerProvider.getScopeSession(containerFirFile.session)
-        val transformer = FirDesignatedBodyResolveTransformerForIDE(
-            designation.iterator(), containerFirFile.session,
-            scopeSession,
-            toPhase,
-            towerDataContextCollector
-        )
-        executeWithoutPCE {
+        val transformer = when (phase) {
+            FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
+                designation.iterator(),
+                containerFirFile.session,
+                scopeSession,
+            )
+            FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
+                designation.iterator(),
+                containerFirFile.session,
+                scopeSession,
+            )
+            FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
+                designation.iterator(),
+                containerFirFile.session,
+                scopeSession,
+                towerDataContextCollector
+            )
+            else -> error("Non-lazy phase $phase")
+        }
+
+        firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(phase) {
             containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
         }
     }
 
+
     private fun FirDeclaration.getNonLocalDeclarationToResolve(provider: FirProvider, moduleFileCache: ModuleFileCache): FirDeclaration {
         if (this is FirFile) return this
         val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration")
@@ -155,6 +185,10 @@ internal class FirLazyDeclarationResolver(
             ?: error("Container for local declaration cannot be null")
         return nonLocalPsi.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache)
     }
+
+    companion object {
+        private val LAST_NON_LAZY_PHASE = FirResolvePhase.STATUS
+    }
 }
 
 

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

@@ -3,7 +3,7 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
+package org.jetbrains.kotlin.idea.fir.low.level.api.trasformers
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -13,32 +13,31 @@ import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 
 internal class FirDesignatedBodyResolveTransformerForIDE(
     private val designation: Iterator<FirElement>,
     session: FirSession,
     scopeSession: ScopeSession,
-    phase: FirResolvePhase,
     private val towerDataContextCollector: FirTowerDataContextCollector? = null
 ) : FirBodyResolveTransformer(
     session,
-    phase = phase,
-    implicitTypeOnly = phase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+    phase = FirResolvePhase.BODY_RESOLVE,
+    implicitTypeOnly = false,
     scopeSession = scopeSession,
     returnTypeCalculator = createReturnTypeCalculatorForIDE(
         session,
         scopeSession,
+        ImplicitBodyResolveComputationSession(),
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     )
 ) {
 
-    override fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) {
-        towerDataContextCollector?.addDeclarationContext(declaration, context.towerDataContext)
-    }
-
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
         if (designation.hasNext()) {
             designation.next().visitNoTransform(this, data)
@@ -48,6 +47,10 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         return super.transformDeclarationContent(declaration, data)
     }
 
+    override fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) {
+        towerDataContextCollector?.addDeclarationContext(declaration, context.towerDataContext)
+    }
+
     override fun onBeforeStatementResolution(statement: FirStatement) {
         towerDataContextCollector?.addStatementContext(statement, context.towerDataContext)
     }

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.trasformers
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveTransformer
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.compose
+
+internal class FirDesignatedContractsResolveTransformerForIDE(
+    private val designation: Iterator<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession,
+) : FirContractResolveTransformer(session, scopeSession) {
+
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+        if (designation.hasNext()) {
+            designation.next().visitNoTransform(this, data)
+            return declaration.compose()
+        }
+        return super.transformDeclarationContent(declaration, data)
+    }
+}

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.trasformers
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitAwareBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+
+internal class FirDesignatedImplicitTypesTransformerForIDE(
+    private val designation: Iterator<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession,
+    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession(),
+) : FirImplicitAwareBodyResolveTransformer(
+    session,
+    implicitBodyResolveComputationSession = implicitBodyResolveComputationSession,
+    phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+    implicitTypeOnly = true,
+    scopeSession = scopeSession,
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(
+        session,
+        scopeSession,
+        implicitBodyResolveComputationSession,
+        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+    )
+) {
+
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+        if (designation.hasNext()) {
+            designation.next().visitNoTransform(this, data)
+            return declaration.compose()
+        }
+        return super.transformDeclarationContent(declaration, data)
+    }
+}