K2: Set correct containingFile in CheckerContext

Previously, it was working for all the checkers but the file-level
DeclarationChecker's because when execution comes to
DeclarationCheck::check for a file, `containingDeclarations` is empty
(it doesn't contain the file itself yet), thus some things that rely on
CheckerContext::containingFile don't work properly.

For example, SimpleDiagnosticsCollectorWithSuppress.report that
effectively doesn't collect diagnostics when the containingFile is null.

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

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

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

@@ -61,8 +61,7 @@ abstract class CheckerContext : DiagnosticContext {
     override val languageVersionSettings: LanguageVersionSettings
         get() = session.languageVersionSettings
 
-    val containingFile: FirFile?
-        get() = if (containingDeclarations.isEmpty()) null else containingDeclarations.first() as FirFile
+    abstract val containingFile: FirFile?
 
     override val containingFilePath: String?
         get() = containingFile?.sourceFile?.path

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.analysis.checkers.context
 
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
@@ -64,4 +65,8 @@ abstract class CheckerContextForProvider(
 
     abstract fun exitContractBody(): CheckerContextForProvider
 
+    abstract fun enterFile(file: FirFile): CheckerContextForProvider
+
+    abstract fun exitFile(file: FirFile): CheckerContextForProvider
+
 }

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

@@ -9,6 +9,7 @@ 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.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
@@ -24,6 +25,7 @@ class MutableCheckerContext private constructor(
     override val getClassCalls: MutableList<FirGetClassCall>,
     override val annotationContainers: MutableList<FirAnnotationContainer>,
     override var isContractBody: Boolean,
+    override var containingFile: FirFile?,
     sessionHolder: SessionHolder,
     returnTypeCalculator: ReturnTypeCalculator,
     override val suppressedDiagnostics: PersistentSet<String>,
@@ -38,6 +40,7 @@ class MutableCheckerContext private constructor(
         mutableListOf(),
         mutableListOf(),
         isContractBody = false,
+        containingFile = null,
         sessionHolder,
         returnTypeCalculator,
         persistentSetOf(),
@@ -54,6 +57,7 @@ class MutableCheckerContext private constructor(
             getClassCalls,
             annotationContainers,
             isContractBody,
+            containingFile,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics,
@@ -113,6 +117,7 @@ class MutableCheckerContext private constructor(
             getClassCalls,
             annotationContainers,
             isContractBody,
+            containingFile,
             sessionHolder,
             returnTypeCalculator,
             suppressedDiagnostics.addAll(diagnosticNames),
@@ -133,4 +138,14 @@ class MutableCheckerContext private constructor(
         isContractBody = false
         return this
     }
+
+    override fun enterFile(file: FirFile): CheckerContextForProvider {
+        containingFile = file
+        return this
+    }
+
+    override fun exitFile(file: FirFile): CheckerContextForProvider {
+        containingFile = file
+        return this
+    }
 }

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

@@ -11,6 +11,7 @@ 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.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
@@ -31,7 +32,8 @@ class PersistentCheckerContext private constructor(
     override val suppressedDiagnostics: PersistentSet<String>,
     allInfosSuppressed: Boolean,
     allWarningsSuppressed: Boolean,
-    allErrorsSuppressed: Boolean
+    allErrorsSuppressed: Boolean,
+    override val containingFile: FirFile?,
 ) : CheckerContextForProvider(sessionHolder, returnTypeCalculator, allInfosSuppressed, allWarningsSuppressed, allErrorsSuppressed) {
     constructor(sessionHolder: SessionHolder, returnTypeCalculator: ReturnTypeCalculator) : this(
         PersistentImplicitReceiverStack(),
@@ -45,7 +47,8 @@ class PersistentCheckerContext private constructor(
         persistentSetOf(),
         allInfosSuppressed = false,
         allWarningsSuppressed = false,
-        allErrorsSuppressed = false
+        allErrorsSuppressed = false,
+        containingFile = null,
     )
 
     override fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): PersistentCheckerContext =
@@ -106,11 +109,12 @@ class PersistentCheckerContext private constructor(
         allWarningsSuppressed: Boolean = this.allWarningsSuppressed,
         allErrorsSuppressed: Boolean = this.allErrorsSuppressed,
         suppressedDiagnostics: PersistentSet<String> = this.suppressedDiagnostics,
+        containingFile: FirFile? = this.containingFile,
     ): PersistentCheckerContext {
         return PersistentCheckerContext(
             implicitReceiverStack, containingDeclarations, qualifiedAccessOrAssignmentsOrAnnotationCalls,
             getClassCalls, annotationContainers, isContractBody, sessionHolder, returnTypeCalculator, suppressedDiagnostics,
-            allInfosSuppressed, allWarningsSuppressed, allErrorsSuppressed,
+            allInfosSuppressed, allWarningsSuppressed, allErrorsSuppressed, containingFile,
         )
     }
 
@@ -123,4 +127,8 @@ class PersistentCheckerContext private constructor(
     override fun enterContractBody(): CheckerContextForProvider = toggleContractBody(newValue = true)
 
     override fun exitContractBody(): CheckerContextForProvider = toggleContractBody(newValue = false)
+
+    override fun enterFile(file: FirFile): CheckerContextForProvider = copy(containingFile = file)
+
+    override fun exitFile(file: FirFile): CheckerContextForProvider = copy(containingFile = null)
 }

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

@@ -161,7 +161,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
     override fun visitFile(file: FirFile, data: Nothing?) {
         withAnnotationContainer(file) {
-            visitWithDeclaration(file)
+            visitWithFile(file)
         }
     }
 
@@ -253,6 +253,15 @@ abstract class AbstractDiagnosticCollectorVisitor(
         }
     }
 
+    protected inline fun visitWithFile(
+        file: FirFile,
+        block: () -> Unit = { visitNestedElements(file) }
+    ) {
+        withFile(file) {
+            visitWithDeclaration(file, block)
+        }
+    }
+
     private fun visitWithDeclarationAndReceiver(declaration: FirDeclaration, labelName: Name?, receiverParameter: FirReceiverParameter?) {
         visitWithDeclaration(declaration) {
             withLabelAndReceiverType(
@@ -321,6 +330,17 @@ abstract class AbstractDiagnosticCollectorVisitor(
         }
     }
 
+    @OptIn(PrivateForInline::class)
+    inline fun <R> withFile(file: FirFile, block: () -> R): R {
+        val existingContext = context
+        context = context.enterFile(file)
+        try {
+            return block()
+        } finally {
+            existingContext.exitFile(file)
+            context = existingContext
+        }
+    }
 
     @OptIn(PrivateForInline::class)
     inline fun <R> withLabelAndReceiverType(