[LL FIR] StateKeeper: provide context

We need this to be able to construct new elements on the fly
during arrangement

^KT-60552
This commit is contained in:
Dmitrii Gridin
2023-07-25 17:21:08 +02:00
committed by Space Team
parent 38ef5be217
commit 65fec8dc19
6 changed files with 93 additions and 84 deletions
@@ -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 <T : FirElementWithResolveState> resolve(target: T, keeper: StateKeeper<T>) {
resolveWithKeeper(target, keeper, ::calculateLazyBodies) {
protected fun <T : FirElementWithResolveState> resolve(target: T, keeper: StateKeeper<T, FirDesignationWithFile>) {
val firDesignation = FirDesignationWithFile(nestedClassesStack, target, resolveTarget.firFile)
resolveWithKeeper(target, firDesignation, keeper, { FirLazyBodiesCalculator.calculateBodies(firDesignation) }) {
rawResolve(target)
}
}
@@ -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<FirAnnotation> = stateKeeper {
add(ANNOTATION_BASE)
private val ANNOTATION: StateKeeper<FirAnnotation, FirSession> = stateKeeper { _, session ->
add(ANNOTATION_BASE, session)
add(FirAnnotation::argumentMapping, FirAnnotation::replaceArgumentMapping)
add(FirAnnotation::typeArgumentsCopied, FirAnnotation::replaceTypeArguments)
}
private val ANNOTATION_BASE: StateKeeper<FirAnnotation> = stateKeeper { annotation ->
private val ANNOTATION_BASE: StateKeeper<FirAnnotation, FirSession> = stateKeeper { annotation, session ->
if (annotation is FirAnnotationCall) {
entity(annotation, ANNOTATION_CALL)
entity(annotation, ANNOTATION_CALL, session)
}
}
private val ANNOTATION_CALL: StateKeeper<FirAnnotationCall> = stateKeeper { annotationCall ->
private val ANNOTATION_CALL: StateKeeper<FirAnnotationCall, FirSession> = stateKeeper { annotationCall, _ ->
add(FirAnnotationCall::calleeReference, FirAnnotationCall::replaceCalleeReference)
val argumentList = annotationCall.argumentList
@@ -100,12 +101,12 @@ internal object AnnotationArgumentMappingStateKeepers {
}
}
val DECLARATION: StateKeeper<FirElementWithResolveState> = stateKeeper { target ->
val DECLARATION: StateKeeper<FirElementWithResolveState, FirSession> = 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
}
@@ -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<FirScript> = stateKeeper {
val SCRIPT: StateKeeper<FirScript, FirDesignationWithFile> = stateKeeper { _, _ ->
// TODO Lazy body is not supported for scripts yet
}
val ANONYMOUS_INITIALIZER: StateKeeper<FirAnonymousInitializer> = stateKeeper {
val ANONYMOUS_INITIALIZER: StateKeeper<FirAnonymousInitializer, FirDesignationWithFile> = stateKeeper { _, _ ->
add(FirAnonymousInitializer::body, FirAnonymousInitializer::replaceBody, ::blockGuard)
add(FirAnonymousInitializer::controlFlowGraphReference, FirAnonymousInitializer::replaceControlFlowGraphReference)
}
val FUNCTION: StateKeeper<FirFunction> = stateKeeper { function ->
val FUNCTION: StateKeeper<FirFunction, FirDesignationWithFile> = 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<FirConstructor> = stateKeeper {
add(FUNCTION)
val CONSTRUCTOR: StateKeeper<FirConstructor, FirDesignationWithFile> = stateKeeper { _, designation ->
add(FUNCTION, designation)
add(FirConstructor::delegatedConstructor, FirConstructor::replaceDelegatedConstructor, ::delegatedConstructorCallGuard)
}
val VARIABLE: StateKeeper<FirVariable> = stateKeeper { variable ->
val VARIABLE: StateKeeper<FirVariable, FirDesignationWithFile> = stateKeeper { variable, _ ->
add(FirVariable::returnTypeRef, FirVariable::replaceReturnTypeRef)
if (!isCallableWithSpecialBody(variable)) {
@@ -194,7 +195,7 @@ internal object BodyStateKeepers {
}
}
private val VALUE_PARAMETER: StateKeeper<FirValueParameter> = stateKeeper { valueParameter ->
private val VALUE_PARAMETER: StateKeeper<FirValueParameter, FirDesignationWithFile> = 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<FirField> = stateKeeper {
add(VARIABLE)
val FIELD: StateKeeper<FirField, FirDesignationWithFile> = stateKeeper { _, designation ->
add(VARIABLE, designation)
add(FirField::controlFlowGraphReference, FirField::replaceControlFlowGraphReference)
}
val PROPERTY: StateKeeper<FirProperty> = stateKeeper { property ->
val PROPERTY: StateKeeper<FirProperty, FirDesignationWithFile> = 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<FirFunction>.preserveContractBlock(function: FirFunction) {
private fun StateKeeperScope<FirFunction, FirDesignationWithFile>.preserveContractBlock(function: FirFunction) {
val oldBody = function.body
if (oldBody == null || oldBody is FirLazyBlock) {
return
@@ -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<FirContractDescriptionOwner> = stateKeeper {
private val CONTRACT_DESCRIPTION_OWNER: StateKeeper<FirContractDescriptionOwner, FirDesignationWithFile> = stateKeeper { _, _ ->
add(FirContractDescriptionOwner::contractDescription, FirContractDescriptionOwner::replaceContractDescription)
}
private val BODY_OWNER: StateKeeper<FirFunction> = stateKeeper { declaration ->
private val BODY_OWNER: StateKeeper<FirFunction, FirDesignationWithFile> = 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<FirSimpleFunction> = stateKeeper {
add(CONTRACT_DESCRIPTION_OWNER)
add(BODY_OWNER)
val SIMPLE_FUNCTION: StateKeeper<FirSimpleFunction, FirDesignationWithFile> = stateKeeper { _, designation ->
add(CONTRACT_DESCRIPTION_OWNER, designation)
add(BODY_OWNER, designation)
}
val CONSTRUCTOR: StateKeeper<FirConstructor> = stateKeeper {
add(CONTRACT_DESCRIPTION_OWNER)
add(BODY_OWNER)
val CONSTRUCTOR: StateKeeper<FirConstructor, FirDesignationWithFile> = stateKeeper { _, designation ->
add(CONTRACT_DESCRIPTION_OWNER, designation)
add(BODY_OWNER, designation)
}
val PROPERTY_ACCESSOR: StateKeeper<FirPropertyAccessor> = stateKeeper {
add(CONTRACT_DESCRIPTION_OWNER)
add(BODY_OWNER)
val PROPERTY_ACCESSOR: StateKeeper<FirPropertyAccessor, FirDesignationWithFile> = stateKeeper { _, designation ->
add(CONTRACT_DESCRIPTION_OWNER, designation)
add(BODY_OWNER, designation)
}
val PROPERTY: StateKeeper<FirProperty> = stateKeeper { property ->
entity(property.getter, PROPERTY_ACCESSOR)
entity(property.setter, PROPERTY_ACCESSOR)
val PROPERTY: StateKeeper<FirProperty, FirDesignationWithFile> = stateKeeper { property, designation ->
entity(property.getter, PROPERTY_ACCESSOR, designation)
entity(property.setter, PROPERTY_ACCESSOR, designation)
}
}
@@ -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<Owner : Any>(private val owner: Owner) {
internal value class StateKeeperScope<Owner : Any, Context : Any>(private val owner: Owner) {
/**
* Defines an entity state preservation rule.
*
@@ -77,9 +78,9 @@ internal value class StateKeeperScope<Owner : Any>(private val owner: Owner) {
* In other words, applies all rules defined in the [keeper] to the current entity.
*/
context(StateKeeperBuilder)
fun add(keeper: StateKeeper<Owner>) {
fun add(keeper: StateKeeper<Owner, Context>, context: Context) {
val owner = this@StateKeeperScope.owner
register(keeper.prepare(owner))
register(keeper.prepare(owner, context))
}
/**
@@ -99,9 +100,9 @@ internal value class StateKeeperScope<Owner : Any>(private val owner: Owner) {
* Does nothing if the [entity] is `null`.
*/
context(StateKeeperBuilder)
internal fun <Entity : Any> entity(entity: Entity?, keeper: StateKeeper<Entity>) {
internal fun <Entity : Any, Context : Any> entity(entity: Entity?, keeper: StateKeeper<Entity, Context>, context: Context) {
if (entity != null) {
StateKeeperScope(entity).add(keeper)
StateKeeperScope<Entity, Context>(entity).add(keeper, context)
}
}
@@ -110,9 +111,13 @@ internal fun <Entity : Any> entity(entity: Entity?, keeper: StateKeeper<Entity>)
* Does nothing if the [entity] is `null`.
*/
context(StateKeeperBuilder)
internal inline fun <Entity : Any> entity(entity: Entity?, block: StateKeeperScope<Entity>.(Entity) -> Unit) {
internal inline fun <Entity : Any, Context : Any> entity(
entity: Entity?,
context: Context,
block: StateKeeperScope<Entity, Context>.(Entity, Context) -> Unit,
) {
if (entity != null) {
StateKeeperScope(entity).block(entity)
StateKeeperScope<Entity, Context>(entity).block(entity, context)
}
}
@@ -122,11 +127,11 @@ internal inline fun <Entity : Any> entity(entity: Entity?, block: StateKeeperSco
* Skips `null` elements in the [list].
*/
context(StateKeeperBuilder)
internal fun <Entity : Any> entityList(list: List<Entity?>?, keeper: StateKeeper<Entity>) {
internal fun <Entity : Any, Context : Any> entityList(list: List<Entity?>?, keeper: StateKeeper<Entity, Context>, context: Context) {
if (list != null) {
for (entity in list) {
if (entity != null) {
StateKeeperScope(entity).add(keeper)
StateKeeperScope<Entity, Context>(entity).add(keeper, context)
}
}
}
@@ -138,11 +143,15 @@ internal fun <Entity : Any> entityList(list: List<Entity?>?, keeper: StateKeeper
* Skips `null` elements in the [list].
*/
context(StateKeeperBuilder)
internal inline fun <Entity : Any> entityList(list: List<Entity?>?, block: StateKeeperScope<Entity>.(Entity) -> Unit) {
internal inline fun <Entity : Any, Context : Any> entityList(
list: List<Entity?>?,
context: Context,
block: StateKeeperScope<Entity, Context>.(Entity, Context) -> Unit,
) {
if (list != null) {
for (entity in list) {
if (entity != null) {
StateKeeperScope(entity).block(entity)
StateKeeperScope<Entity, Context>(entity).block(entity, context)
}
}
}
@@ -156,8 +165,8 @@ internal inline fun <Entity : Any> entityList(list: List<Entity?>?, block: State
* 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 ->
internal fun <Owner : Any, Context : Any> stateKeeper(block: context(StateKeeperBuilder) StateKeeperScope<Owner, Context>.(Owner, Context) -> Unit): StateKeeper<Owner, Context> {
return StateKeeper { owner, context ->
val states = mutableListOf<PreservedState>()
val builder = object : StateKeeperBuilder {
@@ -166,8 +175,8 @@ internal fun <Owner : Any> stateKeeper(block: context(StateKeeperBuilder) StateK
}
}
val scope = StateKeeperScope(owner)
block(builder, scope, owner)
val scope = StateKeeperScope<Owner, Context>(owner)
block(builder, scope, owner, context)
object : PreservedState {
private var isPostProcessed = false
@@ -222,13 +231,13 @@ internal fun <Owner : Any> stateKeeper(block: context(StateKeeperBuilder) StateK
* }
* ```
*/
internal class StateKeeper<in Owner : Any>(val provider: (Owner) -> PreservedState) {
internal class StateKeeper<in Owner : Any, Context : Any>(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 <Target : FirElementWithResolveState, Result> resolveWithKeeper(
internal inline fun <Target : FirElementWithResolveState, Context : Any, Result> resolveWithKeeper(
target: Target,
keeper: StateKeeper<Target>,
prepareTarget: (Target) -> Unit,
action: () -> Result
context: Context,
keeper: StateKeeper<Target, Context>,
prepareTarget: (Target) -> Unit = {},
action: () -> Result,
): Result {
if (!USE_STATE_KEEPER) {
prepareTarget(target)
@@ -274,7 +284,7 @@ internal inline fun <Target : FirElementWithResolveState, Result> resolveWithKee
var preservedState: PreservedState? = null
try {
preservedState = keeper.prepare(target)
preservedState = keeper.prepare(target, context)
prepareTarget(target)
preservedState.postProcess()
return action()
@@ -22,7 +22,7 @@ class StateKeeperTest {
}
}
val keeper: StateKeeper<Foo> = stateKeeper {
val keeper: StateKeeper<Foo, Unit> = stateKeeper { _, _ ->
add(Foo::getValue, Foo::setValue)
}
@@ -41,7 +41,7 @@ class StateKeeperTest {
fun testProperty() {
class Foo(var value: String)
val keeper: StateKeeper<Foo> = stateKeeper {
val keeper: StateKeeper<Foo, Unit> = stateKeeper { _, _ ->
add(Foo::value::get, Foo::value::set)
}
@@ -60,7 +60,7 @@ class StateKeeperTest {
fun testArranger() {
class Foo(var value: String)
val keeper: StateKeeper<Foo> = stateKeeper {
val keeper: StateKeeper<Foo, Unit> = 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<Foo> = stateKeeper {
val fooKeeper: StateKeeper<Foo, Unit> = stateKeeper { _, _ ->
add(Foo::one::get, Foo::one::set)
}
val barKeeper: StateKeeper<Bar> = stateKeeper {
add(fooKeeper)
val barKeeper: StateKeeper<Bar, Unit> = 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<Bar> = stateKeeper { bar ->
val barKeeper: StateKeeper<Bar, Unit> = 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<Foo> = stateKeeper {
val fooKeeper: StateKeeper<Foo, Unit> = stateKeeper { _, _ ->
add(Foo::value::get, Foo::value::set)
}
val barKeeper: StateKeeper<Bar> = stateKeeper { bar ->
val barKeeper: StateKeeper<Bar, Unit> = 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<Foo>)
val barKeeper: StateKeeper<Bar> = stateKeeper { bar ->
val barKeeper: StateKeeper<Bar, Unit> = 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<Foo> = stateKeeper { foo ->
val keeper: StateKeeper<Foo, Unit> = stateKeeper { foo, _ ->
add(Foo::value::get, Foo::value::set) { "Baz" }
postProcess {
@@ -229,8 +229,8 @@ class StateKeeperTest {
}
}
private fun <T : Any> StateKeeper<T>.withRestoration(owner: T, block: () -> Unit) {
val state = prepare(owner)
private fun <T : Any> StateKeeper<T, Unit>.withRestoration(owner: T, block: () -> Unit) {
val state = prepare(owner, Unit)
state.postProcess()
block()
state.restore()