From a7692ba2c5f081a486212a650f0b11927c9a98ba Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Wed, 10 May 2023 10:36:28 +0200
Subject: [PATCH] [LL FIR] do not calculate lazy body recursively by default

To avoid visiting not locked declarations

^KT-56550
---
 .../lazy/resolve/FirLazyBodiesCalculator.kt   | 22 ++++++++++++++-----
 .../AbstractFirLazyBodiesCalculatorTest.kt    |  2 +-
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt
index e7762970cd7..a01b07a9e86 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt
@@ -34,14 +34,14 @@ internal object FirLazyBodiesCalculator {
     fun calculateLazyBodiesInside(designation: FirDesignation) {
         calculateAnnotations(designation.target)
         designation.target.transform<FirElement, PersistentList<FirRegularClass>>(
-            FirLazyBodiesCalculatorTransformer,
+            FirTargetLazyBodiesCalculatorTransformer,
             designation.path.toPersistentList(),
         )
     }
 
-    fun calculateLazyBodies(firFile: FirFile) {
+    fun calculateAllLazyExpressionsInFile(firFile: FirFile) {
         calculateAnnotations(firFile)
-        firFile.transform<FirElement, PersistentList<FirRegularClass>>(FirLazyBodiesCalculatorTransformer, persistentListOf())
+        firFile.transform<FirElement, PersistentList<FirRegularClass>>(FirAllLazyBodiesCalculatorTransformer, persistentListOf())
     }
 
     fun calculateAnnotations(firElement: FirElementWithResolveState) {
@@ -295,12 +295,12 @@ private object FirLazyAnnotationTransformer : FirTransformer<FirLazyAnnotationTr
     }
 }
 
-private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentList<FirRegularClass>>() {
-
+private object FirAllLazyBodiesCalculatorTransformer : FirLazyBodiesCalculatorTransformer() {
     override fun transformFile(file: FirFile, data: PersistentList<FirRegularClass>): FirFile {
         file.declarations.forEach {
             it.transformSingle(this, data)
         }
+
         return file
     }
 
@@ -312,8 +312,15 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
             }
             element.transformChildren(this, newList)
         }
+
         return element
     }
+}
+
+private object FirTargetLazyBodiesCalculatorTransformer : FirLazyBodiesCalculatorTransformer()
+
+private abstract class FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentList<FirRegularClass>>() {
+    override fun <E : FirElement> transformElement(element: E, data: PersistentList<FirRegularClass>): E = element
 
     override fun transformField(field: FirField, data: PersistentList<FirRegularClass>): FirStatement {
         if (field.initializer is FirLazyExpression) {
@@ -331,6 +338,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
             val designation = FirDesignation(data, simpleFunction)
             FirLazyBodiesCalculator.calculateLazyBodiesForFunction(designation)
         }
+
         return simpleFunction
     }
 
@@ -342,6 +350,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
             val designation = FirDesignation(data, constructor)
             FirLazyBodiesCalculator.calculateLazyBodyForConstructor(designation)
         }
+
         return constructor
     }
 
@@ -350,6 +359,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
             val designation = FirDesignation(data, property)
             FirLazyBodiesCalculator.calculateLazyBodyForProperty(designation)
         }
+
         return property
     }
 
@@ -362,6 +372,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
             val designation = FirDesignation(data, enumEntry)
             FirLazyBodiesCalculator.calculateLazyInitializerForEnumEntry(designation)
         }
+
         return enumEntry
     }
 
@@ -372,6 +383,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
             val designation = FirDesignation(data, anonymousInitializer)
             FirLazyBodiesCalculator.calculateLazyBodyForAnonymousInitializer(designation)
         }
+
         return anonymousInitializer
     }
 }
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyBodiesCalculatorTest.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyBodiesCalculatorTest.kt
index 2ee2e6326f3..8ed29ef1a9b 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyBodiesCalculatorTest.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyBodiesCalculatorTest.kt
@@ -42,7 +42,7 @@ abstract class AbstractFirLazyBodiesCalculatorTest : AbstractLowLevelApiSingleFi
                 bodyBuildingMode = BodyBuildingMode.LAZY_BODIES
             ).buildFirFile(ktFile)
 
-            FirLazyBodiesCalculator.calculateLazyBodies(laziedFirFile)
+            FirLazyBodiesCalculator.calculateAllLazyExpressionsInFile(laziedFirFile)
             laziedFirFile.accept(lazyChecker)
 
             val fullFirFile = RawFirBuilder(