[LL API] Replace whole-session invalidation with granular FIR guards

^KT-59297 Fixed

'ProcessCancelledException's are commonly thrown during code analysis,
yet not every place is ready for them. As a FIR tree is a mutable data
structure with no transaction support, sporadic halts during its
mutation might easily leave it in an inconsistent state.

Before this change, whole module sessions (with all dependents) were
invalidated if an exception occurred during their analysis. Not only it
was very ineffective, it also resulted into a concurrency problem.

'ProcessCancelledException's are thrown when an underlying progress
indicator is invalidated. As analysis in each thread might have its own
progress indicator, invalidation in one thread should not stop
analysis in others. However, it is impossible in the current state,
as all threads share the same FIR tree structure.

The change introduces 'FIR guards' – a set of instructions defining
preservation and restoration rules for individual state items. As each
resolution phase is supposed to modify only a (sort of) known subset
of tree elements, restoration of potentially changed state on a failure
effectively returns the tree consistency.

Note that guards are not defined for certain phase transformers. The
reason is that the transformers already update the tree content safely.
E.g., the type reference is updated once it is resolved, and already
resolved type references are either resolved the same way again, or
just skipped on consequent phase runs.

analysis/low-level-api-fir/build.gradle.kts

@@ -65,6 +65,12 @@ sourceSets {
     "test" { projectDefault() }
 }
 
