[FIR] Fix mutable context usage in inline diagnostics

Also some minor fixes

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

@@ -53,6 +53,15 @@ abstract class CheckerContext {
 
     abstract fun dropDeclaration()
 
+    fun <T> withDeclaration(declaration: FirDeclaration, f: (CheckerContext) -> T): T {
+        val newContext = addDeclaration(declaration)
+        try {
+            return f(newContext)
+        } finally {
+            newContext.dropDeclaration()
+        }
+    }
+
     abstract fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): CheckerContext
 
     abstract fun dropQualifiedAccessOrAnnotationCall()

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirInlineDeclarationChecker.kt

@@ -10,13 +10,13 @@ import org.jetbrains.kotlin.descriptors.EffectiveVisibility
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.isInlineOnly
 import org.jetbrains.kotlin.fir.analysis.checkers.unsubstitutedScope
 import org.jetbrains.kotlin.fir.analysis.checkers.util.checkChildrenWithCustomVisitor
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.diagnostics.withSuppressedDiagnostics
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.*
 import org.jetbrains.kotlin.fir.expressions.*
@@ -50,8 +50,10 @@ object FirInlineDeclarationChecker : FirFunctionChecker() {
         if (declaration !is FirPropertyAccessor && declaration !is FirSimpleFunction) return
 
         val effectiveVisibility = declaration.effectiveVisibility
-        checkInlineFunctionBody(declaration, effectiveVisibility, context, reporter)
+        withSuppressedDiagnostics(declaration, context) { ctx ->
-        checkCallableDeclaration(declaration, context, reporter)
+            checkInlineFunctionBody(declaration, effectiveVisibility, ctx, reporter)
+            checkCallableDeclaration(declaration, ctx, reporter)
+        }
     }
 
     private fun checkInlineFunctionBody(
@@ -74,7 +76,9 @@ object FirInlineDeclarationChecker : FirFunctionChecker() {
             context.session,
             reporter
         )
-        body.checkChildrenWithCustomVisitor((context as PersistentCheckerContext).addDeclaration(function), visitor)
+        context.withDeclaration(function) {
+            body.checkChildrenWithCustomVisitor(it, visitor)
+        }
     }
 
     private class Visitor(

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirPublishedApiChecker.kt

@@ -13,14 +13,14 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.declarations.getAnnotationsByFqName
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByFqName
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 
 object FirPublishedApiChecker : FirAnnotatedDeclarationChecker() {
     override fun check(declaration: FirAnnotatedDeclaration, context: CheckerContext, reporter: DiagnosticReporter) {
         if (declaration !is FirMemberDeclaration) return
         if (declaration.visibility == Visibilities.Internal) return
-        val annotation = declaration.getAnnotationsByFqName(StandardNames.FqNames.publishedApi).firstOrNull() ?: return
+        val annotation = declaration.getAnnotationByFqName(StandardNames.FqNames.publishedApi) ?: return
         reporter.reportOn(annotation.source, FirErrors.NON_INTERNAL_PUBLISHED_API, context)
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/declarations/FirAnnotationUtils.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
-import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -94,7 +93,7 @@ fun FirAnnotatedDeclaration.hasAnnotation(classId: ClassId): Boolean {
 
 fun FirAnnotationContainer.getAnnotationByFqName(fqName: FqName): FirAnnotationCall? {
     return annotations.find {
-        it.annotationTypeRef.coneTypeUnsafe<ConeClassLikeType>().lookupTag.classId.asSingleFqName() == fqName
+        it.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.lookupTag?.classId?.asSingleFqName() == fqName
     }
 }