From 3518aa58982a03fb3ca4dacd6c31d74707b56442 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Tue, 25 Jul 2023 19:54:26 +0200
Subject: [PATCH] [LL FIR] ANNOTATIONS_ARGUMENTS_MAPPING: drop intermediate
 state with lazy arguments

This is safer because we can be sure that no one can
read an intermediate result with lazy expressions (related to KT-60387)

^KT-60552 Fixed
---
 .../api/fir/lazy/resolve/FirLazyBodiesCalculator.kt |  4 ++++
 .../LLFirAnnotationArgumentMappingLazyResolver.kt   | 13 +++++++------
 2 files changed, 11 insertions(+), 6 deletions(-)

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 702e0af01bf..c337dc94ca5 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
@@ -57,6 +57,10 @@ internal object FirLazyBodiesCalculator {
 
     fun calculateLazyArgumentsForAnnotation(annotationCall: FirAnnotationCall, session: FirSession): FirArgumentList {
         require(needCalculatingAnnotationCall(annotationCall))
+        return createArgumentsForAnnotation(annotationCall, session)
+    }
+
+    fun createArgumentsForAnnotation(annotationCall: FirAnnotationCall, session: FirSession): FirArgumentList {
         val builder = PsiRawFirBuilder(session, baseScopeProvider = session.kotlinScopeProvider)
         val ktAnnotationEntry = annotationCall.psi as KtAnnotationEntry
         builder.context.packageFqName = ktAnnotationEntry.containingKtFile.packageFqName
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 f547d4f6680..ef07fc917ac 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,10 +7,9 @@ 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.FirLazyBodiesCalculator
 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.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
@@ -62,7 +61,7 @@ private class LLFirAnnotationArgumentsMappingTargetResolver(
     )
 
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
-        resolveWithKeeper(target, target.llFirSession, AnnotationArgumentMappingStateKeepers.DECLARATION, ::calculateAnnotations) {
+        resolveWithKeeper(target, target.llFirSession, AnnotationArgumentMappingStateKeepers.DECLARATION) {
             transformAnnotations(target)
         }
     }
@@ -81,19 +80,21 @@ internal object AnnotationArgumentMappingStateKeepers {
         }
     }
 
-    private val ANNOTATION_CALL: StateKeeper<FirAnnotationCall, FirSession> = stateKeeper { annotationCall, _ ->
+    private val ANNOTATION_CALL: StateKeeper<FirAnnotationCall, FirSession> = stateKeeper { annotationCall, session ->
         add(FirAnnotationCall::calleeReference, FirAnnotationCall::replaceCalleeReference)
 
         val argumentList = annotationCall.argumentList
         if (argumentList !is FirResolvedArgumentList && argumentList !is FirEmptyArgumentList) {
             add(FirAnnotationCall::argumentList, FirAnnotationCall::replaceArgumentList) { oldList ->
+                val newArguments = FirLazyBodiesCalculator.createArgumentsForAnnotation(annotationCall, session).arguments
                 buildArgumentList {
                     source = oldList.source
-                    for (argument in oldList.arguments) {
+                    for ((index, argument) in oldList.arguments.withIndex()) {
                         val replacement = when {
                             argument is FirPropertyAccessExpression && argument.calleeReference.isError() -> argument
-                            else -> expressionGuard(argument)
+                            else -> newArguments[index]
                         }
+
                         arguments.add(replacement)
                     }
                 }