[FIR] Cache files with plugin annotations

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt

@@ -11,7 +11,9 @@ import kotlinx.collections.immutable.PersistentList
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.extensions.predicate.*
+import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.fqName
 
 abstract class FirPredicateBasedProvider : FirSessionComponent {
@@ -39,6 +41,8 @@ abstract class FirPredicateBasedProvider : FirSessionComponent {
 
     abstract fun getOwnersOfDeclaration(declaration: FirAnnotatedDeclaration): List<FirAnnotatedDeclaration>?
 
+    abstract fun fileHasPluginAnnotations(file: FirFile): Boolean
+
     abstract fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>)
     abstract fun registerGeneratedDeclaration(declaration: FirAnnotatedDeclaration, owner: FirAnnotatedDeclaration)
 
@@ -78,6 +82,10 @@ private class FirPredicateBasedProviderImpl(private val session: FirSession) : F
         return this.map { it to cache.ownersForDeclaration.getValue(it) }
     }
 
+    override fun fileHasPluginAnnotations(file: FirFile): Boolean {
+        return file in cache.filesWithPluginAnnotations
+    }
+
     override fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>) {
         cache.ownersForDeclaration[declaration] = owners
         registerOwnersDeclarations(declaration, owners)
@@ -88,6 +96,8 @@ private class FirPredicateBasedProviderImpl(private val session: FirSession) : F
         if (matchingAnnotations.isEmpty()) return
         matchingAnnotations.forEach { cache.declarationByAnnotation.put(it, declaration) }
         cache.annotationsOfDeclaration.putAll(declaration, matchingAnnotations)
+        val file = owners.first() as FirFile
+        cache.filesWithPluginAnnotations += file
     }
 
     override fun getOwnersOfDeclaration(declaration: FirAnnotatedDeclaration): List<FirAnnotatedDeclaration>? {
@@ -172,6 +182,8 @@ private class FirPredicateBasedProviderImpl(private val session: FirSession) : F
         val parentAnnotationsOfDeclaration: LinkedHashMultimap<FirAnnotatedDeclaration, AnnotationFqn> = LinkedHashMultimap.create()
 
         val ownersForDeclaration: MutableMap<FirAnnotatedDeclaration, PersistentList<FirAnnotatedDeclaration>> = mutableMapOf()
+
+        val filesWithPluginAnnotations: MutableSet<FirFile> = mutableSetOf()
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsResolveProcessor.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.registeredPluginAnnotations
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -30,13 +31,14 @@ class FirAnnotationArgumentsResolveProcessor(
 class FirAnnotationArgumentsResolveTransformerAdapter(session: FirSession, scopeSession: ScopeSession) : FirTransformer<Nothing?>() {
     private val transformer = FirAnnotationArgumentsResolveTransformer(session, scopeSession)
     private val hasAnnotations = session.registeredPluginAnnotations.annotations.isNotEmpty()
+    private val predicateBasedProvider = session.predicateBasedProvider
 
     override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
         return element.compose()
     }
 
     override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirDeclaration> {
-        if (!hasAnnotations) return file.compose()
+        if (!hasAnnotations || !predicateBasedProvider.fileHasPluginAnnotations(file)) return file.compose()
         return file.transform(transformer, ResolutionMode.ContextIndependent)
     }
 }