Extract psi- and descriptior-independent suppress cache logic

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Severity
+import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.constants.ArrayValue
@@ -40,44 +41,22 @@ interface DiagnosticSuppressor {
     }
 }
 
-abstract class KotlinSuppressCache {
+abstract class AbstractKotlinSuppressCache<Element> {
-    private val diagnosticSuppressors = ExtensionProvider.create(DiagnosticSuppressor.EP_NAME)
-
     // The cache is weak: we're OK with losing it
-    private val suppressors = ContainerUtil.createConcurrentWeakValueMap<KtAnnotated, Suppressor>()
+    protected val suppressors = ContainerUtil.createConcurrentWeakValueMap<Element, Suppressor<Element>>()
 
-    val filter: (Diagnostic) -> Boolean = { diagnostic: Diagnostic ->
+    fun isSuppressed(element: Element, suppressionKey: String, severity: Severity) =
-        !isSuppressed(DiagnosticSuppressRequest(diagnostic))
+        isSuppressed(StringSuppressRequest(element, severity, suppressionKey.lowercase()))
-    }
 
-    fun isSuppressed(psiElement: PsiElement, suppressionKey: String, severity: Severity) =
+    protected open fun isSuppressed(request: SuppressRequest<Element>): Boolean {
-        isSuppressed(StringSuppressRequest(psiElement, severity, suppressionKey.lowercase()))
 
-    private fun isSuppressed(request: SuppressRequest): Boolean {
+        val annotated = getParentAnnotatedElement(request.element, false) ?: return false
-        // If diagnostics are reported in a synthetic file generated by KtPsiFactory (dummy.kt),
-        // there's no point to present such diagnostics to the user, because the user didn't write this code
-        val element = request.element
-        if (!element.isValid) return true
-
-        val file = element.containingFile
-        if (file is KtFile) {
-            if (file.doNotAnalyze != null) return true
-        }
-
-        if (request is DiagnosticSuppressRequest) {
-            for (suppressor in diagnosticSuppressors.get()) {
-                if (isSuppressedByExtension(suppressor, request.diagnostic)) return true
-            }
-        }
-
-        val annotated = KtStubbedPsiUtil.getPsiOrStubParent(element, KtAnnotated::class.java, false) ?: return false
 
         return isSuppressedByAnnotated(request.suppressKey, request.severity, annotated, 0)
+
     }
 
-    protected open fun isSuppressedByExtension(suppressor: DiagnosticSuppressor, diagnostic: Diagnostic): Boolean {
+    protected abstract fun getParentAnnotatedElement(element: Element, strict: Boolean): Element?
-        return suppressor.isSuppressed(diagnostic)
-    }
 
     /*
        The cache is optimized for the case where no warnings are suppressed (most frequent one)
@@ -108,16 +87,16 @@ abstract class KotlinSuppressCache {
 
        This way we need no more lookups than the number of suppress() annotations from here to the root.
      */
-    private fun isSuppressedByAnnotated(
+    protected open fun isSuppressedByAnnotated(
         suppressionKey: String,
         severity: Severity,
-        annotated: KtAnnotated,
+        annotated: Element,
         debugDepth: Int
     ): Boolean {
         val suppressor = getOrCreateSuppressor(annotated)
         if (suppressor.isSuppressed(suppressionKey, severity)) return true
 
-        val annotatedAbove = KtStubbedPsiUtil.getPsiOrStubParent(suppressor.annotatedElement, KtAnnotated::class.java, true) ?: return false
+        val annotatedAbove = getParentAnnotatedElement(suppressor.annotatedElement, true) ?: return false
 
         val suppressed = isSuppressedByAnnotated(suppressionKey, severity, annotatedAbove, debugDepth + 1)
         val suppressorAbove = suppressors[annotatedAbove]
@@ -128,7 +107,7 @@ abstract class KotlinSuppressCache {
         return suppressed
     }
 
-    private fun getOrCreateSuppressor(annotated: KtAnnotated): Suppressor =
+    protected fun getOrCreateSuppressor(annotated: Element): Suppressor<Element> =
         suppressors.getOrPut(annotated) {
             val strings = getSuppressingStrings(annotated)
             when (strings.size) {
@@ -138,9 +117,74 @@ abstract class KotlinSuppressCache {
             }
         }
 
-    abstract fun getSuppressionAnnotations(annotated: KtAnnotated): List<AnnotationDescriptor>
+    protected abstract fun getSuppressingStrings(annotated: Element): Set<String>
 
-    private fun getSuppressingStrings(annotated: KtAnnotated): Set<String> {
+    companion object {
+        private fun isSuppressedByStrings(key: String, strings: Set<String>, severity: Severity): Boolean =
+            severity == Severity.WARNING && "warnings" in strings || key.lowercase() in strings
+    }
+
+    protected abstract class Suppressor<Element>(val annotatedElement: Element) {
+        abstract fun isSuppressed(suppressionKey: String, severity: Severity): Boolean
+
+        // true is \forall x. other.isSuppressed(x) -> this.isSuppressed(x)
+        abstract fun dominates(other: Suppressor<Element>): Boolean
+    }
+
+    private class EmptySuppressor<Element>(annotated: Element) : Suppressor<Element>(annotated) {
+        override fun isSuppressed(suppressionKey: String, severity: Severity): Boolean = false
+        override fun dominates(other: Suppressor<Element>): Boolean = other is EmptySuppressor
+    }
+
+    private class SingularSuppressor<Element>(annotated: Element, private val string: String) : Suppressor<Element>(annotated) {
+        override fun isSuppressed(suppressionKey: String, severity: Severity): Boolean {
+            return isSuppressedByStrings(suppressionKey, ImmutableSet.of(string), severity)
+        }
+
+        override fun dominates(other: Suppressor<Element>): Boolean {
+            return other is EmptySuppressor || (other is SingularSuppressor && other.string == string)
+        }
+    }
+
+    private class MultiSuppressor<Element>(annotated: Element, private val strings: Set<String>) : Suppressor<Element>(annotated) {
+        override fun isSuppressed(suppressionKey: String, severity: Severity): Boolean {
+            return isSuppressedByStrings(suppressionKey, strings, severity)
+        }
+
+        override fun dominates(other: Suppressor<Element>): Boolean {
+            // it's too costly to check set inclusion
+            return other is EmptySuppressor
+        }
+    }
+
+    protected interface SuppressRequest<Element> {
+        val element: Element
+        val severity: Severity
+        val suppressKey: String
+    }
+
+    private class StringSuppressRequest<Element>(
+        override val element: Element,
+        override val severity: Severity,
+        override val suppressKey: String
+    ) : SuppressRequest<Element>
+}
+
+abstract class KotlinSuppressCache : AbstractKotlinSuppressCache<PsiElement>() {
+
+    private val diagnosticSuppressors = ExtensionProvider.create(DiagnosticSuppressor.EP_NAME)
+
+    val filter: (Diagnostic) -> Boolean = { diagnostic: Diagnostic ->
+        !isSuppressed(DiagnosticSuppressRequest(diagnostic))
+    }
+
+    protected open fun isSuppressedByExtension(suppressor: DiagnosticSuppressor, diagnostic: Diagnostic): Boolean {
+        return suppressor.isSuppressed(diagnostic)
+    }
+
+    abstract fun getSuppressionAnnotations(annotated: PsiElement): List<AnnotationDescriptor>
+
+    override fun getSuppressingStrings(annotated: PsiElement): Set<String> {
         val builder = ImmutableSet.builder<String>()
 
         for (annotationDescriptor in getSuppressionAnnotations(annotated)) {
@@ -165,76 +209,47 @@ abstract class KotlinSuppressCache {
         }
     }
 
-    companion object {
+    override fun isSuppressed(request: SuppressRequest<PsiElement>): Boolean {
-        private fun getDiagnosticSuppressKey(diagnostic: Diagnostic): String =
+        // If diagnostics are reported in a synthetic file generated by KtPsiFactory (dummy.kt),
-            diagnostic.factory.name.lowercase()
+        // there's no point to present such diagnostics to the user, because the user didn't write this code
+        val element = request.element
+        if (!element.isValid) return true
 
-        private fun isSuppressedByStrings(key: String, strings: Set<String>, severity: Severity): Boolean =
+        val file = element.containingFile
-            severity == Severity.WARNING && "warnings" in strings || key in strings
+        if (file is KtFile) {
-    }
+            if (file.doNotAnalyze != null) return true
-
-    private abstract class Suppressor(val annotatedElement: KtAnnotated) {
-        open fun isSuppressed(diagnostic: Diagnostic): Boolean =
-            isSuppressed(getDiagnosticSuppressKey(diagnostic), diagnostic.severity)
-
-        abstract fun isSuppressed(suppressionKey: String, severity: Severity): Boolean
-
-        // true is \forall x. other.isSuppressed(x) -> this.isSuppressed(x)
-        abstract fun dominates(other: Suppressor): Boolean
-    }
-
-    private class EmptySuppressor(annotated: KtAnnotated) : Suppressor(annotated) {
-        override fun isSuppressed(diagnostic: Diagnostic): Boolean = false
-        override fun isSuppressed(suppressionKey: String, severity: Severity): Boolean = false
-        override fun dominates(other: Suppressor): Boolean = other is EmptySuppressor
-    }
-
-    private class SingularSuppressor(annotated: KtAnnotated, private val string: String) : Suppressor(annotated) {
-        override fun isSuppressed(suppressionKey: String, severity: Severity): Boolean {
-            return isSuppressedByStrings(suppressionKey, ImmutableSet.of(string), severity)
         }
 
-        override fun dominates(other: Suppressor): Boolean {
+        if (request is DiagnosticSuppressRequest) {
-            return other is EmptySuppressor || (other is SingularSuppressor && other.string == string)
+            for (suppressor in diagnosticSuppressors.get()) {
+                if (isSuppressedByExtension(suppressor, request.diagnostic)) return true
+            }
         }
+        return super.isSuppressed(request)
     }
 
-    private class MultiSuppressor(annotated: KtAnnotated, private val strings: Set<String>) : Suppressor(annotated) {
+    override fun getParentAnnotatedElement(element: PsiElement, strict: Boolean) =
-        override fun isSuppressed(suppressionKey: String, severity: Severity): Boolean {
+        KtStubbedPsiUtil.getPsiOrStubParent(element, KtAnnotated::class.java, strict)
-            return isSuppressedByStrings(suppressionKey, strings, severity)
-        }
 
-        override fun dominates(other: Suppressor): Boolean {
+    protected class DiagnosticSuppressRequest(val diagnostic: Diagnostic) : SuppressRequest<PsiElement> {
-            // it's too costly to check set inclusion
-            return other is EmptySuppressor
-        }
-    }
-
-    private interface SuppressRequest {
-        val element: PsiElement
-        val severity: Severity
-        val suppressKey: String
-    }
-
-    private class StringSuppressRequest(
-        override val element: PsiElement,
-        override val severity: Severity,
-        override val suppressKey: String
-    ) : SuppressRequest
-
-    private class DiagnosticSuppressRequest(val diagnostic: Diagnostic) : SuppressRequest {
         override val element: PsiElement get() = diagnostic.psiElement
         override val severity: Severity get() = diagnostic.severity
         override val suppressKey: String get() = getDiagnosticSuppressKey(diagnostic)
     }
+
+    companion object {
+        internal fun getDiagnosticSuppressKey(diagnostic: Diagnostic): String =
+            diagnostic.factory.name.lowercase()
+    }
 }
 
 class BindingContextSuppressCache(val context: BindingContext) : KotlinSuppressCache() {
-    override fun getSuppressionAnnotations(annotated: KtAnnotated): List<AnnotationDescriptor> {
+    override fun getSuppressionAnnotations(annotated: PsiElement): List<AnnotationDescriptor> {
         val descriptor = context.get(BindingContext.DECLARATION_TO_DESCRIPTOR, annotated)
 
         return descriptor?.annotations?.toList()
-            ?: annotated.annotationEntries.mapNotNull { context.get(BindingContext.ANNOTATION, it) }
+            ?: (annotated as? KtAnnotated)?.annotationEntries?.mapNotNull { context.get(BindingContext.ANNOTATION, it) }
+            ?: emptyList()
     }
 
     override fun isSuppressedByExtension(suppressor: DiagnosticSuppressor, diagnostic: Diagnostic): Boolean {

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.diagnostics
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.psi.KtAnnotated
@@ -23,8 +24,8 @@ import org.jetbrains.kotlin.resolve.BindingContext
  *   primary use case of this class is when that binding context is cleared and thus is useless after a certain point.
  */
 class PrecomputedSuppressCache(context: BindingContext, files: List<KtFile>) : KotlinSuppressCache() {
-    val storage: Map<KtAnnotated, List<AnnotationDescriptor>> =
+    val storage: Map<PsiElement, List<AnnotationDescriptor>> =
-        mutableMapOf<KtAnnotated, List<AnnotationDescriptor>>().also { storage ->
+        mutableMapOf<PsiElement, List<AnnotationDescriptor>>().also { storage ->
             val visitor = PrecomputingVisitor(storage, BindingContextSuppressCache(context))
             for (file in files) {
                 file.accept(visitor, null)
@@ -32,7 +33,7 @@ class PrecomputedSuppressCache(context: BindingContext, files: List<KtFile>) : K
         }
 
     private class PrecomputingVisitor(
-        val storage: MutableMap<KtAnnotated, List<AnnotationDescriptor>>,
+        val storage: MutableMap<PsiElement, List<AnnotationDescriptor>>,
         val computer: KotlinSuppressCache,
     ) : KtTreeVisitorVoid() {
         override fun visitKtElement(element: KtElement) {
@@ -55,6 +56,6 @@ class PrecomputedSuppressCache(context: BindingContext, files: List<KtFile>) : K
         }
     }
 
-    override fun getSuppressionAnnotations(annotated: KtAnnotated): List<AnnotationDescriptor> =
+    override fun getSuppressionAnnotations(annotated: PsiElement): List<AnnotationDescriptor> =
         storage[annotated].orEmpty()
 }