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 69295b292d4..793c17dca06 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 @@ -69,13 +69,9 @@ internal abstract class LLFirAbstractBodyTargetResolver( } } - protected fun calculateLazyBodies(declaration: FirElementWithResolveState) { - val firDesignation = FirDesignationWithFile(nestedClassesStack, declaration, resolveTarget.firFile) - FirLazyBodiesCalculator.calculateBodies(firDesignation) - } - - protected fun resolve(target: T, keeper: StateKeeper) { - resolveWithKeeper(target, keeper, ::calculateLazyBodies) { + protected fun resolve(target: T, keeper: StateKeeper) { + val firDesignation = FirDesignationWithFile(nestedClassesStack, target, resolveTarget.firFile) + resolveWithKeeper(target, firDesignation, keeper, { FirLazyBodiesCalculator.calculateBodies(firDesignation) }) { rawResolve(target) } } 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 2af3ad4d67c..f547d4f6680 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 @@ -8,6 +8,7 @@ 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.sessions.llFirSession import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkAnnotationArgumentsMappingIsResolved import org.jetbrains.kotlin.analysis.low.level.api.fir.util.expressionGuard import org.jetbrains.kotlin.fir.* @@ -61,26 +62,26 @@ private class LLFirAnnotationArgumentsMappingTargetResolver( ) override fun doLazyResolveUnderLock(target: FirElementWithResolveState) { - resolveWithKeeper(target, AnnotationArgumentMappingStateKeepers.DECLARATION, ::calculateAnnotations) { + resolveWithKeeper(target, target.llFirSession, AnnotationArgumentMappingStateKeepers.DECLARATION, ::calculateAnnotations) { transformAnnotations(target) } } } internal object AnnotationArgumentMappingStateKeepers { - private val ANNOTATION: StateKeeper = stateKeeper { - add(ANNOTATION_BASE) + private val ANNOTATION: StateKeeper = stateKeeper { _, session -> + add(ANNOTATION_BASE, session) add(FirAnnotation::argumentMapping, FirAnnotation::replaceArgumentMapping) add(FirAnnotation::typeArgumentsCopied, FirAnnotation::replaceTypeArguments) } - private val ANNOTATION_BASE: StateKeeper = stateKeeper { annotation -> + private val ANNOTATION_BASE: StateKeeper = stateKeeper { annotation, session -> if (annotation is FirAnnotationCall) { - entity(annotation, ANNOTATION_CALL) + entity(annotation, ANNOTATION_CALL, session) } } - private val ANNOTATION_CALL: StateKeeper = stateKeeper { annotationCall -> + private val ANNOTATION_CALL: StateKeeper = stateKeeper { annotationCall, _ -> add(FirAnnotationCall::calleeReference, FirAnnotationCall::replaceCalleeReference) val argumentList = annotationCall.argumentList @@ -100,12 +101,12 @@ internal object AnnotationArgumentMappingStateKeepers { } } - val DECLARATION: StateKeeper = stateKeeper { target -> + val DECLARATION: StateKeeper = stateKeeper { target, session -> 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 FirAnnotation -> entity(element, ANNOTATION, session) is FirStatement -> return } 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 42041c434d3..33e635bd7f7 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 @@ -5,6 +5,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.analysis.low.level.api.fir.api.targets.LLFirResolveTarget import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider @@ -154,16 +155,16 @@ private class LLFirBodyTargetResolver( } internal object BodyStateKeepers { - val SCRIPT: StateKeeper = stateKeeper { + val SCRIPT: StateKeeper = stateKeeper { _, _ -> // TODO Lazy body is not supported for scripts yet } - val ANONYMOUS_INITIALIZER: StateKeeper = stateKeeper { + val ANONYMOUS_INITIALIZER: StateKeeper = stateKeeper { _, _ -> add(FirAnonymousInitializer::body, FirAnonymousInitializer::replaceBody, ::blockGuard) add(FirAnonymousInitializer::controlFlowGraphReference, FirAnonymousInitializer::replaceControlFlowGraphReference) } - val FUNCTION: StateKeeper = stateKeeper { function -> + val FUNCTION: StateKeeper = stateKeeper { function, designation -> if (function.isCertainlyResolved) { return@stateKeeper } @@ -174,18 +175,18 @@ internal object BodyStateKeepers { preserveContractBlock(function) add(FirFunction::body, FirFunction::replaceBody, ::blockGuard) - entityList(function.valueParameters, VALUE_PARAMETER) + entityList(function.valueParameters, VALUE_PARAMETER, designation) } add(FirFunction::controlFlowGraphReference, FirFunction::replaceControlFlowGraphReference) } - val CONSTRUCTOR: StateKeeper = stateKeeper { - add(FUNCTION) + val CONSTRUCTOR: StateKeeper = stateKeeper { _, designation -> + add(FUNCTION, designation) add(FirConstructor::delegatedConstructor, FirConstructor::replaceDelegatedConstructor, ::delegatedConstructorCallGuard) } - val VARIABLE: StateKeeper = stateKeeper { variable -> + val VARIABLE: StateKeeper = stateKeeper { variable, _ -> add(FirVariable::returnTypeRef, FirVariable::replaceReturnTypeRef) if (!isCallableWithSpecialBody(variable)) { @@ -194,7 +195,7 @@ internal object BodyStateKeepers { } } - private val VALUE_PARAMETER: StateKeeper = stateKeeper { valueParameter -> + private val VALUE_PARAMETER: StateKeeper = stateKeeper { valueParameter, _ -> if (valueParameter.defaultValue != null) { add(FirValueParameter::defaultValue, FirValueParameter::replaceDefaultValue, ::expressionGuard) } @@ -202,31 +203,31 @@ internal object BodyStateKeepers { add(FirValueParameter::controlFlowGraphReference, FirValueParameter::replaceControlFlowGraphReference) } - val FIELD: StateKeeper = stateKeeper { - add(VARIABLE) + val FIELD: StateKeeper = stateKeeper { _, designation -> + add(VARIABLE, designation) add(FirField::controlFlowGraphReference, FirField::replaceControlFlowGraphReference) } - val PROPERTY: StateKeeper = stateKeeper { property -> + val PROPERTY: StateKeeper = stateKeeper { property, designation -> if (property.bodyResolveState >= FirPropertyBodyResolveState.EVERYTHING_RESOLVED) { return@stateKeeper } - add(VARIABLE) + add(VARIABLE, designation) add(FirProperty::bodyResolveState, FirProperty::replaceBodyResolveState) add(FirProperty::returnTypeRef, FirProperty::replaceReturnTypeRef) - entity(property.getterIfUnresolved, FUNCTION) - entity(property.setterIfUnresolved, FUNCTION) - entity(property.backingFieldIfUnresolved, VARIABLE) + entity(property.getterIfUnresolved, FUNCTION, designation) + entity(property.setterIfUnresolved, FUNCTION, designation) + entity(property.backingFieldIfUnresolved, VARIABLE, designation) add(FirProperty::controlFlowGraphReference, FirProperty::replaceControlFlowGraphReference) } } context(StateKeeperBuilder) -private fun StateKeeperScope.preserveContractBlock(function: FirFunction) { +private fun StateKeeperScope.preserveContractBlock(function: FirFunction) { val oldBody = function.body if (oldBody == null || oldBody is FirLazyBlock) { return 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 c3659f72940..7399c60ba1a 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 @@ -5,6 +5,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.analysis.low.level.api.fir.api.targets.LLFirResolveTarget import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider @@ -77,11 +78,11 @@ private class LLFirContractsTargetResolver( } private object ContractStateKeepers { - private val CONTRACT_DESCRIPTION_OWNER: StateKeeper = stateKeeper { + private val CONTRACT_DESCRIPTION_OWNER: StateKeeper = stateKeeper { _, _ -> add(FirContractDescriptionOwner::contractDescription, FirContractDescriptionOwner::replaceContractDescription) } - private val BODY_OWNER: StateKeeper = stateKeeper { declaration -> + private val BODY_OWNER: StateKeeper = stateKeeper { declaration, _ -> if (declaration is FirContractDescriptionOwner && declaration.contractDescription is FirRawContractDescription) { // No need to change the body, contract is declared separately return@stateKeeper @@ -92,23 +93,23 @@ private object ContractStateKeepers { } } - val SIMPLE_FUNCTION: StateKeeper = stateKeeper { - add(CONTRACT_DESCRIPTION_OWNER) - add(BODY_OWNER) + val SIMPLE_FUNCTION: StateKeeper = stateKeeper { _, designation -> + add(CONTRACT_DESCRIPTION_OWNER, designation) + add(BODY_OWNER, designation) } - val CONSTRUCTOR: StateKeeper = stateKeeper { - add(CONTRACT_DESCRIPTION_OWNER) - add(BODY_OWNER) + val CONSTRUCTOR: StateKeeper = stateKeeper { _, designation -> + add(CONTRACT_DESCRIPTION_OWNER, designation) + add(BODY_OWNER, designation) } - val PROPERTY_ACCESSOR: StateKeeper = stateKeeper { - add(CONTRACT_DESCRIPTION_OWNER) - add(BODY_OWNER) + val PROPERTY_ACCESSOR: StateKeeper = stateKeeper { _, designation -> + add(CONTRACT_DESCRIPTION_OWNER, designation) + add(BODY_OWNER, designation) } - val PROPERTY: StateKeeper = stateKeeper { property -> - entity(property.getter, PROPERTY_ACCESSOR) - entity(property.setter, PROPERTY_ACCESSOR) + val PROPERTY: StateKeeper = stateKeeper { property, designation -> + entity(property.getter, PROPERTY_ACCESSOR, designation) + entity(property.setter, PROPERTY_ACCESSOR, designation) } } \ 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 index 85de10e3cab..ee0ecd130d9 100644 --- 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 @@ -8,6 +8,7 @@ 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.FirSession import org.jetbrains.kotlin.fir.declarations.FirDeclaration import org.jetbrains.kotlin.fir.expressions.FirStatement import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid @@ -29,7 +30,7 @@ internal interface StateKeeperBuilder { @JvmInline @StateKeeperDsl -internal value class StateKeeperScope(private val owner: Owner) { +internal value class StateKeeperScope(private val owner: Owner) { /** * Defines an entity state preservation rule. * @@ -77,9 +78,9 @@ internal value class StateKeeperScope(private val owner: Owner) { * In other words, applies all rules defined in the [keeper] to the current entity. */ context(StateKeeperBuilder) - fun add(keeper: StateKeeper) { + fun add(keeper: StateKeeper, context: Context) { val owner = this@StateKeeperScope.owner - register(keeper.prepare(owner)) + register(keeper.prepare(owner, context)) } /** @@ -99,9 +100,9 @@ internal value class StateKeeperScope(private val owner: Owner) { * Does nothing if the [entity] is `null`. */ context(StateKeeperBuilder) -internal fun entity(entity: Entity?, keeper: StateKeeper) { +internal fun entity(entity: Entity?, keeper: StateKeeper, context: Context) { if (entity != null) { - StateKeeperScope(entity).add(keeper) + StateKeeperScope(entity).add(keeper, context) } } @@ -110,9 +111,13 @@ internal fun entity(entity: Entity?, keeper: StateKeeper) * Does nothing if the [entity] is `null`. */ context(StateKeeperBuilder) -internal inline fun entity(entity: Entity?, block: StateKeeperScope.(Entity) -> Unit) { +internal inline fun entity( + entity: Entity?, + context: Context, + block: StateKeeperScope.(Entity, Context) -> Unit, +) { if (entity != null) { - StateKeeperScope(entity).block(entity) + StateKeeperScope(entity).block(entity, context) } } @@ -122,11 +127,11 @@ internal inline fun entity(entity: Entity?, block: StateKeeperSco * Skips `null` elements in the [list]. */ context(StateKeeperBuilder) -internal fun entityList(list: List?, keeper: StateKeeper) { +internal fun entityList(list: List?, keeper: StateKeeper, context: Context) { if (list != null) { for (entity in list) { if (entity != null) { - StateKeeperScope(entity).add(keeper) + StateKeeperScope(entity).add(keeper, context) } } } @@ -138,11 +143,15 @@ internal fun entityList(list: List?, keeper: StateKeeper * Skips `null` elements in the [list]. */ context(StateKeeperBuilder) -internal inline fun entityList(list: List?, block: StateKeeperScope.(Entity) -> Unit) { +internal inline fun entityList( + list: List?, + context: Context, + block: StateKeeperScope.(Entity, Context) -> Unit, +) { if (list != null) { for (entity in list) { if (entity != null) { - StateKeeperScope(entity).block(entity) + StateKeeperScope(entity).block(entity, context) } } } @@ -156,8 +165,8 @@ internal inline fun entityList(list: List?, block: State * The function collects rules for each individual owner separately. * Nested owners can be handled inside [entity] or [entityList] blocks. */ -internal fun stateKeeper(block: context(StateKeeperBuilder) StateKeeperScope.(Owner) -> Unit): StateKeeper { - return StateKeeper { owner -> +internal fun stateKeeper(block: context(StateKeeperBuilder) StateKeeperScope.(Owner, Context) -> Unit): StateKeeper { + return StateKeeper { owner, context -> val states = mutableListOf() val builder = object : StateKeeperBuilder { @@ -166,8 +175,8 @@ internal fun stateKeeper(block: context(StateKeeperBuilder) StateK } } - val scope = StateKeeperScope(owner) - block(builder, scope, owner) + val scope = StateKeeperScope(owner) + block(builder, scope, owner, context) object : PreservedState { private var isPostProcessed = false @@ -222,13 +231,13 @@ internal fun stateKeeper(block: context(StateKeeperBuilder) StateK * } * ``` */ -internal class StateKeeper(val provider: (Owner) -> PreservedState) { +internal class StateKeeper(val provider: (Owner, Context) -> 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) + fun prepare(owner: Owner, context: Context): PreservedState { + return provider(owner, context) } } @@ -251,11 +260,12 @@ internal interface PreservedState { fun restore() } -internal inline fun resolveWithKeeper( +internal inline fun resolveWithKeeper( target: Target, - keeper: StateKeeper, - prepareTarget: (Target) -> Unit, - action: () -> Result + context: Context, + keeper: StateKeeper, + prepareTarget: (Target) -> Unit = {}, + action: () -> Result, ): Result { if (!USE_STATE_KEEPER) { prepareTarget(target) @@ -274,7 +284,7 @@ internal inline fun resolveWithKee var preservedState: PreservedState? = null try { - preservedState = keeper.prepare(target) + preservedState = keeper.prepare(target, context) prepareTarget(target) preservedState.postProcess() return action() diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/StateKeeperTest.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/StateKeeperTest.kt index 6a510ce4588..3e5a1784110 100644 --- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/StateKeeperTest.kt +++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/StateKeeperTest.kt @@ -22,7 +22,7 @@ class StateKeeperTest { } } - val keeper: StateKeeper = stateKeeper { + val keeper: StateKeeper = stateKeeper { _, _ -> add(Foo::getValue, Foo::setValue) } @@ -41,7 +41,7 @@ class StateKeeperTest { fun testProperty() { class Foo(var value: String) - val keeper: StateKeeper = stateKeeper { + val keeper: StateKeeper = stateKeeper { _, _ -> add(Foo::value::get, Foo::value::set) } @@ -60,7 +60,7 @@ class StateKeeperTest { fun testArranger() { class Foo(var value: String) - val keeper: StateKeeper = stateKeeper { + val keeper: StateKeeper = stateKeeper { _, _ -> add(Foo::value::get, Foo::value::set) { "Baz" } } @@ -80,12 +80,12 @@ class StateKeeperTest { open class Foo(var one: String) class Bar(one: String, var two: Int) : Foo(one) - val fooKeeper: StateKeeper = stateKeeper { + val fooKeeper: StateKeeper = stateKeeper { _, _ -> add(Foo::one::get, Foo::one::set) } - val barKeeper: StateKeeper = stateKeeper { - add(fooKeeper) + val barKeeper: StateKeeper = stateKeeper { _, context -> + add(fooKeeper, context) add(Bar::two::get, Bar::two::set) } @@ -109,9 +109,9 @@ class StateKeeperTest { data class Foo(var value: Int) class Bar(var one: String, var foo: Foo) - val barKeeper: StateKeeper = stateKeeper { bar -> + val barKeeper: StateKeeper = stateKeeper { bar, context -> add(Bar::one::get, Bar::one::set) - entity(bar.foo) { + entity(bar.foo, context) { _, _ -> add(Foo::value::get, Foo::value::set) } } @@ -140,13 +140,13 @@ class StateKeeperTest { data class Foo(var value: Int) class Bar(var one: String, var foo: Foo) - val fooKeeper: StateKeeper = stateKeeper { + val fooKeeper: StateKeeper = stateKeeper { _, _ -> add(Foo::value::get, Foo::value::set) } - val barKeeper: StateKeeper = stateKeeper { bar -> + val barKeeper: StateKeeper = stateKeeper { bar, context -> add(Bar::one::get, Bar::one::set) - entity(bar.foo, fooKeeper) + entity(bar.foo, fooKeeper, context) } val foo = Foo(1) @@ -173,9 +173,9 @@ class StateKeeperTest { data class Foo(var value: Int) class Bar(var one: String, var foos: List) - val barKeeper: StateKeeper = stateKeeper { bar -> + val barKeeper: StateKeeper = stateKeeper { bar, context -> add(Bar::one::get, Bar::one::set) - entityList(bar.foos) { + entityList(bar.foos, context) { _, _ -> add(Foo::value::get, Foo::value::set) { it + 1 } } } @@ -209,7 +209,7 @@ class StateKeeperTest { fun testPostProcess() { class Foo(var value: String) - val keeper: StateKeeper = stateKeeper { foo -> + val keeper: StateKeeper = stateKeeper { foo, _ -> add(Foo::value::get, Foo::value::set) { "Baz" } postProcess { @@ -229,8 +229,8 @@ class StateKeeperTest { } } -private fun StateKeeper.withRestoration(owner: T, block: () -> Unit) { - val state = prepare(owner) +private fun StateKeeper.withRestoration(owner: T, block: () -> Unit) { + val state = prepare(owner, Unit) state.postProcess() block() state.restore()