[FIR] FirResolvePhase: introduce isItAllowedToCallLazyResolveTo

To explicitly declare when it is possible to call lazy resolve

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt

@@ -200,7 +200,7 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
         actionUnderLock: () -> Unit,
         actionOnCycle: () -> Unit,
     ) {
-        checker.lazyResolveToPhaseInside(phase, isJumpingPhase = true) {
+        checker.lazyResolveToPhaseInside(phase) {
             target.withJumpingLockImpl(phase, actionUnderLock, actionOnCycle)
         }
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirLazyResolveContractChecker.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
 
 import com.intellij.openapi.diagnostic.Logger
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.isItAllowedToCallLazyResolveTo
 import org.jetbrains.kotlin.fir.symbols.FirLazyResolveContractViolationException
 
 /**
@@ -18,8 +19,9 @@ import org.jetbrains.kotlin.fir.symbols.FirLazyResolveContractViolationException
 internal class LLFirLazyResolveContractChecker {
     private val currentTransformerPhase = ThreadLocal.withInitial<FirResolvePhase?> { null }
 
-    inline fun lazyResolveToPhaseInside(phase: FirResolvePhase, isJumpingPhase: Boolean = false, resolve: () -> Unit) {
+    inline fun lazyResolveToPhaseInside(phase: FirResolvePhase, resolve: () -> Unit) {
-        checkIfCanLazyResolveToPhase(phase, isJumpingPhase)
+        checkIfCanLazyResolveToPhase(phase)
+
         val previousPhase = currentTransformerPhase.get()
         currentTransformerPhase.set(phase)
         try {
@@ -29,10 +31,10 @@ internal class LLFirLazyResolveContractChecker {
         }
     }
 
-    internal fun checkIfCanLazyResolveToPhase(requestedPhase: FirResolvePhase, isJumpingPhase: Boolean) {
+    private fun checkIfCanLazyResolveToPhase(requestedPhase: FirResolvePhase) {
         val currentPhase = currentTransformerPhase.get() ?: return
 
-        if (requestedPhase > currentPhase || !isJumpingPhase && requestedPhase == currentPhase) {
+        if (!currentPhase.isItAllowedToCallLazyResolveTo(requestedPhase)) {
             val exception = FirLazyResolveContractViolationException(currentPhase = currentPhase, requestedPhase = requestedPhase)
             if (System.getProperty("kotlin.suppress.lazy.resolve.contract.violation") != null) {
                 LoggerHolder.LOG.warn(exception)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirResolvePhase.kt

@@ -46,6 +46,7 @@ package org.jetbrains.kotlin.fir.declarations
  * 2. The compiler **must not request and cannot rely on any information from the higher phases**.
  *   For example, during the [STATUS] phase, we cannot access information regarding implicit types
  *   as they will only be calculated during the [IMPLICIT_TYPES_BODY_RESOLVE] phase.
+ *   See [isItAllowedToCallLazyResolveTo] as a reference.
  *
  * 3. The compiler **can request and rely on the information from the current phase only during *jumping phases***.
  *   For example, during the [TYPES] phase,
@@ -227,3 +228,28 @@ val FirResolvePhase.isBodyResolve: Boolean
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> true
         else -> false
     }
+
+/**
+ * See [FirResolvePhase] KDoc for more details about resolution contacts.
+ *
+ * @param this The current phase
+ * @param requestedPhase The requested phase
+ *
+ * @see FirResolvePhase
+ * @see org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
+ * @see org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
+ */
+fun FirResolvePhase.isItAllowedToCallLazyResolveTo(requestedPhase: FirResolvePhase): Boolean = when {
+    // It is fine to call lazy resolution for all phases less than our
+    this > requestedPhase -> true
+
+    // It is legal only in specific cases
+    this == requestedPhase -> when (requestedPhase) {
+        // The resolver can jump into Java initializer during this phase,
+        // which can jump into the Kotlin world again, and we cannot provide the initial context
+        FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> true
+        else -> false
+    }
+
+    else -> false
+}

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirCompilerLazyDeclarationResolverWithPhaseChecking.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 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.
  */
 
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.test.frontend.fir.handlers
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.isItAllowedToCallLazyResolveTo
 import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.symbols.FirLazyResolveContractViolationException
 
@@ -53,7 +54,7 @@ class FirCompilerLazyDeclarationResolverWithPhaseChecking : FirLazyDeclarationRe
         // However, now we keep more strict invariant here,
         // because we don't want to hide some cases when transformers violate phase contract directly
         // but due to usage of already resolved stdlib classes we don't see it
-        if (requestedPhase >= currentPhase) {
+        if (!currentPhase.isItAllowedToCallLazyResolveTo(requestedPhase)) {
             exceptions += FirLazyResolveContractViolationException(
                 currentPhase = currentPhase,
                 requestedPhase = requestedPhase,