[FIR] Add stack of all FIR elements to CheckerContext

^KT-56472 Fixed

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

@@ -35,6 +35,7 @@ abstract class CheckerContext : DiagnosticContext {
     abstract val qualifiedAccessOrAssignmentsOrAnnotationCalls: List<FirStatement>
     abstract val getClassCalls: List<FirGetClassCall>
     abstract val annotationContainers: List<FirAnnotationContainer>
+    abstract val containingElements: List<FirElement>
     abstract val isContractBody: Boolean
 
     // Suppress

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.context
 
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
@@ -69,4 +70,7 @@ abstract class CheckerContextForProvider(
 
     abstract fun exitFile(file: FirFile): CheckerContextForProvider
 
+    abstract fun addElement(element: FirElement): CheckerContextForProvider
+
+    abstract fun dropElement()
 }

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

@@ -8,6 +8,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.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
@@ -24,6 +25,7 @@ class MutableCheckerContext private constructor(
     override val qualifiedAccessOrAssignmentsOrAnnotationCalls: MutableList<FirStatement>,
     override val getClassCalls: MutableList<FirGetClassCall>,
     override val annotationContainers: MutableList<FirAnnotationContainer>,
+    override val containingElements: MutableList<FirElement>,
     override var isContractBody: Boolean,
     override var containingFile: FirFile?,
     sessionHolder: SessionHolder,
@@ -39,6 +41,7 @@ class MutableCheckerContext private constructor(
         mutableListOf(),
         mutableListOf(),
         mutableListOf(),
+        mutableListOf(),
         isContractBody = false,
         containingFile = null,
         sessionHolder,
@@ -56,6 +59,7 @@ class MutableCheckerContext private constructor(
             qualifiedAccessOrAssignmentsOrAnnotationCalls,
             getClassCalls,
             annotationContainers,
+            containingElements,
             isContractBody,
             containingFile,
             sessionHolder,
@@ -73,7 +77,7 @@ class MutableCheckerContext private constructor(
     }
 
     override fun dropDeclaration() {
-        containingDeclarations.removeAt(containingDeclarations.size - 1)
+        containingDeclarations.removeLast()
     }
 
     override fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): MutableCheckerContext {
@@ -82,7 +86,7 @@ class MutableCheckerContext private constructor(
     }
 
     override fun dropQualifiedAccessOrAnnotationCall() {
-        qualifiedAccessOrAssignmentsOrAnnotationCalls.removeAt(qualifiedAccessOrAssignmentsOrAnnotationCalls.size - 1)
+        qualifiedAccessOrAssignmentsOrAnnotationCalls.removeLast()
     }
 
     override fun addGetClassCall(getClassCall: FirGetClassCall): MutableCheckerContext {
@@ -91,7 +95,7 @@ class MutableCheckerContext private constructor(
     }
 
     override fun dropGetClassCall() {
-        getClassCalls.removeAt(getClassCalls.size - 1)
+        getClassCalls.removeLast()
     }
 
     override fun addAnnotationContainer(annotationContainer: FirAnnotationContainer): CheckerContextForProvider {
@@ -100,7 +104,17 @@ class MutableCheckerContext private constructor(
     }
 
     override fun dropAnnotationContainer() {
-        annotationContainers.removeAt(annotationContainers.size - 1)
+        annotationContainers.removeLast()
+    }
+
+    override fun addElement(element: FirElement): CheckerContextForProvider {
+        assert(containingElements.lastOrNull() !== element)
+        containingElements.add(element)
+        return this
+    }
+
+    override fun dropElement() {
+        containingElements.removeLast()
     }
 
     override fun addSuppressedDiagnostics(
@@ -116,6 +130,7 @@ class MutableCheckerContext private constructor(
             qualifiedAccessOrAssignmentsOrAnnotationCalls,
             getClassCalls,
             annotationContainers,
+            containingElements,
             isContractBody,
             containingFile,
             sessionHolder,

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

@@ -10,6 +10,7 @@ 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.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
@@ -26,6 +27,7 @@ class PersistentCheckerContext private constructor(
     override val qualifiedAccessOrAssignmentsOrAnnotationCalls: PersistentList<FirStatement>,
     override val getClassCalls: PersistentList<FirGetClassCall>,
     override val annotationContainers: PersistentList<FirAnnotationContainer>,
+    override val containingElements: PersistentList<FirElement>,
     override val isContractBody: Boolean,
     sessionHolder: SessionHolder,
     returnTypeCalculator: ReturnTypeCalculator,
@@ -41,6 +43,7 @@ class PersistentCheckerContext private constructor(
         persistentListOf(),
         persistentListOf(),
         persistentListOf(),
+        persistentListOf(),
         isContractBody = false,
         sessionHolder,
         returnTypeCalculator,
@@ -57,8 +60,7 @@ class PersistentCheckerContext private constructor(
     override fun addDeclaration(declaration: FirDeclaration): PersistentCheckerContext =
         copy(containingDeclarations = containingDeclarations.add(declaration))
 
-    override fun dropDeclaration() {
+    override fun dropDeclaration() {}
-    }
 
     override fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): PersistentCheckerContext =
         copy(
@@ -66,22 +68,22 @@ class PersistentCheckerContext private constructor(
             qualifiedAccessOrAssignmentsOrAnnotationCalls.add(qualifiedAccessOrAnnotationCall)
         )
 
-    override fun dropQualifiedAccessOrAnnotationCall() {
+    override fun dropQualifiedAccessOrAnnotationCall() {}
-    }
 
     override fun addGetClassCall(getClassCall: FirGetClassCall): PersistentCheckerContext =
-        copy(
+        copy(getClassCalls = getClassCalls.add(getClassCall))
-            getClassCalls = getClassCalls.add(getClassCall),
-        )
 
-    override fun dropGetClassCall() {
+    override fun dropGetClassCall() {}
-    }
 
     override fun addAnnotationContainer(annotationContainer: FirAnnotationContainer): PersistentCheckerContext =
         copy(annotationContainers = annotationContainers.add(annotationContainer))
 
-    override fun dropAnnotationContainer() {
+    override fun dropAnnotationContainer() {}
-    }
+
+    override fun addElement(element: FirElement): PersistentCheckerContext =
+        copy(containingElements = containingElements.add(element))
+
+    override fun dropElement() {}
 
     override fun addSuppressedDiagnostics(
         diagnosticNames: Collection<String>,
@@ -103,6 +105,7 @@ class PersistentCheckerContext private constructor(
         qualifiedAccessOrAssignmentsOrAnnotationCalls: PersistentList<FirStatement> = this.qualifiedAccessOrAssignmentsOrAnnotationCalls,
         getClassCalls: PersistentList<FirGetClassCall> = this.getClassCalls,
         annotationContainers: PersistentList<FirAnnotationContainer> = this.annotationContainers,
+        containingElements: PersistentList<FirElement> = this.containingElements,
         containingDeclarations: PersistentList<FirDeclaration> = this.containingDeclarations,
         isContractBody: Boolean = this.isContractBody,
         allInfosSuppressed: Boolean = this.allInfosSuppressed,
@@ -112,8 +115,16 @@ class PersistentCheckerContext private constructor(
         containingFile: FirFile? = this.containingFile,
     ): PersistentCheckerContext {
         return PersistentCheckerContext(
-            implicitReceiverStack, containingDeclarations, qualifiedAccessOrAssignmentsOrAnnotationCalls,
+            implicitReceiverStack,
-            getClassCalls, annotationContainers, isContractBody, sessionHolder, returnTypeCalculator, suppressedDiagnostics,
+            containingDeclarations,
+            qualifiedAccessOrAssignmentsOrAnnotationCalls,
+            getClassCalls,
+            annotationContainers,
+            containingElements,
+            isContractBody,
+            sessionHolder,
+            returnTypeCalculator,
+            suppressedDiagnostics,
             allInfosSuppressed, allWarningsSuppressed, allErrorsSuppressed, containingFile,
         )
     }

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.analysis.collectors
 
-import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.PrivateForInline
@@ -35,16 +34,16 @@ abstract class AbstractDiagnosticCollectorVisitor(
     protected abstract fun checkElement(element: FirElement)
 
     override fun visitElement(element: FirElement, data: Nothing?) {
-        if (element is FirAnnotationContainer) {
+        when (element) {
-            withAnnotationContainer(element) {
+            is FirAnnotationContainer -> withAnnotationContainer(element) {
+                checkElement(element)
+                visitNestedElements(element)
+            }
+            else -> withElement(element) {
                 checkElement(element)
                 visitNestedElements(element)
             }
-            return
         }
-
-        checkElement(element)
-        visitNestedElements(element)
     }
 
     override fun visitAnnotationContainer(annotationContainer: FirAnnotationContainer, data: Nothing?) {
@@ -166,18 +165,18 @@ abstract class AbstractDiagnosticCollectorVisitor(
     }
 
     override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer, data: Nothing?) {
-        visitWithDeclaration(anonymousInitializer)
+        withElement(anonymousInitializer) {
+            visitWithDeclaration(anonymousInitializer)
+        }
     }
 
     override fun visitBlock(block: FirBlock, data: Nothing?) {
-        withAnnotationContainer(block) {
+        if (block is FirContractCallBlock) {
-            if (block is FirContractCallBlock) {
+            insideContractBody {
-                insideContractBody {
-                    visitExpression(block, data)
-                }
-            } else {
                 visitExpression(block, data)
             }
+        } else {
+            visitExpression(block, data)
         }
     }
 
@@ -341,6 +340,19 @@ abstract class AbstractDiagnosticCollectorVisitor(
             context = existingContext
         }
     }
+    @OptIn(PrivateForInline::class)
+    inline fun <T> withElement(element: FirElement, block: () -> T): T {
+        val existingContext = context
+        context = context.addElement(element)
+        return try {
+            whileAnalysing(context.session, element) {
+                block()
+            }
+        } finally {
+            existingContext.dropElement()
+            context = existingContext
+        }
+    }
 
     @OptIn(PrivateForInline::class)
     inline fun <R> withLabelAndReceiverType(
@@ -367,19 +379,21 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
     @OptIn(PrivateForInline::class)
     inline fun <R> withAnnotationContainer(annotationContainer: FirAnnotationContainer, block: () -> R): R {
-        val existingContext = context
+        return withElement(annotationContainer) {
-        addSuppressedDiagnosticsToContext(annotationContainer)
+            val existingContext = context
-        val notEmptyAnnotations = annotationContainer.annotations.isNotEmpty()
+            addSuppressedDiagnosticsToContext(annotationContainer)
-        if (notEmptyAnnotations) {
+            val notEmptyAnnotations = annotationContainer.annotations.isNotEmpty()
-            context = context.addAnnotationContainer(annotationContainer)
-        }
-        return try {
-            block()
-        } finally {
             if (notEmptyAnnotations) {
-                existingContext.dropAnnotationContainer()
+                context = context.addAnnotationContainer(annotationContainer)
+            }
+            try {
+                block()
+            } finally {
+                if (notEmptyAnnotations) {
+                    existingContext.dropAnnotationContainer()
+                }
+                context = existingContext
             }
-            context = existingContext
         }
     }