Lazy diagnostics API in frontend

Relates to #KT-37702

compiler/frontend/src/org/jetbrains/kotlin/analyzer/AnalysisResult.kt

@@ -52,9 +52,9 @@ open class AnalysisResult protected constructor(
     fun isError(): Boolean = this is InternalError || this is CompilationError
 
     fun throwIfError() {
-        when {
-            this is InternalError -> throw IllegalStateException("failed to analyze: " + error, error)
-            this is CompilationError -> throw CompilationErrorException()
+        when (this) {
+            is InternalError -> throw IllegalStateException("failed to analyze: $error", error)
+            is CompilationError -> throw CompilationErrorException()
         }
     }
 

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/DiagnosticSink.java

@@ -53,6 +53,15 @@ public interface DiagnosticSink {
         }
     };
 
+    interface DiagnosticsCallback {
+        void callback(Diagnostic diagnostic);
+    }
+
     void report(@NotNull Diagnostic diagnostic);
+
+    default void setCallback(@NotNull DiagnosticsCallback callback) { }
+
+    default void resetCallback() { }
+
     boolean wantsDiagnostics();
 }

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/rendering/RenderingContext.kt

@@ -48,16 +48,16 @@ sealed class RenderingContext {
         }
 
         @JvmStatic
-        fun fromDiagnostic(d: Diagnostic): RenderingContext {
-            val parameters = when (d) {
-                is SimpleDiagnostic<*> -> listOf()
-                is DiagnosticWithParameters1<*, *> -> listOf(d.a)
-                is DiagnosticWithParameters2<*, *, *> -> listOf(d.a, d.b)
-                is DiagnosticWithParameters3<*, *, *, *> -> listOf(d.a, d.b, d.c)
-                is ParametrizedDiagnostic<*> -> error("Unexpected diagnostic: ${d::class.java}")
-                else -> listOf()
-            }
-            return Impl(parameters)
+        fun parameters(d: Diagnostic): List<Any> = when (d) {
+            is SimpleDiagnostic<*> -> listOf()
+            is DiagnosticWithParameters1<*, *> -> listOf(d.a)
+            is DiagnosticWithParameters2<*, *, *> -> listOf(d.a, d.b)
+            is DiagnosticWithParameters3<*, *, *, *> -> listOf(d.a, d.b, d.c)
+            is ParametrizedDiagnostic<*> -> error("Unexpected diagnostic: ${d::class.java}")
+            else -> listOf()
         }
+
+        @JvmStatic
+        fun fromDiagnostic(d: Diagnostic): RenderingContext = Impl(parameters(d))
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/AbstractFilteringTrace.kt

@@ -37,6 +37,8 @@ abstract class AbstractFilteringTrace(
     }
 
     override fun report(diagnostic: Diagnostic) {
+        diagnosticsCallback?.callback(diagnostic)
+
         parentTrace.report(diagnostic)
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/DelegatingBindingTrace.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.resolve
 import com.google.common.collect.ImmutableMap
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.resolve.diagnostics.BindingContextSuppressCache
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
@@ -156,6 +157,18 @@ open class DelegatingBindingTrace(
         mutableDiagnostics.report(diagnostic)
     }
 
+    protected var diagnosticsCallback: DiagnosticSink.DiagnosticsCallback? = null
+
+    override fun setCallback(callback: DiagnosticSink.DiagnosticsCallback) {
+        diagnosticsCallback = callback
+        mutableDiagnostics?.setCallback(callback)
+    }
+
+    override fun resetCallback() {
+        diagnosticsCallback = null
+        mutableDiagnostics?.resetCallback()
+    }
+
     override fun wantsDiagnostics(): Boolean = mutableDiagnostics != null
 
     override fun toString(): String = name

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

@@ -16,9 +16,10 @@
 
 package org.jetbrains.kotlin.resolve.diagnostics
 
-import com.intellij.psi.PsiElement
 import com.intellij.openapi.util.ModificationTracker
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.diagnostics.DiagnosticSink.DiagnosticsCallback
 import org.jetbrains.kotlin.diagnostics.GenericDiagnostics
 
 interface Diagnostics : GenericDiagnostics<Diagnostic> {
@@ -37,6 +38,9 @@ interface Diagnostics : GenericDiagnostics<Diagnostic> {
 
     fun noSuppression(): Diagnostics
 
+    fun setCallback(callback: DiagnosticsCallback) {}
+    fun resetCallback() {}
+
     companion object {
         val EMPTY: Diagnostics = object : Diagnostics {
             override fun noSuppression(): Diagnostics = this

compiler/frontend/src/org/jetbrains/kotlin/resolve/diagnostics/DiagnosticsWithSuppression.java

@@ -23,6 +23,7 @@ import kotlin.collections.CollectionsKt;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.TestOnly;
 import org.jetbrains.kotlin.diagnostics.Diagnostic;
+import org.jetbrains.kotlin.diagnostics.DiagnosticSink;
 
 import java.util.Collection;
 import java.util.Iterator;
@@ -73,6 +74,16 @@ public class DiagnosticsWithSuppression implements Diagnostics {
         throw new IllegalStateException("Trying to obtain modification tracker for readonly DiagnosticsWithSuppression.");
     }
 
+    @Override
+    public void setCallback(@NotNull DiagnosticSink.DiagnosticsCallback callback) {
+
+    }
+
+    @Override
+    public void resetCallback() {
+
+    }
+
     @TestOnly
     @NotNull
     public Collection<Diagnostic> getDiagnostics() {

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

@@ -24,10 +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.KtFile
-import org.jetbrains.kotlin.psi.KtStubbedPsiUtil
-import org.jetbrains.kotlin.psi.doNotAnalyze
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.constants.ArrayValue
 import org.jetbrains.kotlin.resolve.constants.StringValue
@@ -111,7 +108,12 @@ abstract class KotlinSuppressCache {
 
        This way we need no more lookups than the number of suppress() annotations from here to the root.
      */
-    private fun isSuppressedByAnnotated(suppressionKey: String, severity: Severity, annotated: KtAnnotated, debugDepth: Int): Boolean {
+    private fun isSuppressedByAnnotated(
+        suppressionKey: String,
+        severity: Severity,
+        annotated: KtAnnotated,
+        debugDepth: Int
+    ): Boolean {
         val suppressor = getOrCreateSuppressor(annotated)
         if (suppressor.isSuppressed(suppressionKey, severity)) return true
 
@@ -140,6 +142,7 @@ abstract class KotlinSuppressCache {
 
     private fun getSuppressingStrings(annotated: KtAnnotated): Set<String> {
         val builder = ImmutableSet.builder<String>()
+
         for (annotationDescriptor in getSuppressionAnnotations(annotated)) {
             processAnnotation(builder, annotationDescriptor)
         }

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

@@ -22,12 +22,16 @@ import com.intellij.psi.util.CachedValueProvider
 import com.intellij.util.CachedValueImpl
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.diagnostics.DiagnosticSink
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.KtStubbedPsiUtil
 
 class MutableDiagnosticsWithSuppression(
     private val suppressCache: KotlinSuppressCache,
     private val delegateDiagnostics: Diagnostics,
 ) : Diagnostics {
     private val diagnosticList = ArrayList<Diagnostic>()
+    private var diagnosticsCallback: DiagnosticSink.DiagnosticsCallback? = null
 
     //NOTE: CachedValuesManager is not used because it requires Project passed to this object
     private val cache = CachedValueImpl {
@@ -43,16 +47,38 @@ class MutableDiagnosticsWithSuppression(
     override fun forElement(psiElement: PsiElement) = readonlyView().forElement(psiElement)
     override fun noSuppression() = readonlyView().noSuppression()
 
+    override fun setCallback(callback: DiagnosticSink.DiagnosticsCallback) {
+        assert(diagnosticsCallback == null) { "diagnostic callback has been already registered" }
+        diagnosticsCallback = callback
+        delegateDiagnostics.setCallback(callback)
+    }
+
+    override fun resetCallback() {
+        diagnosticsCallback = null
+        delegateDiagnostics.resetCallback()
+    }
+
     //essential that this list is readonly
     fun getOwnDiagnostics(): List<Diagnostic> {
         return diagnosticList
     }
 
     fun report(diagnostic: Diagnostic) {
+        onFlyDiagnosticsCallback(diagnostic)?.callback(diagnostic)
+
         diagnosticList.add(diagnostic)
         modificationTracker.incModificationCount()
     }
 
+    private fun onFlyDiagnosticsCallback(diagnostic: Diagnostic): DiagnosticSink.DiagnosticsCallback? =
+        diagnosticsCallback.takeIf {
+            diagnosticsCallback != null &&
+                    // Due to a potential recursion in filter.invoke (via LazyAnnotations) do not try to report
+                    // diagnostic on fly if it happened in annotations
+                    KtStubbedPsiUtil.getPsiOrStubParent(diagnostic.psiElement, KtAnnotationEntry::class.java, false) == null &&
+                    suppressCache.filter.invoke(diagnostic)
+        }
+
     fun clear() {
         diagnosticList.clear()
         modificationTracker.incModificationCount()

compiler/frontend/src/org/jetbrains/kotlin/storage/LockBasedLazyResolveStorageManager.kt

@@ -22,7 +22,6 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.util.slicedMap.ReadOnlySlice
 import org.jetbrains.kotlin.util.slicedMap.WritableSlice
-import org.jetbrains.kotlin.resolve.TraceEntryFilter
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.psi.KtExpression