[FIR] move lazy expression assertions to more common class

^KT-56543

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

@@ -9,6 +9,9 @@ import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
@@ -50,6 +53,21 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         }
     }
 
+    override fun transformLazyExpression(lazyExpression: FirLazyExpression, data: ResolutionMode): FirStatement {
+        suppressOrThrowError("FirLazyExpression should be calculated before accessing")
+        return lazyExpression
+    }
+
+    override fun transformLazyBlock(lazyBlock: FirLazyBlock, data: ResolutionMode): FirStatement {
+        suppressOrThrowError("FirLazyBlock should be calculated before accessing")
+        return lazyBlock
+    }
+
+    private fun suppressOrThrowError(message: String) {
+        if (System.getProperty("kotlin.suppress.lazy.expression.access").toBoolean()) return
+        error(message)
+    }
+
     protected inline val localScopes: List<FirLocalScope> get() = components.localScopes
 
     protected inline val noExpectedType: FirTypeRef get() = components.noExpectedType

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

@@ -7,9 +7,6 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.PrivateForInline
-import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
-import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 
@@ -38,12 +35,4 @@ abstract class FirPartialBodyResolveTransformer(
     override fun <E : FirElement> transformElement(element: E, data: ResolutionMode): E {
         return element.transform(transformer, data)
     }
-
-    override fun transformLazyExpression(lazyExpression: FirLazyExpression, data: ResolutionMode): FirStatement {
-        error("FirLazyExpression should be calculated before accessing")
-    }
-
-    override fun transformLazyBlock(lazyBlock: FirLazyBlock, data: ResolutionMode): FirStatement {
-        error("FirLazyBlock should be calculated before accessing")
-    }
 }

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractLazyBodyIsNotTouchedTest.kt

@@ -17,6 +17,7 @@ abstract class AbstractLazyBodyIsNotTouchedTest : AbstractFirBaseDiagnosticsTest
     override fun runAnalysis(testDataFile: File, testFiles: List<TestFile>, firFilesPerSession: Map<FirSession, List<FirFile>>) {
         val phases = FirResolvePhase.values()
             .dropWhile { it <= FirResolvePhase.RAW_FIR }
+            .filterNot { it == FirResolvePhase.COMPILER_REQUIRED_ANNOTATIONS }
             .takeWhile { it < FirResolvePhase.ARGUMENTS_OF_ANNOTATIONS }
 
         for ((session, firFiles) in firFilesPerSession) {