FIR: store all annotationContainers in checker context

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -34953,6 +34953,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnFileWithVeryExperimentalMarker.kt");
             }
 
+            @Test
+            @TestMetadata("useExperimentalOnStatement.kt")
+            public void testUseExperimentalOnStatement() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnStatement.kt");
+            }
+
             @Test
             @TestMetadata("useExperimentalOnWholeModule.kt")
             public void testUseExperimentalOnWholeModule() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -34953,6 +34953,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnFileWithVeryExperimentalMarker.kt");
             }
 
+            @Test
+            @TestMetadata("useExperimentalOnStatement.kt")
+            public void testUseExperimentalOnStatement() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnStatement.kt");
+            }
+
             @Test
             @TestMetadata("useExperimentalOnWholeModule.kt")
             public void testUseExperimentalOnWholeModule() throws Exception {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/context/CheckerContext.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.analysis.checkers.context
 
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
@@ -28,6 +29,7 @@ abstract class CheckerContext {
     abstract val containingDeclarations: List<FirDeclaration>
     abstract val qualifiedAccessOrAnnotationCalls: List<FirStatement>
     abstract val getClassCalls: List<FirGetClassCall>
+    abstract val annotationContainers: List<FirAnnotationContainer>
 
     // Suppress
     abstract val suppressedDiagnostics: Set<String>
@@ -59,6 +61,10 @@ abstract class CheckerContext {
 
     abstract fun dropGetClassCall()
 
+    abstract fun addAnnotationContainer(annotationContainer: FirAnnotationContainer): CheckerContext
+
+    abstract fun dropAnnotationContainer()
+
     fun isDiagnosticSuppressed(diagnostic: FirDiagnostic): Boolean {
         val factory = diagnostic.factory
         val name = factory.name

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/context/MutableCheckerContext.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.analysis.checkers.context
 
 import kotlinx.collections.immutable.PersistentSet
 import kotlinx.collections.immutable.persistentSetOf
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
@@ -21,6 +22,7 @@ class MutableCheckerContext private constructor(
     override val containingDeclarations: MutableList<FirDeclaration>,
     override val qualifiedAccessOrAnnotationCalls: MutableList<FirStatement>,
     override val getClassCalls: MutableList<FirGetClassCall>,
+    override val annotationContainers: MutableList<FirAnnotationContainer>,
     sessionHolder: SessionHolder,
     returnTypeCalculator: ReturnTypeCalculator,
     override val suppressedDiagnostics: PersistentSet<String>,
@@ -33,6 +35,7 @@ class MutableCheckerContext private constructor(
         mutableListOf(),
         mutableListOf(),
         mutableListOf(),
+        mutableListOf(),
         sessionHolder,
         returnTypeCalculator,
         persistentSetOf(),
@@ -47,6 +50,7 @@ class MutableCheckerContext private constructor(
             containingDeclarations,
             qualifiedAccessOrAnnotationCalls,
             getClassCalls,
+            annotationContainers,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics,
@@ -83,6 +87,15 @@ class MutableCheckerContext private constructor(
         getClassCalls.removeAt(getClassCalls.size - 1)
     }
 
+    override fun addAnnotationContainer(annotationContainer: FirAnnotationContainer): CheckerContext {
+        annotationContainers.add(annotationContainer)
+        return this
+    }
+
+    override fun dropAnnotationContainer() {
+        annotationContainers.removeAt(annotationContainers.size - 1)
+    }
+
     override fun addSuppressedDiagnostics(
         diagnosticNames: Collection<String>,
         allInfosSuppressed: Boolean,
@@ -95,6 +108,7 @@ class MutableCheckerContext private constructor(
             containingDeclarations,
             qualifiedAccessOrAnnotationCalls,
             getClassCalls,
+            annotationContainers,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics.addAll(diagnosticNames),

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/context/PersistentCheckerContext.kt

@@ -9,6 +9,7 @@ import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.PersistentSet
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.persistentSetOf
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
@@ -23,6 +24,7 @@ class PersistentCheckerContext private constructor(
     override val containingDeclarations: PersistentList<FirDeclaration>,
     override val qualifiedAccessOrAnnotationCalls: PersistentList<FirStatement>,
     override val getClassCalls: PersistentList<FirGetClassCall>,
+    override val annotationContainers: PersistentList<FirAnnotationContainer>,
     sessionHolder: SessionHolder,
     returnTypeCalculator: ReturnTypeCalculator,
     override val suppressedDiagnostics: PersistentSet<String>,
@@ -35,6 +37,7 @@ class PersistentCheckerContext private constructor(
         persistentListOf(),
         persistentListOf(),
         persistentListOf(),
+        persistentListOf(),
         sessionHolder,
         returnTypeCalculator,
         persistentSetOf(),
@@ -49,6 +52,7 @@ class PersistentCheckerContext private constructor(
             containingDeclarations,
             qualifiedAccessOrAnnotationCalls,
             getClassCalls,
+            annotationContainers,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics,
@@ -64,6 +68,7 @@ class PersistentCheckerContext private constructor(
             containingDeclarations.add(declaration),
             qualifiedAccessOrAnnotationCalls,
             getClassCalls,
+            annotationContainers,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics,
@@ -82,6 +87,7 @@ class PersistentCheckerContext private constructor(
             containingDeclarations,
             this.qualifiedAccessOrAnnotationCalls.add(qualifiedAccessOrAnnotationCall),
             getClassCalls,
+            annotationContainers,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics,
@@ -100,6 +106,7 @@ class PersistentCheckerContext private constructor(
             containingDeclarations,
             qualifiedAccessOrAnnotationCalls,
             getClassCalls.add(getClassCall),
+            annotationContainers,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics,
@@ -112,6 +119,25 @@ class PersistentCheckerContext private constructor(
     override fun dropGetClassCall() {
     }
 
+    override fun addAnnotationContainer(annotationContainer: FirAnnotationContainer): PersistentCheckerContext {
+        return PersistentCheckerContext(
+            implicitReceiverStack,
+            containingDeclarations,
+            qualifiedAccessOrAnnotationCalls,
+            getClassCalls,
+            annotationContainers.add(annotationContainer),
+            sessionHolder,
+            returnTypeCalculator,
+            suppressedDiagnostics,
+            allInfosSuppressed,
+            allWarningsSuppressed,
+            allErrorsSuppressed
+        )
+    }
+
+    override fun dropAnnotationContainer() {
+    }
+
     override fun addSuppressedDiagnostics(
         diagnosticNames: Collection<String>,
         allInfosSuppressed: Boolean,
@@ -124,6 +150,7 @@ class PersistentCheckerContext private constructor(
             containingDeclarations,
             qualifiedAccessOrAnnotationCalls,
             getClassCalls,
+            annotationContainers,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics.addAll(diagnosticNames),

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirOptInUsageBaseChecker.kt

@@ -195,7 +195,7 @@ internal object FirOptInUsageBaseChecker {
         reporter: DiagnosticReporter
     ) {
         for ((annotationFqName, severity, message) in experimentalities) {
-            if (!isExperimentalityAcceptableInContext(annotationFqName, element, context)) {
+            if (!isExperimentalityAcceptableInContext(annotationFqName, context)) {
                 val diagnostic = when (severity) {
                     Experimentality.Severity.WARNING -> FirErrors.EXPERIMENTAL_API_USAGE
                     Experimentality.Severity.ERROR -> FirErrors.EXPERIMENTAL_API_USAGE_ERROR
@@ -211,7 +211,6 @@ internal object FirOptInUsageBaseChecker {
 
     private fun isExperimentalityAcceptableInContext(
         annotationFqName: FqName,
-        element: FirElement,
         context: CheckerContext
     ): Boolean {
         val languageVersionSettings = context.session.languageVersionSettings
@@ -219,19 +218,12 @@ internal object FirOptInUsageBaseChecker {
         if (fqNameAsString in languageVersionSettings.getFlag(AnalysisFlags.useExperimental)) {
             return true
         }
-        for (declaration in context.containingDeclarations) {
+        for (annotationContainer in context.annotationContainers) {
-            if (declaration !is FirAnnotatedDeclaration) continue
+            if (annotationContainer.isExperimentalityAcceptable(annotationFqName)) {
-            if (declaration.isExperimentalityAcceptable(annotationFqName)) {
                 return true
             }
         }
-        for (accessOrAnnotation in context.qualifiedAccessOrAnnotationCalls) {
+        return false
-            if (accessOrAnnotation.isExperimentalityAcceptable(annotationFqName)) {
-                return true
-            }
-        }
-        if (element !is FirAnnotationContainer) return false
-        return element.isExperimentalityAcceptable(annotationFqName)
     }
 
     private fun FirAnnotationContainer.isExperimentalityAcceptable(annotationFqName: FqName): Boolean {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollectorVisitor.kt

@@ -43,14 +43,14 @@ abstract class AbstractDiagnosticCollectorVisitor(
     }
 
     override fun visitAnnotationContainer(annotationContainer: FirAnnotationContainer, data: Nothing?) {
-        withSuppressedDiagnostics(annotationContainer) {
+        withAnnotationContainer(annotationContainer) {
             checkElement(annotationContainer)
             visitNestedElements(annotationContainer)
         }
     }
 
     private fun visitJump(loopJump: FirLoopJump) {
-        withSuppressedDiagnostics(loopJump) {
+        withAnnotationContainer(loopJump) {
             checkElement(loopJump)
             loopJump.target.labeledElement.takeIf { it is FirErrorLoop }?.accept(this, null)
         }
@@ -72,7 +72,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
     }
 
     override fun visitRegularClass(regularClass: FirRegularClass, data: Nothing?) {
-        withSuppressedDiagnostics(regularClass) {
+        withAnnotationContainer(regularClass) {
             visitClassAndChildren(regularClass, regularClass.defaultType())
         }
     }
@@ -82,19 +82,19 @@ abstract class AbstractDiagnosticCollectorVisitor(
     }
 
     override fun visitAnonymousObject(anonymousObject: FirAnonymousObject, data: Nothing?) {
-        withSuppressedDiagnostics(anonymousObject) {
+        withAnnotationContainer(anonymousObject) {
             visitClassAndChildren(anonymousObject, anonymousObject.defaultType())
         }
     }
 
     override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?) {
-        withSuppressedDiagnostics(simpleFunction) {
+        withAnnotationContainer(simpleFunction) {
             visitWithDeclarationAndReceiver(simpleFunction, simpleFunction.name, simpleFunction.receiverTypeRef)
         }
     }
 
     override fun visitConstructor(constructor: FirConstructor, data: Nothing?) {
-        withSuppressedDiagnostics(constructor) {
+        withAnnotationContainer(constructor) {
             visitWithDeclaration(constructor)
         }
     }
@@ -104,7 +104,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
     }
 
     override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Nothing?) {
-        withSuppressedDiagnostics(anonymousFunction) {
+        withAnnotationContainer(anonymousFunction) {
             val labelName = anonymousFunction.label?.name?.let { Name.identifier(it) }
             visitWithDeclarationAndReceiver(
                 anonymousFunction,
@@ -115,13 +115,13 @@ abstract class AbstractDiagnosticCollectorVisitor(
     }
 
     override fun visitProperty(property: FirProperty, data: Nothing?) {
-        withSuppressedDiagnostics(property) {
+        withAnnotationContainer(property) {
             visitWithDeclaration(property)
         }
     }
 
     override fun visitTypeAlias(typeAlias: FirTypeAlias, data: Nothing?) {
-        withSuppressedDiagnostics(typeAlias) {
+        withAnnotationContainer(typeAlias) {
             visitWithDeclaration(typeAlias)
         }
     }
@@ -129,26 +129,26 @@ abstract class AbstractDiagnosticCollectorVisitor(
     override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: Nothing?) {
         if (propertyAccessor !is FirDefaultPropertyAccessor) {
             val property = context.containingDeclarations.last() as FirProperty
-            withSuppressedDiagnostics(propertyAccessor) {
+            withAnnotationContainer(propertyAccessor) {
                 visitWithDeclarationAndReceiver(propertyAccessor, property.name, property.receiverTypeRef)
             }
         }
     }
 
     override fun visitValueParameter(valueParameter: FirValueParameter, data: Nothing?) {
-        withSuppressedDiagnostics(valueParameter) {
+        withAnnotationContainer(valueParameter) {
             visitWithDeclaration(valueParameter)
         }
     }
 
     override fun visitEnumEntry(enumEntry: FirEnumEntry, data: Nothing?) {
-        withSuppressedDiagnostics(enumEntry) {
+        withAnnotationContainer(enumEntry) {
             visitWithDeclaration(enumEntry)
         }
     }
 
     override fun visitFile(file: FirFile, data: Nothing?) {
-        withSuppressedDiagnostics(file) {
+        withAnnotationContainer(file) {
             visitWithDeclaration(file)
         }
     }
@@ -158,14 +158,14 @@ abstract class AbstractDiagnosticCollectorVisitor(
     }
 
     override fun visitBlock(block: FirBlock, data: Nothing?) {
-        withSuppressedDiagnostics(block) {
+        withAnnotationContainer(block) {
             visitExpression(block, data)
         }
     }
 
     override fun visitTypeRef(typeRef: FirTypeRef, data: Nothing?) {
         if (typeRef.source != null && typeRef.source?.kind !is FirFakeSourceElementKind) {
-            withSuppressedDiagnostics(typeRef) {
+            withAnnotationContainer(typeRef) {
                 checkElement(typeRef)
                 visitNestedElements(typeRef)
             }
@@ -185,7 +185,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
         //if we don't visit resolved type we can't make any diagnostics on them
         //so here we check resolvedTypeRef
         if (resolvedTypeRef.type !is ConeClassErrorType) {
-            withSuppressedDiagnostics(resolvedTypeRef) {
+            withAnnotationContainer(resolvedTypeRef) {
                 checkElement(resolvedTypeRef)
             }
         }
@@ -312,12 +312,18 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
 
     @OptIn(PrivateForInline::class)
-    inline fun <R> withSuppressedDiagnostics(annotationContainer: FirAnnotationContainer, block: () -> R): R {
+    inline fun <R> withAnnotationContainer(annotationContainer: FirAnnotationContainer, block: () -> R): R {
         val existingContext = context
         addSuppressedDiagnosticsToContext(annotationContainer)
+        if (annotationContainer.annotations.isNotEmpty()) {
+            context = context.addAnnotationContainer(annotationContainer)
+        }
         return try {
             block()
         } finally {
+            if (annotationContainer.annotations.isNotEmpty()) {
+                existingContext.dropAnnotationContainer()
+            }
             context = existingContext
         }
     }

compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnStatement.kt

@@ -0,0 +1,19 @@
+// FIR_IDENTICAL
+// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
+
+@RequiresOptIn(level = RequiresOptIn.Level.ERROR)
+@Retention(AnnotationRetention.BINARY)
+annotation class ExperimentalAPI
+
+@ExperimentalAPI
+fun function(): String = ""
+
+fun use(): String {
+    @OptIn(ExperimentalAPI::class)
+    for (i in 1..2) {
+        function()
+    }
+
+    @OptIn(ExperimentalAPI::class)
+    return function()
+}

compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnStatement.txt

@@ -0,0 +1,11 @@
+package
+
+@ExperimentalAPI public fun function(): kotlin.String
+public fun use(): kotlin.String
+
+@kotlin.RequiresOptIn(level = Level.ERROR) @kotlin.annotation.Retention(value = AnnotationRetention.BINARY) public final annotation class ExperimentalAPI : kotlin.Annotation {
+    public constructor ExperimentalAPI()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -35049,6 +35049,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnFileWithVeryExperimentalMarker.kt");
             }
 
+            @Test
+            @TestMetadata("useExperimentalOnStatement.kt")
+            public void testUseExperimentalOnStatement() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnStatement.kt");
+            }
+
             @Test
             @TestMetadata("useExperimentalOnWholeModule.kt")
             public void testUseExperimentalOnWholeModule() throws Exception {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FileStructureElementDiagnosticRetriever.kt

@@ -109,7 +109,7 @@ internal object FileDiagnosticRetriever : FileStructureElementDiagnosticRetrieve
         components,
     ) {
         override fun visitFile(file: FirFile, data: Nothing?) {
-            withSuppressedDiagnostics(file) {
+            withAnnotationContainer(file) {
                 visitWithDeclaration(file) {
                     file.annotations.forEach { it.accept(this, data) }
                     file.packageDirective.accept(this, data)

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -34953,6 +34953,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnFileWithVeryExperimentalMarker.kt");
             }
 
+            @Test
+            @TestMetadata("useExperimentalOnStatement.kt")
+            public void testUseExperimentalOnStatement() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/experimental/useExperimentalOnStatement.kt");
+            }
+
             @Test
             @TestMetadata("useExperimentalOnWholeModule.kt")
             public void testUseExperimentalOnWholeModule() throws Exception {