From 22b5d940f14f003153f05d7634e4e25efbec26b8 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Wed, 5 Apr 2023 15:01:02 +0200
Subject: [PATCH] [FIR] move lazy expression assertions to more common class

^KT-56543
---
 .../FirAbstractBodyResolveTransformer.kt       | 18 ++++++++++++++++++
 .../FirPartialBodyResolveTransformer.kt        | 11 -----------
 .../fir/AbstractLazyBodyIsNotTouchedTest.kt    |  1 +
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt
index 1bec1cf736b..57a19ddd1cc 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt
+++ b/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
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirPartialBodyResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirPartialBodyResolveTransformer.kt
index dfaf866ada7..54c0721b622 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirPartialBodyResolveTransformer.kt
+++ b/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")
-    }
 }
diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractLazyBodyIsNotTouchedTest.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractLazyBodyIsNotTouchedTest.kt
index 195aa6319fc..92dbcc2495e 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractLazyBodyIsNotTouchedTest.kt
+++ b/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) {