[LL FIR] StateKeeper: provide context

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

^KT-60552

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) {
+    protected fun <T : FirElementWithResolveState> resolve(target: T, keeper: StateKeeper<T, FirDesignationWithFile>) {
-        val firDesignation = FirDesignationWithFile(nestedClassesStack, declaration, resolveTarget.firFile)
+        val firDesignation = FirDesignationWithFile(nestedClassesStack, target, resolveTarget.firFile)
-        FirLazyBodiesCalculator.calculateBodies(firDesignation)
+        resolveWithKeeper(target, firDesignation, keeper, { FirLazyBodiesCalculator.calculateBodies(firDesignation) }) {
-    }
-
-    protected fun <T : FirElementWithResolveState> resolve(target: T, keeper: StateKeeper<T>) {
-        resolveWithKeeper(target, keeper, ::calculateLazyBodies) {
             rawResolve(target)
         }
     }

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<FirAnnotation> = stateKeeper {
+    private val ANNOTATION: StateKeeper<FirAnnotation, FirSession> = stateKeeper { _, session ->
-        add(ANNOTATION_BASE)
+        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
                 }
 

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<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 {
+    val CONSTRUCTOR: StateKeeper<FirConstructor, FirDesignationWithFile> = stateKeeper { _, designation ->
-        add(FUNCTION)
+        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 {
+    val FIELD: StateKeeper<FirField, FirDesignationWithFile> = stateKeeper { _, designation ->
-        add(VARIABLE)
+        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.getterIfUnresolved, FUNCTION, designation)
-        entity(property.setterIfUnresolved, FUNCTION)
+        entity(property.setterIfUnresolved, FUNCTION, designation)
-        entity(property.backingFieldIfUnresolved, VARIABLE)
+        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

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<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 {
+    val SIMPLE_FUNCTION: StateKeeper<FirSimpleFunction, FirDesignationWithFile> = stateKeeper { _, designation ->
-        add(CONTRACT_DESCRIPTION_OWNER)
+        add(CONTRACT_DESCRIPTION_OWNER, designation)
-        add(BODY_OWNER)
+        add(BODY_OWNER, designation)
     }
 
-    val CONSTRUCTOR: StateKeeper<FirConstructor> = stateKeeper {
+    val CONSTRUCTOR: StateKeeper<FirConstructor, FirDesignationWithFile> = stateKeeper { _, designation ->
-        add(CONTRACT_DESCRIPTION_OWNER)
+        add(CONTRACT_DESCRIPTION_OWNER, designation)
-        add(BODY_OWNER)
+        add(BODY_OWNER, designation)
     }
 
-    val PROPERTY_ACCESSOR: StateKeeper<FirPropertyAccessor> = stateKeeper {
+    val PROPERTY_ACCESSOR: StateKeeper<FirPropertyAccessor, FirDesignationWithFile> = stateKeeper { _, designation ->
-        add(CONTRACT_DESCRIPTION_OWNER)
+        add(CONTRACT_DESCRIPTION_OWNER, designation)
-        add(BODY_OWNER)
+        add(BODY_OWNER, designation)
     }
 
-    val PROPERTY: StateKeeper<FirProperty> = stateKeeper { property ->
+    val PROPERTY: StateKeeper<FirProperty, FirDesignationWithFile> = stateKeeper { property, designation ->
-        entity(property.getter, PROPERTY_ACCESSOR)
+        entity(property.getter, PROPERTY_ACCESSOR, designation)
-        entity(property.setter, PROPERTY_ACCESSOR)
+        entity(property.setter, PROPERTY_ACCESSOR, designation)
     }
 }

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<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> {
+internal fun <Owner : Any, Context : Any> stateKeeper(block: context(StateKeeperBuilder) StateKeeperScope<Owner, Context>.(Owner, Context) -> Unit): StateKeeper<Owner, Context> {
-    return StateKeeper { owner ->
+    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)
+        val scope = StateKeeperScope<Owner, Context>(owner)
-        block(builder, scope, 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 {
+    fun prepare(owner: Owner, context: Context): PreservedState {
-        return provider(owner)
+        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>,
+    context: Context,
-    prepareTarget: (Target) -> Unit,
+    keeper: StateKeeper<Target, Context>,
-    action: () -> Result
+    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()

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<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 {
+        val barKeeper: StateKeeper<Bar, Unit> = stateKeeper { _, context ->
-            add(fooKeeper)
+            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) {
+private fun <T : Any> StateKeeper<T, Unit>.withRestoration(owner: T, block: () -> Unit) {
-    val state = prepare(owner)
+    val state = prepare(owner, Unit)
     state.postProcess()
     block()
     state.restore()