[LL FIR] mark entire delegate as lazy expression

To simplify modification and restoring

^KT-59687

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

@@ -139,9 +139,9 @@ private fun replaceLazyInitializer(target: FirVariable, copy: FirVariable) {
     }
 }
 
-private fun replaceLazyExpression(target: FirWrappedExpression, copy: FirWrappedExpression) {
-    if (target.expression is FirLazyExpression) {
-        target.replaceExpression(copy.expression)
+private fun replaceLazyDelegate(target: FirVariable, copy: FirVariable) {
+    if (target.delegate is FirLazyExpression) {
+        target.replaceDelegate(copy.delegate)
     }
 }
 
@@ -187,17 +187,12 @@ private fun calculateLazyBodyForProperty(designation: FirDesignation) {
     }
 
     replaceLazyInitializer(firProperty, newProperty)
+    replaceLazyDelegate(firProperty, newProperty)
 
     firProperty.getExplicitBackingField()?.let { backingField ->
         val newBackingField = newProperty.getExplicitBackingField()!!
         replaceLazyInitializer(backingField, newBackingField)
     }
-
-    (firProperty.delegate as? FirWrappedDelegateExpression)?.let { delegate ->
-        val newDelegate = newProperty.delegate as FirWrappedDelegateExpression
-        replaceLazyExpression(delegate, newDelegate)
-        delegate.replaceDelegateProvider(newDelegate.delegateProvider)
-    }
 }
 
 private fun calculateLazyInitializerForEnumEntry(designation: FirDesignation) {
@@ -258,7 +253,7 @@ private fun needCalculatingLazyBodyForProperty(firProperty: FirProperty): Boolea
     firProperty.getter?.let { needCalculatingLazyBodyForFunction(it) } == true
             || firProperty.setter?.let { needCalculatingLazyBodyForFunction(it) } == true
             || firProperty.initializer is FirLazyExpression
-            || (firProperty.delegate as? FirWrappedDelegateExpression)?.expression is FirLazyExpression
+            || firProperty.delegate is FirLazyExpression
             || firProperty.getExplicitBackingField()?.initializer is FirLazyExpression
 
 private enum class FirLazyAnnotationTransformerScope {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -223,6 +223,7 @@ internal object BodyStateKeepers {
 
         if (!isCallableWithSpecialBody(variable)) {
             add(FirVariable::initializerIfUnresolved, FirVariable::replaceInitializer, ::expressionGuard)
+            add(FirVariable::delegateIfUnresolved, FirVariable::replaceDelegate, ::expressionGuard)
         }
     }
 
@@ -253,11 +254,6 @@ internal object BodyStateKeepers {
         entity(property.setterIfUnresolved, FUNCTION)
         entity(property.backingFieldIfUnresolved, VARIABLE)
 
-        entity(property.delegateIfUnresolved) {
-            add(FirWrappedDelegateExpression::expression, FirWrappedDelegateExpression::replaceExpression, ::expressionGuard)
-            add(FirWrappedDelegateExpression::delegateProvider, FirWrappedDelegateExpression::replaceDelegateProvider, ::expressionGuard)
-        }
-
         add(FirProperty::controlFlowGraphReference, FirProperty::replaceControlFlowGraphReference)
     }
 }
@@ -324,6 +320,12 @@ private val FirVariable.initializerIfUnresolved: FirExpression?
         else -> initializer
     }
 
+private val FirVariable.delegateIfUnresolved: FirExpression?
+    get() = when (this) {
+        is FirProperty -> if (bodyResolveState < FirPropertyBodyResolveState.EVERYTHING_RESOLVED) delegate else null
+        else -> delegate
+    }
+
 private val FirProperty.backingFieldIfUnresolved: FirBackingField?
     get() = if (bodyResolveState < FirPropertyBodyResolveState.INITIALIZER_RESOLVED) getExplicitBackingField() else null
 
@@ -333,9 +335,6 @@ private val FirProperty.getterIfUnresolved: FirPropertyAccessor?
 private val FirProperty.setterIfUnresolved: FirPropertyAccessor?
     get() = if (bodyResolveState < FirPropertyBodyResolveState.EVERYTHING_RESOLVED) setter else null
 
-private val FirProperty.delegateIfUnresolved: FirWrappedDelegateExpression?
-    get() = if (bodyResolveState < FirPropertyBodyResolveState.EVERYTHING_RESOLVED) delegate as? FirWrappedDelegateExpression else null
-
 private fun delegatedConstructorCallGuard(fir: FirDelegatedConstructorCall): FirDelegatedConstructorCall {
     if (fir is FirLazyDelegatedConstructorCall) {
         return fir

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt

@@ -1908,12 +1908,18 @@ open class PsiRawFirBuilder(
                                 expression = delegateExpression
                             }
 
+                            val lazyDelegateExpression: FirLazyExpression? = buildOrLazy(
+                                build = { null },
+                                lazy = { buildLazyExpression { source = delegateBuilder.source } },
+                            )
+
                             generateAccessorsByDelegate(
                                 delegateBuilder,
                                 baseModuleData,
                                 ownerRegularOrAnonymousObjectSymbol,
                                 context,
                                 isExtension = receiverTypeReference != null,
+                                lazyDelegateExpression = lazyDelegateExpression,
                             )
                         }
                     }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.contracts.FirLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.contracts.builder.buildLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirReceiverParameter
 import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.declarations.builder.*
@@ -313,6 +314,7 @@ fun <T> FirPropertyBuilder.generateAccessorsByDelegate(
     ownerRegularOrAnonymousObjectSymbol: FirClassSymbol<*>?,
     context: Context<T>,
     isExtension: Boolean,
+    lazyDelegateExpression: FirLazyExpression? = null,
 ) {
     if (delegateBuilder == null) return
     val delegateFieldSymbol = FirDelegateFieldSymbol(symbol.callableId).also {
@@ -402,17 +404,21 @@ fun <T> FirPropertyBuilder.generateAccessorsByDelegate(
         }
     }
 
-    delegateBuilder.delegateProvider = buildFunctionCall {
-        explicitReceiver = delegateBuilder.expression
-        calleeReference = buildSimpleNamedReference {
+    delegate = lazyDelegateExpression ?: run {
+        delegateBuilder.delegateProvider = buildFunctionCall {
+            explicitReceiver = delegateBuilder.expression
+            calleeReference = buildSimpleNamedReference {
+                source = fakeSource
+                name = OperatorNameConventions.PROVIDE_DELEGATE
+            }
+            argumentList = buildBinaryArgumentList(thisRef(forDispatchReceiver = true), propertyRef())
+            origin = FirFunctionCallOrigin.Operator
             source = fakeSource
-            name = OperatorNameConventions.PROVIDE_DELEGATE
         }
-        argumentList = buildBinaryArgumentList(thisRef(forDispatchReceiver = true), propertyRef())
-        origin = FirFunctionCallOrigin.Operator
-        source = fakeSource
+
+        delegateBuilder.build()
     }
-    delegate = delegateBuilder.build()
+
     if (getter == null || getter is FirDefaultPropertyAccessor) {
         val annotations = getter?.annotations
         val returnTarget = FirFunctionTarget(null, isLambda = false)