From 932eab8dfc196a32f78adc3e179d3c7ac7eb1bc2 Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Fri, 15 Oct 2021 12:43:36 +0200
Subject: [PATCH] Implement semi-lazy PSI-based suppress cache for IR

calculates suppressions per root (KtFile) basis on demand.
---
 .../kotlin/codegen/state/GenerationState.kt   |  3 +-
 .../diagnostics/PrecomputedSuppressCache.kt   | 49 +++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/state/GenerationState.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/state/GenerationState.kt
index b6c7d616c0d..1ef4b0b148b 100644
--- a/compiler/backend/src/org/jetbrains/kotlin/codegen/state/GenerationState.kt
+++ b/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
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/diagnostics/PrecomputedSuppressCache.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/diagnostics/PrecomputedSuppressCache.kt
index 739ecfc2cb2..2a29f0a052a 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/diagnostics/PrecomputedSuppressCache.kt
+++ b/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)
+    }
+}