+kotlin {
+    compilerOptions {
+        freeCompilerArgs.add("-Xcontext-receivers")
+    }
+}
+
 projectTest(jUnitMode = JUnitMode.JUnit5) {
     dependsOn(":dist")
     workingDir = rootDir

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

@@ -5,12 +5,15 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
 
+import com.intellij.psi.PsiElement
 import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.toPersistentList
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.builder.PsiRawFirBuilder
+import org.jetbrains.kotlin.fir.contracts.FirRawContractDescription
+import org.jetbrains.kotlin.fir.contracts.impl.FirEmptyContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.getExplicitBackingField
 import org.jetbrains.kotlin.fir.expressions.*
@@ -18,16 +21,11 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirLazyDelegatedConstructorCall
 import org.jetbrains.kotlin.fir.extensions.registeredPluginAnnotations
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.CompilerRequiredAnnotationsHelper
 import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.FirUserTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.fir.types.customAnnotations
-import org.jetbrains.kotlin.fir.types.forEachType
 
 internal object FirLazyBodiesCalculator {
     fun calculateBodies(designation: FirDesignation) {
@@ -70,89 +68,134 @@ internal object FirLazyBodiesCalculator {
         firAnnotationCall.argumentList.arguments.any { it is FirLazyExpression }
 }
 
-private fun replaceValueParameterDefaultValues(valueParameters: List<FirValueParameter>, newValueParameters: List<FirValueParameter>) {
-    require(valueParameters.size == newValueParameters.size)
-    for ((valueParameter, newValueParameter) in valueParameters.zip(newValueParameters)) {
-        if (newValueParameter.defaultValue != null) {
+private inline fun <reified T : FirDeclaration> revive(
+    designation: FirDesignation,
+    psiFactory: (FirDesignation) -> PsiElement? = { it.target.psi }
+): T {
+    val session = designation.target.moduleData.session
+
+    return RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
+        session = session,
+        scopeProvider = session.kotlinScopeProvider,
+        designation = designation,
+        rootNonLocalDeclaration = psiFactory(designation) as KtDeclaration,
+    ) as T
+}
+
+private fun replaceLazyValueParameters(target: FirFunction, copy: FirFunction) {
+    val targetParameters = target.valueParameters
+    val copyParameters = copy.valueParameters
+    require(targetParameters.size == copyParameters.size)
+
+    for ((valueParameter, newValueParameter) in targetParameters.zip(copyParameters)) {
+        if (valueParameter.defaultValue is FirLazyExpression) {
             valueParameter.replaceDefaultValue(newValueParameter.defaultValue)
         }
     }
 }
 
+private fun replaceLazyBody(target: FirFunction, copy: FirFunction) {
+    if (target.body is FirLazyBlock) {
+        target.replaceBody(copy.body)
+    }
+}
+
+private fun replaceLazyContractDescription(target: FirContractDescriptionOwner, copy: FirContractDescriptionOwner) {
+    val shouldReplace = when (val currentContractDescription = target.contractDescription) {
+        is FirRawContractDescription -> currentContractDescription.rawEffects.any { it is FirLazyExpression }
+        is FirEmptyContractDescription -> copy.contractDescription !is FirEmptyContractDescription
+        else -> false
+    }
+
+    if (shouldReplace) {
+        target.replaceContractDescription(copy.contractDescription)
+    }
+}
+
+private fun replaceLazyDelegatedConstructor(target: FirConstructor, copy: FirConstructor) {
+    val targetCall = target.delegatedConstructor
+    val copyCall = copy.delegatedConstructor
+
+    when (targetCall) {
+        is FirLazyDelegatedConstructorCall -> {
+            require(copyCall !is FirMultiDelegatedConstructorCall)
+            target.replaceDelegatedConstructor(copyCall)
+        }
+        is FirMultiDelegatedConstructorCall -> {
+            require(copyCall is FirMultiDelegatedConstructorCall)
+            require(targetCall.delegatedConstructorCalls.size == copyCall.delegatedConstructorCalls.size)
+
+            val newCalls = targetCall.delegatedConstructorCalls.zip(copyCall.delegatedConstructorCalls)
+                .map { (target, copy) -> target.takeUnless { it is FirLazyDelegatedConstructorCall } ?: copy }
+
+            targetCall.replaceDelegatedConstructorCalls(newCalls)
+        }
+    }
+}
+
+private fun replaceLazyInitializer(target: FirVariable, copy: FirVariable) {
+    if (target.initializer is FirLazyExpression) {
+        target.replaceInitializer(copy.initializer)
+    }
+}
+
+private fun replaceLazyExpression(target: FirWrappedExpression, copy: FirWrappedExpression) {
+    if (target.expression is FirLazyExpression) {
+        target.replaceExpression(copy.expression)
+    }
+}
+
 private fun calculateLazyBodiesForFunction(designation: FirDesignation) {
     val simpleFunction = designation.target as FirSimpleFunction
     require(needCalculatingLazyBodyForFunction(simpleFunction))
-    val newFunction = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
-        session = simpleFunction.moduleData.session,
-        scopeProvider = simpleFunction.moduleData.session.kotlinScopeProvider,
-        designation = designation,
-        rootNonLocalDeclaration = simpleFunction.psi as KtNamedFunction,
-    ) as FirSimpleFunction
 
-    simpleFunction.apply {
-        replaceBody(newFunction.body)
-        replaceContractDescription(newFunction.contractDescription)
-        replaceValueParameterDefaultValues(valueParameters, newFunction.valueParameters)
-    }
+    val newSimpleFunction = revive<FirSimpleFunction>(designation)
+
+    replaceLazyContractDescription(simpleFunction, newSimpleFunction)
+    replaceLazyBody(simpleFunction, newSimpleFunction)
+    replaceLazyValueParameters(simpleFunction, newSimpleFunction)
 }
 
 private fun calculateLazyBodyForConstructor(designation: FirDesignation) {
     val constructor = designation.target as FirConstructor
     require(needCalculatingLazyBodyForConstructor(constructor))
 
-    val newConstructor = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
-        session = constructor.moduleData.session,
-        scopeProvider = constructor.moduleData.session.kotlinScopeProvider,
-        designation = designation,
-        rootNonLocalDeclaration = constructor.psi as KtDeclaration,
-    ) as FirConstructor
+    val newConstructor = revive<FirConstructor>(designation)
 
-    constructor.apply {
-        replaceBody(newConstructor.body)
-        replaceContractDescription(newConstructor.contractDescription)
-        replaceDelegatedConstructor(newConstructor.delegatedConstructor)
-        replaceValueParameterDefaultValues(valueParameters, newConstructor.valueParameters)
-    }
+    replaceLazyContractDescription(constructor, newConstructor)
+    replaceLazyBody(constructor, newConstructor)
+    replaceLazyDelegatedConstructor(constructor, newConstructor)
+    replaceLazyValueParameters(constructor, newConstructor)
 }
 
 private fun calculateLazyBodyForProperty(designation: FirDesignation) {
     val firProperty = designation.target as FirProperty
     if (!needCalculatingLazyBodyForProperty(firProperty)) return
 
-    val newProperty = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
-        session = firProperty.moduleData.session,
-        scopeProvider = firProperty.moduleData.session.kotlinScopeProvider,
-        designation = designation,
-        rootNonLocalDeclaration = firProperty.psi as KtProperty
-    ) as FirProperty
+    val newProperty = revive<FirProperty>(designation)
 
-    firProperty.getter?.takeIf { it.body is FirLazyBlock }?.let { getter ->
+    firProperty.getter?.let { getter ->
         val newGetter = newProperty.getter!!
-        getter.replaceBody(newGetter.body)
-        getter.replaceContractDescription(newGetter.contractDescription)
+        replaceLazyContractDescription(getter, newGetter)
+        replaceLazyBody(getter, newGetter)
     }
 
-    firProperty.setter?.takeIf { it.body is FirLazyBlock }?.let { setter ->
+    firProperty.setter?.let { setter ->
         val newSetter = newProperty.setter!!
-        setter.replaceBody(newSetter.body)
-        setter.replaceContractDescription(newSetter.contractDescription)
+        replaceLazyContractDescription(setter, newSetter)
+        replaceLazyBody(setter, newSetter)
     }
 
-    if (firProperty.initializer is FirLazyExpression) {
-        firProperty.replaceInitializer(newProperty.initializer)
+    replaceLazyInitializer(firProperty, newProperty)
+
+    firProperty.getExplicitBackingField()?.let { backingField ->
+        val newBackingField = newProperty.getExplicitBackingField()!!
+        replaceLazyInitializer(backingField, newBackingField)
     }
 
-    firProperty.getExplicitBackingField()?.takeIf { it.initializer is FirLazyExpression }?.let { backingField ->
-        val newInitializer = newProperty.getExplicitBackingField()?.initializer
-        backingField.replaceInitializer(newInitializer)
-    }
-
-    val delegate = firProperty.delegate as? FirWrappedDelegateExpression
-    val delegateExpression = delegate?.expression
-    if (delegateExpression is FirLazyExpression) {
-        val newDelegate = newProperty.delegate as? FirWrappedDelegateExpression
-        check(newDelegate != null) { "Invalid replacement delegate" }
-        delegate.replaceExpression(newDelegate.expression)
+    (firProperty.delegate as? FirWrappedDelegateExpression)?.let { delegate ->
+        val newDelegate = newProperty.delegate as FirWrappedDelegateExpression
+        replaceLazyExpression(delegate, newDelegate)
         delegate.replaceDelegateProvider(newDelegate.delegateProvider)
     }
 }
@@ -161,41 +204,25 @@ private fun calculateLazyInitializerForEnumEntry(designation: FirDesignation) {
     val enumEntry = designation.target as FirEnumEntry
     require(enumEntry.initializer is FirLazyExpression)
 
-    val newEntry = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
-        session = enumEntry.moduleData.session,
-        scopeProvider = enumEntry.moduleData.session.kotlinScopeProvider,
-        designation = designation,
-        rootNonLocalDeclaration = enumEntry.psi as KtEnumEntry,
-    ) as FirEnumEntry
-
-    enumEntry.apply {
-        replaceInitializer(newEntry.initializer)
-    }
+    val newEnumEntry = revive<FirEnumEntry>(designation)
+    enumEntry.replaceInitializer(newEnumEntry.initializer)
 }
 
 private fun calculateLazyBodyForAnonymousInitializer(designation: FirDesignation) {
     val initializer = designation.target as FirAnonymousInitializer
     require(initializer.body is FirLazyBlock)
 
-    val newInitializer = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
-        session = initializer.moduleData.session,
-        scopeProvider = initializer.moduleData.session.kotlinScopeProvider,
-        designation = designation,
-        rootNonLocalDeclaration = initializer.psi as KtAnonymousInitializer,
-    ) as FirAnonymousInitializer
-
-    initializer.apply {
-        replaceBody(newInitializer.body)
-    }
+    val newInitializer = revive<FirAnonymousInitializer>(designation)
+    initializer.replaceBody(newInitializer.body)
 }
 
 private fun needCalculatingLazyBodyForConstructor(firConstructor: FirConstructor): Boolean {
     if (needCalculatingLazyBodyForFunction(firConstructor) || firConstructor.delegatedConstructor is FirLazyDelegatedConstructorCall) {
         return true
     }
-    val deleatedConstructor = firConstructor.delegatedConstructor
-    if (deleatedConstructor is FirMultiDelegatedConstructorCall) {
-        for (delegated in deleatedConstructor.delegatedConstructorCalls) {
+    val delegatedConstructor = firConstructor.delegatedConstructor
+    if (delegatedConstructor is FirMultiDelegatedConstructorCall) {
+        for (delegated in delegatedConstructor.delegatedConstructorCalls) {
             if (delegated is FirLazyDelegatedConstructorCall) {
                 return true
             }
@@ -207,19 +234,25 @@ private fun needCalculatingLazyBodyForConstructor(firConstructor: FirConstructor
 private fun calculateLazyBodiesForField(designation: FirDesignation) {
     val field = designation.target as FirField
     require(field.initializer is FirLazyExpression)
-    val newField = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
-        session = field.moduleData.session,
-        scopeProvider = field.moduleData.session.kotlinScopeProvider,
-        designation = designation,
-        rootNonLocalDeclaration = designation.path.last().psi as KtClassOrObject,
-    ) as FirField
-    field.apply {
-        replaceInitializer(newField.initializer)
-    }
+
+    val newField = revive<FirField>(designation) { it.path.last().psi }
+    field.replaceInitializer(newField.initializer)
 }
 
-private fun needCalculatingLazyBodyForFunction(firFunction: FirFunction): Boolean =
-    firFunction.body is FirLazyBlock || firFunction.valueParameters.any { it.defaultValue is FirLazyExpression }
+private fun needCalculatingLazyBodyForContractDescriptionOwner(firContractOwner: FirContractDescriptionOwner): Boolean {
+    val contractDescription = firContractOwner.contractDescription
+    if (contractDescription is FirRawContractDescription) {
+        return contractDescription.rawEffects.any { it is FirLazyExpression }
+    }
+
+    return false
+}
+
+private fun needCalculatingLazyBodyForFunction(firFunction: FirFunction): Boolean {
+    return (firFunction.body is FirLazyBlock
+            || firFunction.valueParameters.any { it.defaultValue is FirLazyExpression })
+            || (firFunction is FirContractDescriptionOwner && needCalculatingLazyBodyForContractDescriptionOwner(firFunction))
+}
 
 private fun needCalculatingLazyBodyForProperty(firProperty: FirProperty): Boolean =
     firProperty.getter?.let { needCalculatingLazyBodyForFunction(it) } == true

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

@@ -12,8 +12,10 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveT
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirSingleResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.InvalidSessionException
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirLazyResolverRunner
+import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.USE_STATE_KEEPER
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.withOnAirDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
@@ -205,7 +207,7 @@ private fun handleExceptionFromResolve(
     }
 
     val session = firDeclarationToResolve.llFirSession
-    session.invalidate()
+    invalidateFailedSessionIfNeeded(session)
 
     val moduleData = firDeclarationToResolve.llFirModuleData
     val module = moduleData.ktModule
@@ -240,7 +242,7 @@ private fun handleExceptionFromResolve(
     }
 
     val session = designation.firFile.llFirSession
-    session.invalidate()
+    invalidateFailedSessionIfNeeded(session)
 
     val moduleData = session.llFirModuleData
     val module = moduleData.ktModule
@@ -260,4 +262,10 @@ private fun handleExceptionFromResolve(
         withEntry("moduleData", moduleData) { it.toString() }
         withEntry("firDesignationToResolve", designation) { it.toString() }
     }
+}
+
+private fun invalidateFailedSessionIfNeeded(session: LLFirSession) {
+    if (!USE_STATE_KEEPER) {
+        session.invalidate()
+    }
 }

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

@@ -11,14 +11,14 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.LLFirRetu
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformerDispatcher
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 internal abstract class LLFirAbstractBodyTargetResolver(
     resolveTarget: LLFirResolveTarget,
@@ -65,4 +65,14 @@ internal abstract class LLFirAbstractBodyTargetResolver(
         val firDesignation = FirDesignationWithFile(nestedClassesStack, declaration, resolveTarget.firFile)
         FirLazyBodiesCalculator.calculateBodies(firDesignation)
     }
+
+    protected fun <T : FirElementWithResolveState> resolve(target: T, keeper: StateKeeper<T>) {
+        resolveWithKeeper(target, keeper, ::calculateLazyBodies) {
+            rawResolve(target)
+        }
+    }
+
+    protected open fun rawResolve(target: FirElementWithResolveState) {
+        target.transformSingle(transformer, ResolutionMode.ContextIndependent)
+    }
 }

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

@@ -7,17 +7,22 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator.calculateAnnotations
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirPhaseUpdater
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkAnnotationArgumentsMappingIsResolved
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkPhase
-import org.jetbrains.kotlin.fir.FirAnnotationContainer
-import org.jetbrains.kotlin.fir.FirElementWithResolveState
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.expressionGuard
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.builder.buildArgumentList
+import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
+import org.jetbrains.kotlin.fir.references.isError
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirAnnotationArgumentsMappingTransformer
+import org.jetbrains.kotlin.fir.types.FirTypeProjection
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 
 internal object LLFirAnnotationArgumentMappingLazyResolver : LLFirLazyResolver(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING) {
     override fun resolve(
@@ -66,6 +71,61 @@ private class LLFirAnnotationArgumentsMappingTargetResolver(
     )
 
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
-        transformAnnotations(target)
+        resolveWithKeeper(target, AnnotationArgumentMappingStateKeepers.DECLARATION, ::calculateAnnotations) {
+            transformAnnotations(target)
+        }
     }
 }
+
+internal object AnnotationArgumentMappingStateKeepers {
+    private val ANNOTATION: StateKeeper<FirAnnotation> = stateKeeper {
+        add(ANNOTATION_BASE)
+        add(FirAnnotation::argumentMapping, FirAnnotation::replaceArgumentMapping)
+        add(FirAnnotation::typeArgumentsCopied, FirAnnotation::replaceTypeArguments)
+    }
+
+    private val ANNOTATION_BASE: StateKeeper<FirAnnotation> = stateKeeper { annotation ->
+        if (annotation is FirAnnotationCall) {
+            entity(annotation, ANNOTATION_CALL)
+        }
+    }
+
+    private val ANNOTATION_CALL: StateKeeper<FirAnnotationCall> = stateKeeper { annotationCall ->
+        add(FirAnnotationCall::calleeReference, FirAnnotationCall::replaceCalleeReference)
+
+        val argumentList = annotationCall.argumentList
+        if (argumentList !is FirResolvedArgumentList && argumentList !is FirEmptyArgumentList) {
+            add(FirAnnotationCall::argumentList, FirAnnotationCall::replaceArgumentList) { oldList ->
+                buildArgumentList {
+                    source = oldList.source
+                    for (argument in oldList.arguments) {
+                        val replacement = when {
+                            argument is FirPropertyAccessExpression && argument.calleeReference.isError() -> argument
+                            else -> expressionGuard(argument)
+                        }
+                        arguments.add(replacement)
+                    }
+                }
+            }
+        }
+    }
+
+    val DECLARATION: StateKeeper<FirElementWithResolveState> = stateKeeper { target ->
+        val visitor = object : FirVisitorVoid() {
+            override fun visitElement(element: FirElement) {
+                when (element) {
+                    is FirDeclaration -> if (element !== target) return // Avoid nested declarations
+                    is FirAnnotation -> entity(element, ANNOTATION)
+                    is FirStatement -> return
+                }
+
+                element.acceptChildren(this)
+            }
+        }
+
+        target.accept(visitor)
+    }
+}
+
+private val FirAnnotation.typeArgumentsCopied: List<FirTypeProjection>
+    get() = if (typeArguments.isEmpty()) emptyList() else ArrayList(typeArguments)

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

@@ -10,25 +10,31 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirEle
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirPhaseUpdater
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkDelegatedConstructorIsResolved
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkBodyIsResolved
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkDefaultValueIsResolved
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkInitializerIsResolved
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkPhase
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirEntry
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.*
 import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.checkWithAttachmentBuilder
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.declarations.utils.getExplicitBackingField
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.builder.buildLazyDelegatedConstructorCall
+import org.jetbrains.kotlin.fir.expressions.builder.buildMultiDelegatedConstructorCall
+import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirLazyDelegatedConstructorCall
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.FirSuperReference
+import org.jetbrains.kotlin.fir.references.FirThisReference
+import org.jetbrains.kotlin.fir.references.builder.buildExplicitSuperReference
+import org.jetbrains.kotlin.fir.references.builder.buildExplicitThisReference
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.dfa.FirControlFlowGraphReferenceImpl
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
+import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractsDslNames
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
-import org.jetbrains.kotlin.fir.visitors.transformSingle
-
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 
 internal object LLFirBodyLazyResolver : LLFirLazyResolver(FirResolvePhase.BODY_RESOLVE) {
     override fun resolve(
@@ -151,18 +157,201 @@ private class LLFirBodyTargetResolver(
         }
 
         when (target) {
-            is FirRegularClass -> {
-                error("should be resolved in ${::doResolveWithoutLock.name}")
-            }
-            is FirDanglingModifierList, is FirFileAnnotationsContainer, is FirTypeAlias -> {
-                // no bodies here
-            }
-            is FirCallableDeclaration, is FirAnonymousInitializer, is FirScript -> {
-                calculateLazyBodies(target)
-                target.transformSingle(transformer, ResolutionMode.ContextIndependent)
+            is FirRegularClass -> error("Should have been resolved in ${::doResolveWithoutLock.name}")
+            is FirConstructor -> resolve(target, BodyStateKeepers.CONSTRUCTOR)
+            is FirFunction -> resolve(target, BodyStateKeepers.FUNCTION)
+            is FirProperty -> resolve(target, BodyStateKeepers.PROPERTY)
+            is FirField -> resolve(target, BodyStateKeepers.FIELD)
+            is FirVariable -> resolve(target, BodyStateKeepers.VARIABLE)
+            is FirAnonymousInitializer -> resolve(target, BodyStateKeepers.ANONYMOUS_INITIALIZER)
+            is FirScript -> resolve(target, BodyStateKeepers.SCRIPT)
+            is FirDanglingModifierList,
+            is FirFileAnnotationsContainer,
+            is FirTypeAlias -> {
+                // No bodies here
             }
             else -> throwUnexpectedFirElementError(target)
         }
     }
 }
 
+internal object BodyStateKeepers {
+    val SCRIPT: StateKeeper<FirScript> = stateKeeper {
+        // TODO Lazy body is not supported for scripts yet
+    }
+
+    val ANONYMOUS_INITIALIZER: StateKeeper<FirAnonymousInitializer> = stateKeeper {
+        add(FirAnonymousInitializer::body, FirAnonymousInitializer::replaceBody, ::blockGuard)
+        add(FirAnonymousInitializer::controlFlowGraphReference, FirAnonymousInitializer::replaceControlFlowGraphReference)
+    }
+
+    val FUNCTION: StateKeeper<FirFunction> = stateKeeper { function ->
+        if (function.isCertainlyResolved) {
+            return@stateKeeper
+        }
+
+        add(FirFunction::returnTypeRef, FirFunction::replaceReturnTypeRef)
+
+        if (!isCallableWithSpecialBody(function)) {
+            preserveContractBlock(function)
+
+            add(FirFunction::body, FirFunction::replaceBody, ::blockGuard)
+            entityList(function.valueParameters, VALUE_PARAMETER)
+        }
+
+        add(FirFunction::controlFlowGraphReference, FirFunction::replaceControlFlowGraphReference)
+    }
+
+    val CONSTRUCTOR: StateKeeper<FirConstructor> = stateKeeper {
+        add(FUNCTION)
+        add(FirConstructor::delegatedConstructor, FirConstructor::replaceDelegatedConstructor, ::delegatedConstructorCallGuard)
+    }
+
+    val VARIABLE: StateKeeper<FirVariable> = stateKeeper { variable ->
+        add(FirVariable::returnTypeRef, FirVariable::replaceReturnTypeRef)
+
+        if (!isCallableWithSpecialBody(variable)) {
+            add(FirVariable::initializerIfUnresolved, FirVariable::replaceInitializer, ::expressionGuard)
+        }
+    }
+
+    private val VALUE_PARAMETER: StateKeeper<FirValueParameter> = stateKeeper { valueParameter ->
+        if (valueParameter.defaultValue != null) {
+            add(FirValueParameter::defaultValue, FirValueParameter::replaceDefaultValue, ::expressionGuard)
+        }
+
+        add(FirValueParameter::controlFlowGraphReference, FirValueParameter::replaceControlFlowGraphReference)
+    }
+
+    val FIELD: StateKeeper<FirField> = stateKeeper {
+        add(VARIABLE)
+        add(FirField::controlFlowGraphReference, FirField::replaceControlFlowGraphReference)
+    }
+
+    val PROPERTY: StateKeeper<FirProperty> = stateKeeper { property ->
+        if (property.bodyResolveState >= FirPropertyBodyResolveState.EVERYTHING_RESOLVED) {
+            return@stateKeeper
+        }
+
+        add(VARIABLE)
+
+        add(FirProperty::bodyResolveState, FirProperty::replaceBodyResolveState)
+        add(FirProperty::returnTypeRef, FirProperty::replaceReturnTypeRef)
+
+        entity(property.getterIfUnresolved, FUNCTION)
+        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)
+    }
+}
+
+context(StateKeeperBuilder)
+private fun StateKeeperScope<FirFunction>.preserveContractBlock(function: FirFunction) {
+    val oldBody = function.body
+    if (oldBody == null || oldBody is FirLazyBlock) {
+        return
+    }
+
+    val oldFirstStatement = oldBody.statements.firstOrNull() ?: return
+
+    // The old body starts with a contract definition
+    if (oldFirstStatement is FirContractCallBlock) {
+        if (oldFirstStatement.call.calleeReference is FirResolvedNamedReference) {
+            postProcess {
+                val newBody = function.body
+                if (newBody != null && newBody.statements.isNotEmpty()) {
+                    // Replace the newly created (and not yet resolved) contract block with the old, resolved one
+                    newBody.replaceFirstStatement<FirContractCallBlock> { oldFirstStatement }
+                }
+            }
+        }
+
+        return
+    }
+
+    // The old body starts with a contract-like call (but it's not a proper contract definition)
+    if (oldFirstStatement is FirFunctionCall && oldFirstStatement.calleeReference.name == FirContractsDslNames.CONTRACT.callableName) {
+        postProcess {
+            val newBody = function.body
+            if (newBody != null && newBody.statements.isNotEmpty()) {
+                val newFirstStatement = newBody.statements.first()
+                if (newFirstStatement is FirContractCallBlock) {
+                    // We already know that the function doesn't have a contract, so we can safely unwrap the contract block
+                    newBody.replaceFirstStatement<FirContractCallBlock> { newFirstStatement.call }
+                }
+            }
+        }
+    }
+}
+
+private val FirFunction.isCertainlyResolved: Boolean
+    get() {
+        if (this is FirPropertyAccessor) {
+            val requiredState = when {
+                isSetter -> FirPropertyBodyResolveState.EVERYTHING_RESOLVED
+                else -> FirPropertyBodyResolveState.INITIALIZER_AND_GETTER_RESOLVED
+            }
+
+            if (propertySymbol.fir.bodyResolveState >= requiredState) {
+                return true
+            }
+        }
+
+        val body = this.body ?: return false // Not completely sure
+        return body !is FirLazyBlock && body.typeRef is FirResolvedTypeRef
+    }
+
+private val FirVariable.initializerIfUnresolved: FirExpression?
+    get() = when (this) {
+        is FirProperty -> if (bodyResolveState < FirPropertyBodyResolveState.INITIALIZER_RESOLVED) initializer else null
+        else -> initializer
+    }
+
+private val FirProperty.backingFieldIfUnresolved: FirBackingField?
+    get() = if (bodyResolveState < FirPropertyBodyResolveState.INITIALIZER_RESOLVED) getExplicitBackingField() else null
+
+private val FirProperty.getterIfUnresolved: FirPropertyAccessor?
+    get() = if (bodyResolveState < FirPropertyBodyResolveState.INITIALIZER_AND_GETTER_RESOLVED) getter else null
+
+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
+    } else if (fir is FirMultiDelegatedConstructorCall) {
+        return buildMultiDelegatedConstructorCall {
+            for (delegatedConstructorCall in fir.delegatedConstructorCalls) {
+                delegatedConstructorCalls.add(delegatedConstructorCallGuard(delegatedConstructorCall))
+            }
+        }
+    }
+
+    return buildLazyDelegatedConstructorCall {
+        constructedTypeRef = fir.constructedTypeRef
+        when (val originalCalleeReference = fir.calleeReference) {
+            is FirThisReference -> {
+                isThis = true
+                calleeReference = buildExplicitThisReference {
+                    source = null
+                }
+            }
+            is FirSuperReference -> {
+                isThis = false
+                calleeReference = buildExplicitSuperReference {
+                    source = originalCalleeReference.source
+                    superTypeRef = originalCalleeReference.superTypeRef
+                }
+            }
+        }
+    }
+}

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

@@ -10,16 +10,16 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirEle
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirPhaseUpdater
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkPhase
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.blockGuard
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isCallableWithSpecialBody
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.contracts.FirRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveTransformer
-import org.jetbrains.kotlin.fir.visitors.transformSingle
-
 
 internal object LLFirContractsLazyResolver : LLFirLazyResolver(FirResolvePhase.CONTRACTS) {
     override fun resolve(
@@ -61,15 +61,61 @@ private class LLFirContractsTargetResolver(
 
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
         when (target) {
-            is FirRegularClass, is FirAnonymousInitializer, is FirDanglingModifierList, is FirFileAnnotationsContainer, is FirTypeAlias, is FirScript -> {
-                // no contracts here
-            }
-            is FirCallableDeclaration -> {
-                // TODO calculate bodies only when in-body contract is present
-                calculateLazyBodies(target)
-                target.transformSingle(transformer, ResolutionMode.ContextIndependent)
+            is FirSimpleFunction -> resolve(target, ContractStateKeepers.SIMPLE_FUNCTION)
+            is FirConstructor -> resolve(target, ContractStateKeepers.CONSTRUCTOR)
+            is FirProperty -> resolve(target, ContractStateKeepers.PROPERTY)
+            is FirPropertyAccessor -> resolve(target, ContractStateKeepers.PROPERTY_ACCESSOR)
+            is FirRegularClass,
+            is FirTypeAlias,
+            is FirVariable,
+            is FirFunction,
+            is FirAnonymousInitializer,
+            is FirScript,
+            is FirFileAnnotationsContainer,
+            is FirDanglingModifierList -> {
+                // No contracts here
+                check(target !is FirContractDescriptionOwner) {
+                    "Unexpected contract description owner: $target (${target.javaClass.name})"
+                }
             }
             else -> throwUnexpectedFirElementError(target)
         }
     }
 }
+
+private object ContractStateKeepers {
+    private val CONTRACT_DESCRIPTION_OWNER: StateKeeper<FirContractDescriptionOwner> = stateKeeper {
+        add(FirContractDescriptionOwner::contractDescription, FirContractDescriptionOwner::replaceContractDescription)
+    }
+
+    private val BODY_OWNER: StateKeeper<FirFunction> = stateKeeper { declaration ->
+        if (declaration is FirContractDescriptionOwner && declaration.contractDescription is FirRawContractDescription) {
+            // No need to change the body, contract is declared separately
+            return@stateKeeper
+        }
+
+        if (!isCallableWithSpecialBody(declaration)) {
+            add(FirFunction::body, FirFunction::replaceBody, ::blockGuard)
+        }
+    }
+
+    val SIMPLE_FUNCTION: StateKeeper<FirSimpleFunction> = stateKeeper {
+        add(CONTRACT_DESCRIPTION_OWNER)
+        add(BODY_OWNER)
+    }
+
+    val CONSTRUCTOR: StateKeeper<FirConstructor> = stateKeeper {
+        add(CONTRACT_DESCRIPTION_OWNER)
+        add(BODY_OWNER)
+    }
+
+    val PROPERTY_ACCESSOR: StateKeeper<FirPropertyAccessor> = stateKeeper {
+        add(CONTRACT_DESCRIPTION_OWNER)
+        add(BODY_OWNER)
+    }
+
+    val PROPERTY: StateKeeper<FirProperty> = stateKeeper { property ->
+        entity(property.getter, PROPERTY_ACCESSOR)
+        entity(property.setter, PROPERTY_ACCESSOR)
+    }
+}

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

@@ -15,13 +15,11 @@ import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitAwareBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.scopes.fakeOverrideSubstitution
-import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 internal object LLFirImplicitTypesLazyResolver : LLFirLazyResolver(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
     override fun resolve(
@@ -82,28 +80,27 @@ internal class LLFirImplicitBodyTargetResolver(
 
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
         when (target) {
+            is FirFunction -> resolve(target, BodyStateKeepers.FUNCTION)
+            is FirProperty -> resolve(target, BodyStateKeepers.PROPERTY)
+            is FirVariable -> resolve(target, BodyStateKeepers.VARIABLE)
             is FirRegularClass,
-            is FirDanglingModifierList,
-            is FirAnonymousInitializer,
-            is FirFileAnnotationsContainer,
             is FirTypeAlias,
             is FirScript,
-            is FirConstructor,
-            -> {
-                // no implicit bodies here
+            is FirAnonymousInitializer,
+            is FirDanglingModifierList,
+            is FirFileAnnotationsContainer -> {
+                // No implicit bodies here
             }
-
-            is FirCallableDeclaration -> {
-                if (target.attributes.fakeOverrideSubstitution != null) {
-                    transformer.returnTypeCalculator.fakeOverrideTypeCalculator.computeReturnType(target)
-                } else {
-                    calculateLazyBodies(target)
-                    target.transformSingle(transformer, ResolutionMode.ContextIndependent)
-                }
-            }
-
             else -> throwUnexpectedFirElementError(target)
         }
     }
-}
 
+    override fun rawResolve(target: FirElementWithResolveState) {
+        if (target is FirCallableDeclaration && target.attributes.fakeOverrideSubstitution != null) {
+            transformer.returnTypeCalculator.fakeOverrideTypeCalculator.computeReturnType(target)
+            return
+        }
+
+        super.rawResolve(target)
+    }
+}

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

@@ -0,0 +1,307 @@
+/*
+ * Copyright 2010-2023 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
+
+import com.intellij.openapi.util.registry.Registry
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+
+@DslMarker
+internal annotation class StateKeeperDsl
+
+/**
+ * Post-processors run after arrangers, but before the action block.
+ * It is a work-around for cases when an arranger cannot properly tweak the state by itself
+ * (for instance, when such tweak depends on a non-local mutation).
+ */
+internal typealias PostProcessor = () -> Unit
+
+@StateKeeperDsl
+internal interface StateKeeperBuilder {
+    fun register(state: PreservedState)
+}
+
+@JvmInline
+@StateKeeperDsl
+internal value class StateKeeperScope<Owner : Any>(private val owner: Owner) {
+    /**
+     * Defines an entity state preservation rule.
+     *
+     * @param provider a function that returns the current entity state (a getter).
+     * @param mutator a function that modifies the entity state (a setter).
+     * @param arranger a function that provides a tweaked entity state. Such a state is then applied using a [mutator].
+     */
+    context(StateKeeperBuilder)
+    inline fun <Value> add(provider: (Owner) -> Value, crossinline mutator: (Owner, Value) -> Unit, arranger: (Value & Any) -> Value) {
+        val owner = this@StateKeeperScope.owner
+
+        val storedValue = provider(owner)
+        if (storedValue != null) {
+            val arrangedValue = arranger(storedValue)
+            if (arrangedValue !== storedValue) {
+                mutator(owner, arrangedValue)
+            }
+        }
+
+        register(object : PreservedState {
+            override fun restore() = mutator(owner, storedValue)
+            override fun postProcess() {}
+        })
+    }
+
+    /**
+     * Defines an entity state preservation rule.
+     *
+     * @param provider a function that returns the current entity state (a getter).
+     * @param mutator a function that modifies the entity state (a setter).
+     */
+    context(StateKeeperBuilder)
+    inline fun <Value> add(provider: (Owner) -> Value, crossinline mutator: (Owner, Value) -> Unit) {
+        val owner = this@StateKeeperScope.owner
+        val storedValue = provider(owner)
+
+        register(object : PreservedState {
+            override fun restore() = mutator(owner, storedValue)
+            override fun postProcess() {}
+        })
+    }
+
+    /**
+     * Defines an entity state preservation rule by delegating to a given [keeper].
+     * In other words, applies all rules defined in the [keeper] to the current entity.
+     */
+    context(StateKeeperBuilder)
+    fun add(keeper: StateKeeper<Owner>) {
+        val owner = this@StateKeeperScope.owner
+        register(keeper.prepare(owner))
+    }
+
+    /**
+     * Defines a post-processor.
+     */
+    context(StateKeeperBuilder)
+    fun postProcess(block: PostProcessor) {
+        register(object : PreservedState {
+            override fun postProcess() = block()
+            override fun restore() {}
+        })
+    }
+}
+
+/**
+ * Registers a given [entity] using the delegate [keeper].
+ * Does nothing if the [entity] is `null`.
+ */
+context(StateKeeperBuilder)
+internal fun <Entity : Any> entity(entity: Entity?, keeper: StateKeeper<Entity>) {
+    if (entity != null) {
+        StateKeeperScope(entity).add(keeper)
+    }
+}
+
+/**
+ * Registers a given [entity] using the building [block].
+ * Does nothing if the [entity] is `null`.
+ */
+context(StateKeeperBuilder)
+internal inline fun <Entity : Any> entity(entity: Entity?, block: StateKeeperScope<Entity>.(Entity) -> Unit) {
+    if (entity != null) {
+        StateKeeperScope(entity).block(entity)
+    }
+}
+
+/**
+ * Registers all entities in a given [list] sequentially using the delegate [keeper].
+ * Does nothing if the [list] is `null`.
+ * Skips `null` elements in the [list].
+ */
+context(StateKeeperBuilder)
+internal fun <Entity : Any> entityList(list: List<Entity?>?, keeper: StateKeeper<Entity>) {
+    if (list != null) {
+        for (entity in list) {
+            if (entity != null) {
+                StateKeeperScope(entity).add(keeper)
+            }
+        }
+    }
+}
+
+/**
+ * Registers all entities in a given [list] sequentially using the building [block].
+ * Does nothing if the [list] is `null`.
+ * Skips `null` elements in the [list].
+ */
+context(StateKeeperBuilder)
+internal inline fun <Entity : Any> entityList(list: List<Entity?>?, block: StateKeeperScope<Entity>.(Entity) -> Unit) {
+    if (list != null) {
+        for (entity in list) {
+            if (entity != null) {
+                StateKeeperScope(entity).block(entity)
+            }
+        }
+    }
+}
+
+/**
+ * Defines a [StateKeeper] using a builder DSL.
+ * This function is supposed to be the main entry point for [StateKeeper] creation.
+ *
+ * @param block a function that defines state preservation rules.
+ *  The function collects rules for each individual owner separately.
+ *  Nested owners can be handled inside [entity] or [entityList] blocks.
+ */
+internal fun <Owner : Any> stateKeeper(block: context(StateKeeperBuilder) StateKeeperScope<Owner>.(Owner) -> Unit): StateKeeper<Owner> {
+    return StateKeeper { owner ->
+        val states = mutableListOf<PreservedState>()
+
+        val builder = object : StateKeeperBuilder {
+            override fun register(state: PreservedState) {
+                states += state
+            }
+        }
+
+        val scope = StateKeeperScope(owner)
+        block(builder, scope, owner)
+
+        object : PreservedState {
+            private var isPostProcessed = false
+            private var isRestored = false
+
+            override fun postProcess() {
+                if (isPostProcessed) {
+                    return
+                }
+
+                isPostProcessed = true
+                states.forEach { it.postProcess() }
+            }
+
+            override fun restore() {
+                if (isRestored) {
+                    return
+                }
+
+                isRestored = true
+                states.forEach { it.restore() }
+            }
+        }
+    }
+}
+
+/**
+ * [StateKeeper] backs up parts of an object state, and allows to restore it if errors occur during that object mutation.
+ * This is not a complete object dumper. All preservation rules are explicitly defined (usually in the [stateKeeper] DSL).
+ *
+ * State keeper provides the following life cycle:
+ * 1. Preparation. Back up the state.
+ * 2. Arrangement. For rules with arrangers, use the state provider by the arranger.
+ * 3. Post-processing. Call supplied post-processors for additional state tweaks.
+ * 4. Action. Perform the potentially failing action.
+ * 5. Restoration (optional). Restore the state if the action failed.
+ *
+ * Preparation and arrangement steps are run by calling the [prepare] function.
+ * Post-processing is run by calling [PreservedState.postProcess], potentially after some manual tweaks.
+ * Restoration is run by calling [PreservedState.restore] if the action failed.
+ *
+ * @sample
+ * ```
+ * var state: PreservedState? = null
+ *
+ * try {
+ *     state = keeper.prepare(owner)
+ *     state.postProcess()
+ *     action(owner)
+ * } catch (e: Throwable) {
+ *     state?.restore()
+ * }
+ * ```
+ */
+internal class StateKeeper<in Owner : Any>(val provider: (Owner) -> PreservedState) {
+    /**
+     * Backs up the [owner] state by calling providers, and potentially tweaks the object state with arrangers.
+     * Note that post-processors are not run during preparation.
+     */
+    fun prepare(owner: Owner): PreservedState {
+        return provider(owner)
+    }
+}
+
+internal val USE_STATE_KEEPER: Boolean = Registry.`is`("kotlin.analysis.use.state.keeper", true)
+
+/**
+ * State preserved by a [StateKeeper].
+ */
+internal interface PreservedState {
+    /**
+     * Performs post-processing of the state, according to supplied [PostProcessor]s.
+     * This function must be called before the potentially failing action block.
+     */
+    fun postProcess()
+
+    /**
+     * Restores the state, according to rules defined by the [StateKeeper].
+     * This function must be called when the action block fails.
+     */
+    fun restore()
+}
+
+internal inline fun <Target : FirElementWithResolveState, Result> resolveWithKeeper(
+    target: Target,
+    keeper: StateKeeper<Target>,
+    prepareTarget: (Target) -> Unit,
+    action: () -> Result
+): Result {
+    if (!USE_STATE_KEEPER) {
+        prepareTarget(target)
+        return action()
+    }
+
+    val onAirAnalysisTarget = target.moduleData.session.onAirAnalysisTarget
+    if (onAirAnalysisTarget != null && target in onAirAnalysisTarget) {
+        // Several arrangers reset declaration bodies to lazy ones, and then re-construct the FIR tree from the backing PSI.
+        // This won't work for on-air analysis, as the tree is manually modified. However, it doesn't seem that tree guards are needed
+        // in the first place, as results of the on-air analysis aren't going to be shared.
+        prepareTarget(target)
+        return action()
+    }
+
+    var preservedState: PreservedState? = null
+
+    try {
+        preservedState = keeper.prepare(target)
+        prepareTarget(target)
+        preservedState.postProcess()
+        return action()
+    } catch (e: Throwable) {
+        preservedState?.restore()
+        throw e
+    }
+}
+
+private operator fun FirElementWithResolveState.contains(child: FirElementWithResolveState): Boolean {
+    if (this == child) {
+        return true
+    }
+
+    var isFound = false
+
+    accept(object : FirVisitorVoid() {
+        override fun visitElement(element: FirElement) {
+            if (!isFound) {
+                if (element == child) {
+                    isFound = true
+                } else if (element is FirDeclaration || element !is FirStatement) {
+                    element.acceptChildren(this)
+                }
+            }
+        }
+    })
+
+    return isFound
+}

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
@@ -24,6 +25,9 @@ internal interface SyntheticFirClassProvider {
     }
 }
 
+internal val FirSession.onAirAnalysisTarget: FirElementWithResolveState?
+    get() = onAirProviderForThread.get()?.target
+
 /**
  * Injects a designation-based class file provider to 'LLFirProvider'.
  * Used for on-air analysis.
@@ -44,6 +48,7 @@ internal fun withOnAirDesignation(designation: FirDesignationWithFile, block: ()
 
 private class OnAirSyntheticFirClassProvider private constructor(
     private val firFile: FirFile,
+    val target: FirElementWithResolveState,
     private val classes: Map<ClassId, FirClassLikeDeclaration>
 ) : SyntheticFirClassProvider {
     val session: FirSession
@@ -73,7 +78,7 @@ private class OnAirSyntheticFirClassProvider private constructor(
             }
 
             nodeInfoCollector.visitElement(firElement)
-            return OnAirSyntheticFirClassProvider(firFile, nodeInfoCollector.classes)
+            return OnAirSyntheticFirClassProvider(firFile, designation.target, nodeInfoCollector.classes)
         }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/stateKeeperUtils.kt

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2023 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.util
+
+import org.jetbrains.kotlin.KtFakeSourceElement
+import org.jetbrains.kotlin.KtFakeSourceElementKind
+import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.builder.buildLazyBlock
+import org.jetbrains.kotlin.fir.expressions.builder.buildLazyExpression
+
+internal fun blockGuard(fir: FirBlock): FirBlock {
+    if (isLazyStatement(fir)) {
+        return fir
+    }
+
+    return buildLazyBlock()
+}
+
+internal fun expressionGuard(fir: FirExpression): FirExpression {
+    if (isLazyStatement(fir)) {
+        return fir
+    }
+
+    return buildLazyExpression {
+        source = fir.source
+    }
+}
+
+private fun isLazyStatement(fir: FirStatement): Boolean {
+    return fir is FirLazyExpression || fir is FirLazyBlock
+}
+
+private val SPECIAL_BODY_CALLABLE_SOURCE_KINDS = setOf(
+    KtFakeSourceElementKind.DefaultAccessor,
+    KtFakeSourceElementKind.DelegatedPropertyAccessor,
+    KtFakeSourceElementKind.ImplicitConstructor,
+    KtFakeSourceElementKind.PropertyFromParameter,
+    KtFakeSourceElementKind.DataClassGeneratedMembers,
+    KtFakeSourceElementKind.EnumGeneratedDeclaration,
+)
+
+internal fun isCallableWithSpecialBody(fir: FirCallableDeclaration): Boolean {
+    val source = fir.source as? KtFakeSourceElement ?: return false
+    return source.kind in SPECIAL_BODY_CALLABLE_SOURCE_KINDS
+}

analysis/low-level-api-fir/testdata/lazyResolve/delegatedField.txt

@@ -232,7 +232,7 @@ FILE: [ResolvedTo(IMPORTS)] delegatedField.kt
             LAZY_super<<implicit>>
         }
 
-        private final [ResolvedTo(CONTRACTS)] field $$delegate_0: R|one/Foo| = prop#
+        private final [ResolvedTo(CONTRACTS)] field $$delegate_0: R|one/Foo| = LAZY_EXPRESSION
 
         public final [ResolvedTo(STATUS)] fun baz([ResolvedTo(STATUS)] s: R|kotlin/String|): R|kotlin/Unit| { LAZY_BLOCK }
 

analysis/low-level-api-fir/testdata/lazyResolve/enumEntry.txt

@@ -246,13 +246,7 @@ FILE: [ResolvedTo(IMPORTS)] enumEntry.kt
             LAZY_super<R|kotlin/Enum<Foo>|>
         }
 
-        @R|Anno|[Types]() public final static [ResolvedTo(CONTRACTS)] [ContainingClassKey=Foo] enum entry ResolveMe: R|Foo| = object : R|Foo| {
-            private [ResolvedTo(RAW_FIR)] [ContainingClassKey=<anonymous>] constructor(): R|<anonymous>| {
-                super<R|Foo|>()
-            }
-
-        }
-
+        @R|Anno|[Types]() public final static [ResolvedTo(CONTRACTS)] [ContainingClassKey=Foo] enum entry ResolveMe: R|Foo| = LAZY_EXPRESSION
         public final static [ResolvedTo(STATUS)] [ContainingClassKey=Foo] fun values(): R|kotlin/Array<Foo>| {
         }