diff --git a/analysis/low-level-api-fir/build.gradle.kts b/analysis/low-level-api-fir/build.gradle.kts
index 9e8cec824c1..cb12f3fa995 100644
--- a/analysis/low-level-api-fir/build.gradle.kts
+++ b/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
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 62ce22c56f9..fba8232f7c9 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
@@ -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
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt
index 6d0362c34ec..eda36338b43 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt
+++ b/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()
+    }
 }
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAbstractBodyTargetResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAbstractBodyTargetResolver.kt
index 040adf04bd0..149dec77ee1 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAbstractBodyTargetResolver.kt
+++ b/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)
+    }
 }
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentMappingLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentMappingLazyResolver.kt
index fa21636291b..414640932b6 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentMappingLazyResolver.kt
+++ b/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)
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
index 50ff1fa8c06..04d69dec681 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
+++ b/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
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirContractsLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirContractsLazyResolver.kt
index 14b557f3ecd..bd97697dd4e 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirContractsLazyResolver.kt
+++ b/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)
+    }
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirImplicitTypesLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirImplicitTypesLazyResolver.kt
index c3d42dde6d8..5d91ff8bffb 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirImplicitTypesLazyResolver.kt
+++ b/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)
+    }
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/StateKeeper.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/StateKeeper.kt
new file mode 100644
index 00000000000..85de10e3cab
--- /dev/null
+++ b/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
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/SyntheticFirClassProvider.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/SyntheticFirClassProvider.kt
index 87ae7d3af34..c563135351d 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/SyntheticFirClassProvider.kt
+++ b/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)
         }
     }
 }
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/stateKeeperUtils.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/stateKeeperUtils.kt
new file mode 100644
index 00000000000..c5003e0e892
--- /dev/null
+++ b/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
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/testdata/lazyResolve/delegatedField.txt b/analysis/low-level-api-fir/testdata/lazyResolve/delegatedField.txt
index 0f87233baba..29357fcc2cd 100644
--- a/analysis/low-level-api-fir/testdata/lazyResolve/delegatedField.txt
+++ b/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 }
 
diff --git a/analysis/low-level-api-fir/testdata/lazyResolve/enumEntry.txt b/analysis/low-level-api-fir/testdata/lazyResolve/enumEntry.txt
index dec8897bb31..11c4fb7c156 100644
--- a/analysis/low-level-api-fir/testdata/lazyResolve/enumEntry.txt
+++ b/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>| {
         }