From ee794f86641de11a00649190fffa3f169477cda3 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Tue, 15 Sep 2020 17:46:48 +0300
Subject: [PATCH] FIR IDE: add check canceled between diagnostics

---
 .../fir/analysis/collectors/AbstractDiagnosticCollector.kt    | 3 +++
 .../api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt     | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollector.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollector.kt
index 10ae41e657d..522612e3483 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollector.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollector.kt
@@ -62,9 +62,12 @@ abstract class AbstractDiagnosticCollector(
         componentsInitialized = true
     }
 
+    protected open fun beforeCollecting() {}
+
     private inner class Visitor : FirDefaultVisitor<Unit, Nothing?>() {
         private fun <T : FirElement> T.runComponents() {
             components.forEach {
+                beforeCollecting()
                 this.accept(it, context)
             }
         }
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt
index 07256c7b2a2..5f2f746faee 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt
@@ -48,6 +48,10 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
         reporter = Reporter()
     }
 
+    override fun beforeCollecting() {
+        checkCanceled()
+    }
+
     override fun getCollectedDiagnostics(): Iterable<FirDiagnostic<*>> {
         // Not necessary in IDE
         return emptyList()