[FIR] withSuppressedDiagnostics now takes FirAnnotationContainer instead of FirElement

Reduce allocations in withSuppressedDiagnostics

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

@@ -80,7 +80,7 @@ object FirClassLiteralChecker : FirGetClassCallChecker() {
             @OptIn(SymbolInternals::class)
             val typeParameters = (symbol?.fir as? FirTypeParameterRefsOwner)?.typeParameters
             // Among type parameter references, only count actual type parameter while discarding [FirOuterClassTypeParameterRef]
-            val expectedTypeArgumentSize = typeParameters?.filterIsInstance<FirTypeParameter>()?.size ?: 0
+            val expectedTypeArgumentSize = typeParameters?.count { it is FirTypeParameter } ?: 0
             if (expectedTypeArgumentSize != argument.typeArguments.size) {
                 // Will be reported as WRONG_NUMBER_OF_TYPE_ARGUMENTS
                 return

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/DiagnosticReporter.kt

@@ -144,12 +144,77 @@ fun <F : AbstractFirDiagnosticFactory> FirDiagnosticFactoryForDeprecation<F>.cho
     }
 }
 
+fun DiagnosticReporter.reportOnWithSuppression(
+    element: FirAnnotationContainer,
+    factory: FirDiagnosticFactory0,
+    context: CheckerContext,
+    positioningStrategy: SourceElementPositioningStrategy? = null
+) {
+    withSuppressedDiagnostics(element, context) {
+        reportOn(element.source, factory, it, positioningStrategy)
+    }
+}
+
+fun <A : Any> DiagnosticReporter.reportOnWithSuppression(
+    element: FirAnnotationContainer,
+    factory: FirDiagnosticFactory1<A>,
+    a: A,
+    context: CheckerContext,
+    positioningStrategy: SourceElementPositioningStrategy? = null
+) {
+    withSuppressedDiagnostics(element, context) {
+        reportOn(element.source, factory, a, it, positioningStrategy)
+    }
+}
+
+fun <A : Any, B : Any> DiagnosticReporter.reportOnWithSuppression(
+    element: FirAnnotationContainer,
+    factory: FirDiagnosticFactory2<A, B>,
+    a: A,
+    b: B,
+    context: CheckerContext,
+    positioningStrategy: SourceElementPositioningStrategy? = null
+) {
+    withSuppressedDiagnostics(element, context) {
+        reportOn(element.source, factory, a, b, it, positioningStrategy)
+    }
+}
+
+fun <A : Any, B : Any, C : Any> DiagnosticReporter.reportOnWithSuppression(
+    element: FirAnnotationContainer,
+    factory: FirDiagnosticFactory3<A, B, C>,
+    a: A,
+    b: B,
+    c: C,
+    context: CheckerContext,
+    positioningStrategy: SourceElementPositioningStrategy? = null
+) {
+    withSuppressedDiagnostics(element, context) {
+        reportOn(element.source, factory, a, b, c, it, positioningStrategy)
+    }
+}
+
+fun <A : Any, B : Any, C : Any, D : Any> DiagnosticReporter.reportOnWithSuppression(
+    element: FirAnnotationContainer,
+    factory: FirDiagnosticFactory4<A, B, C, D>,
+    a: A,
+    b: B,
+    c: C,
+    d: D,
+    context: CheckerContext,
+    positioningStrategy: SourceElementPositioningStrategy? = null
+) {
+    withSuppressedDiagnostics(element, context) {
+        reportOn(element.source, factory, a, b, c, d, it, positioningStrategy)
+    }
+}
+
 inline fun withSuppressedDiagnostics(
-    element: FirElement,
+    annotationContainer: FirAnnotationContainer,
     context: CheckerContext,
     f: (CheckerContext) -> Unit
 ) {
-    val arguments = (element as? FirAnnotationContainer)?.let { AbstractDiagnosticCollector.getDiagnosticsSuppressedForContainer(it) }
+    val arguments = AbstractDiagnosticCollector.getDiagnosticsSuppressedForContainer(annotationContainer)
     if (arguments != null) {
         f(
             context.addSuppressedDiagnostics(
@@ -164,68 +229,3 @@ inline fun withSuppressedDiagnostics(
     f(context)
 }
 
-fun DiagnosticReporter.reportOnWithSuppression(
-    element: FirElement,
-    factory: FirDiagnosticFactory0,
-    context: CheckerContext,
-    positioningStrategy: SourceElementPositioningStrategy? = null
-) {
-    withSuppressedDiagnostics(element, context) {
-        reportOn(element.source, factory, it, positioningStrategy)
-    }
-}
-
-fun <A : Any> DiagnosticReporter.reportOnWithSuppression(
-    element: FirElement,
-    factory: FirDiagnosticFactory1<A>,
-    a: A,
-    context: CheckerContext,
-    positioningStrategy: SourceElementPositioningStrategy? = null
-) {
-    withSuppressedDiagnostics(element, context) {
-        reportOn(element.source, factory, a, it, positioningStrategy)
-    }
-}
-
-fun <A : Any, B : Any> DiagnosticReporter.reportOnWithSuppression(
-    element: FirElement,
-    factory: FirDiagnosticFactory2<A, B>,
-    a: A,
-    b: B,
-    context: CheckerContext,
-    positioningStrategy: SourceElementPositioningStrategy? = null
-) {
-    withSuppressedDiagnostics(element, context) {
-        reportOn(element.source, factory, a, b, it, positioningStrategy)
-    }
-}
-
-fun <A : Any, B : Any, C : Any> DiagnosticReporter.reportOnWithSuppression(
-    element: FirElement,
-    factory: FirDiagnosticFactory3<A, B, C>,
-    a: A,
-    b: B,
-    c: C,
-    context: CheckerContext,
-    positioningStrategy: SourceElementPositioningStrategy? = null
-) {
-    withSuppressedDiagnostics(element, context) {
-        reportOn(element.source, factory, a, b, c, it, positioningStrategy)
-    }
-}
-
-fun <A : Any, B : Any, C : Any, D : Any> DiagnosticReporter.reportOnWithSuppression(
-    element: FirElement,
-    factory: FirDiagnosticFactory4<A, B, C, D>,
-    a: A,
-    b: B,
-    c: C,
-    d: D,
-    context: CheckerContext,
-    positioningStrategy: SourceElementPositioningStrategy? = null
-) {
-    withSuppressedDiagnostics(element, context) {
-        reportOn(element.source, factory, a, b, c, d, it, positioningStrategy)
-    }
-}
-