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 7049e1ea1e9..e0060e5a68d 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
@@ -16,8 +16,6 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.fir.resolve.defaultType
-import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
-import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
@@ -26,7 +24,6 @@ import org.jetbrains.kotlin.name.Name
 abstract class AbstractDiagnosticCollector(
     override val session: FirSession,
     override val scopeSession: ScopeSession = ScopeSession(),
-    returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve()
 ) : SessionHolder {
     fun collectDiagnostics(firFile: FirFile): List<FirDiagnostic<*>> {
         if (!componentsInitialized) {
@@ -41,13 +38,11 @@ abstract class AbstractDiagnosticCollector(
     protected abstract fun getCollectedDiagnostics(): List<FirDiagnostic<*>>
     abstract val reporter: DiagnosticReporter
 
-    private val components: MutableList<AbstractDiagnosticCollectorComponent> = mutableListOf()
+    protected val components: MutableList<AbstractDiagnosticCollectorComponent> = mutableListOf()
     private var componentsInitialized = false
 
-    @Suppress("LeakingThis")
-    private val visitor = Visitor(PersistentCheckerContext(this, returnTypeCalculator))
+    protected abstract val visitor: DiagnosticCollectingVisitor
 
-    private var currentAction = DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
 
     fun initializeComponents(vararg components: AbstractDiagnosticCollectorComponent) {
         if (componentsInitialized) {
@@ -57,10 +52,16 @@ abstract class AbstractDiagnosticCollector(
         componentsInitialized = true
     }
 
-    protected open fun beforeRunningAllComponentsOnElement(element: FirElement) {}
-    protected open fun beforeRunningSingleComponentOnElement(element: FirElement) {}
 
-    private inner class Visitor(context: PersistentCheckerContext) : AbstractDiagnosticCollectorVisitor(context) {
+    open class DiagnosticCollectingVisitor(
+        context: PersistentCheckerContext,
+        private val components: List<AbstractDiagnosticCollectorComponent>
+    ) : AbstractDiagnosticCollectorVisitor(context) {
+        private var currentAction = DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
+
+        protected open fun beforeRunningAllComponentsOnElement(element: FirElement) {}
+        protected open fun beforeRunningSingleComponentOnElement(element: FirElement) {}
+
         private fun <T : FirElement> T.runComponents() {
             if (currentAction.checkInCurrentDeclaration) {
                 beforeRunningAllComponentsOnElement(this)
@@ -260,20 +261,20 @@ abstract class AbstractDiagnosticCollector(
                 getClassCall.acceptChildren(this, null)
             }
         }
-    }
 
-    protected open fun getDeclarationActionOnDeclarationEnter(declaration: FirDeclaration): DiagnosticCollectorDeclarationAction =
-        DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
+        protected open fun getDeclarationActionOnDeclarationEnter(declaration: FirDeclaration): DiagnosticCollectorDeclarationAction =
+            DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
 
-    protected open fun onDeclarationExit(declaration: FirDeclaration) {}
+        protected open fun onDeclarationExit(declaration: FirDeclaration) {}
 
-    private inline fun <R> withDiagnosticsAction(action: DiagnosticCollectorDeclarationAction, block: () -> R): R {
-        val oldAction = currentAction
-        currentAction = action
-        return try {
-            block()
-        } finally {
-            currentAction = oldAction
+        private inline fun <R> withDiagnosticsAction(action: DiagnosticCollectorDeclarationAction, block: () -> R): R {
+            val oldAction = currentAction
+            currentAction = action
+            return try {
+                block()
+            } finally {
+                currentAction = oldAction
+            }
         }
     }
 
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/SimpleDiagnosticsCollector.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/SimpleDiagnosticsCollector.kt
index 331eb909c6f..5d0bf7e0e68 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/SimpleDiagnosticsCollector.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/SimpleDiagnosticsCollector.kt
@@ -6,17 +6,26 @@
 package org.jetbrains.kotlin.fir.analysis.collectors
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.impl.BaseDiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.impl.DiagnosticReporterWithSuppress
 import org.jetbrains.kotlin.fir.analysis.diagnostics.impl.SimpleDiagnosticReporter
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 
 class SimpleDiagnosticsCollector(
     session: FirSession,
     scopeSession: ScopeSession,
     private val disableSuppress: Boolean = false
 ) : AbstractDiagnosticCollector(session, scopeSession) {
+    override val visitor: DiagnosticCollectingVisitor = DiagnosticCollectingVisitor(
+        PersistentCheckerContext(
+            this,
+            ReturnTypeCalculatorForFullBodyResolve()
+        ),
+        components
+    )
     override var reporter = createDiagnosticReporter()
         private set
 
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 11ad90ce80a..01c4b7bfe5a 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
@@ -10,13 +10,11 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.CheckersComponent
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
-import org.jetbrains.kotlin.fir.analysis.collectors.components.ControlFlowAnalysisDiagnosticComponent
-import org.jetbrains.kotlin.fir.analysis.collectors.components.DeclarationCheckersDiagnosticComponent
-import org.jetbrains.kotlin.fir.analysis.collectors.components.ErrorNodeDiagnosticCollectorComponent
-import org.jetbrains.kotlin.fir.analysis.collectors.components.ExpressionCheckersDiagnosticComponent
+import org.jetbrains.kotlin.fir.analysis.collectors.components.*
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
@@ -34,17 +32,25 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
     session: FirSession,
     useExtendedCheckers: Boolean,
 ) : AbstractDiagnosticCollector(
-    session,
-    returnTypeCalculator = createReturnTypeCalculatorForIDE(
-        session,
-        ScopeSession(),
-        ImplicitBodyResolveComputationSession(),
-        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
-    )
+    session
 ) {
     private val beforeElementDiagnosticCollectionHandler: BeforeElementDiagnosticCollectionHandler? =
         session.beforeElementDiagnosticCollectionHandler
 
+    @Suppress("LeakingThis")
+    override val visitor = run {
+        val returnTypeCalculator = createReturnTypeCalculatorForIDE(
+            session,
+            ScopeSession(),
+            ImplicitBodyResolveComputationSession(),
+            ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+        )
+        CollectingVisitor(
+            PersistentCheckerContext(this, returnTypeCalculator),
+            components
+        )
+    }
+
     init {
         val declarationCheckers = CheckersFactory.createDeclarationCheckers(useExtendedCheckers)
         val expressionCheckers = CheckersFactory.createExpressionCheckers(useExtendedCheckers)
@@ -74,13 +80,19 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
         reporter = Reporter()
     }
 
-    override fun beforeRunningSingleComponentOnElement(element: FirElement) {
-        checkCanceled()
+    inner class CollectingVisitor(
+        context: PersistentCheckerContext,
+        components: List<AbstractDiagnosticCollectorComponent>
+    ) : DiagnosticCollectingVisitor(context, components) {
+        override fun beforeRunningSingleComponentOnElement(element: FirElement) {
+            checkCanceled()
+        }
+
+        override fun beforeRunningAllComponentsOnElement(element: FirElement) {
+            beforeElementDiagnosticCollectionHandler?.beforeCollectingForElement(element)
+        }
     }
 
-    override fun beforeRunningAllComponentsOnElement(element: FirElement) {
-        beforeElementDiagnosticCollectionHandler?.beforeCollectingForElement(element)
-    }
 
     override fun getCollectedDiagnostics(): List<FirDiagnostic<*>> {
         // Not necessary in IDE