Implement semi-lazy PSI-based suppress cache for IR

calculates suppressions per root (KtFile) basis on demand.

compiler/backend/src/org/jetbrains/kotlin/codegen/state/GenerationState.kt

@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.deprecation.DeprecationResolver
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
+import org.jetbrains.kotlin.resolve.diagnostics.OnDemandSuppressCache
 import org.jetbrains.kotlin.resolve.diagnostics.PrecomputedSuppressCache
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
@@ -198,7 +199,7 @@ class GenerationState private constructor(
     val extraJvmDiagnosticsTrace: BindingTrace =
         DelegatingBindingTrace(
             originalFrontendBindingContext, "For extra diagnostics in ${this::class.java}", false,
-            customSuppressCache = if (isIrBackend) PrecomputedSuppressCache(originalFrontendBindingContext, files) else null,
+            customSuppressCache = if (isIrBackend) OnDemandSuppressCache(originalFrontendBindingContext) else null,
         )
 
     private val interceptedBuilderFactory: ClassBuilderFactory

compiler/frontend/src/org/jetbrains/kotlin/resolve/diagnostics/PrecomputedSuppressCache.kt

@@ -59,3 +59,52 @@ class PrecomputedSuppressCache(context: BindingContext, files: List<KtFile>) : K
     override fun getSuppressionAnnotations(annotated: PsiElement): List<AnnotationDescriptor> =
         storage[annotated].orEmpty()
 }
+
+class OnDemandSuppressCache(private val context: BindingContext) : KotlinSuppressCache() {
+
+    private val processedRoots = mutableSetOf<KtFile>()
+
+    private val storage = mutableMapOf<PsiElement, List<AnnotationDescriptor>>()
+
+    @Synchronized
+    private fun ensureRootProcessed(rootElement: PsiElement) {
+        require(rootElement is KtFile)
+        if (!processedRoots.contains(rootElement)) {
+            val visitor = PrecomputingVisitor(storage, BindingContextSuppressCache(context))
+            rootElement.accept(visitor, null)
+            processedRoots.add(rootElement)
+        }
+    }
+
+    private class PrecomputingVisitor(
+        val storage: MutableMap<PsiElement, List<AnnotationDescriptor>>,
+        val computer: KotlinSuppressCache,
+    ) : KtTreeVisitorVoid() {
+        override fun visitKtElement(element: KtElement) {
+            super.visitKtElement(element)
+            if (element is KtAnnotated) {
+                computeAnnotations(element)
+            }
+        }
+
+        override fun visitKtFile(file: KtFile) {
+            super.visitKtFile(file)
+            computeAnnotations(file)
+        }
+
+        private fun computeAnnotations(element: KtAnnotated) {
+            val suppressions = computer.getSuppressionAnnotations(element).filter { it.fqName == StandardNames.FqNames.suppress }
+            if (suppressions.isNotEmpty()) {
+                storage[element] = suppressions
+            }
+        }
+    }
+
+    override fun getSuppressionAnnotations(annotated: PsiElement): List<AnnotationDescriptor> =
+        storage[annotated].orEmpty()
+
+    override fun getClosestAnnotatedAncestorElement(element: PsiElement, rootElement: PsiElement, excludeSelf: Boolean): PsiElement? {
+        ensureRootProcessed(rootElement)
+        return super.getClosestAnnotatedAncestorElement(element, rootElement, excludeSelf)
+    }
+}