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 134d4fc..e85751e 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
@@ -17,6 +17,7 @@
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -85,3 +86,19 @@
         }
     }
 }
+
+private object AnnotationArgumentMappingStateKeepers {
+    val ANNOTATION: StateKeeper<FirAnnotation> = stateKeeper {
+        add(FirAnnotation::typeRef, FirAnnotation::replaceTypeRef)
+        add(FirAnnotation::annotationTypeRef, FirAnnotation::replaceAnnotationTypeRef)
+        add(FirAnnotation::argumentMapping, FirAnnotation::replaceArgumentMapping)
+    }
+
+    val FUNCTION: StateKeeper<FirFunction> = stateKeeper { function ->
+        entityList(function.annotations, ANNOTATION)
+
+        for (valueParameter in function.valueParameters) {
+            entityList(valueParameter.annotations, ANNOTATION)
+        }
+    }
+}
\ 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 318e0aa..aab93bb 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
@@ -72,6 +72,17 @@
     }
 }
 
+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)
+            }
+        }
+    }
+}
+
 internal fun <Owner : Any> stateKeeper(block: context(StateKeeperBuilder) StateKeeperScope<Owner>.(Owner) -> Unit): StateKeeper<Owner> {
     return StateKeeper { owner ->
         val states = mutableListOf<PreservedState>()