[FIR] extract exception handing into a service to have additional implementation for the ide

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/sessionFactoryHelpers.kt

@@ -14,11 +14,13 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirIdePredica
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirIdeRegisteredPluginAnnotations
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSourcesSession
-import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirExceptionHandler
 import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.moduleScopeProvider
 import org.jetbrains.kotlin.analysis.providers.createAnnotationResolver
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
+import org.jetbrains.kotlin.fir.FirExceptionHandler
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
@@ -39,6 +41,8 @@ internal fun LLFirSession.registerIdeComponents(project: Project) {
     register(IdeSessionComponents::class, IdeSessionComponents.create(this))
     register(FirCachesFactory::class, FirThreadSafeCachesFactory)
     register(SealedClassInheritorsProvider::class, project.createSealedInheritorsProvider())
+    register(SealedClassInheritorsProvider::class, project.createSealedInheritorsProvider())
+    register(FirExceptionHandler::class, LLFirExceptionHandler)
 }
 
 internal inline fun createCompositeSymbolProvider(

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/LLFirExceptionHandler.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.util
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirExceptionHandler
+import org.jetbrains.kotlin.fir.declarations.FirFile
+
+internal object LLFirExceptionHandler : FirExceptionHandler() {
+    override fun handleExceptionOnElementAnalysis(element: FirElement, throwable: Throwable): Nothing {
+        throw throwable
+    }
+
+    override fun handleExceptionOnFileAnalysis(file: FirFile, throwable: Throwable): Nothing {
+        throw throwable
+    }
+}

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt

@@ -77,6 +77,7 @@ fun FirSession.registerCliCompilerOnlyComponents() {
     register(FirCachesFactory::class, FirThreadUnsafeCachesFactory)
     register(SealedClassInheritorsProvider::class, SealedClassInheritorsProviderImpl)
     register(FirLazyDeclarationResolver::class, FirDummyCompilerLazyDeclarationResolver)
+    register(FirExceptionHandler::class, FirCliExceptionHandler)
 
     register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotationsImpl(this))
     register(FirPredicateBasedProvider::class, FirPredicateBasedProviderImpl(this))

compiler/fir/tree/src/org/jetbrains/kotlin/fir/Utils.kt

@@ -26,6 +26,8 @@ import org.jetbrains.kotlin.fir.types.builder.*
 import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.util.wrapIntoFileAnalysisExceptionIfNeeded
+import org.jetbrains.kotlin.util.wrapIntoSourceCodeAnalysisExceptionIfNeeded
 
 // TODO: rewrite
 fun FirBlock.returnExpressions(): List<FirExpression> = listOfNotNull(statements.lastOrNull() as? FirExpression)
@@ -137,16 +139,40 @@ fun FirDeclarationStatus.copy(
     }
 }
 
-inline fun <R> whileAnalysing(@Suppress("UNUSED_PARAMETER") session: FirSession, element: FirElement, block: () -> R) =
-    org.jetbrains.kotlin.util.whileAnalysing(element.source, block)
+inline fun <R> whileAnalysing(session: FirSession, element: FirElement, block: () -> R): R {
+    return try {
+        block()
+    } catch (throwable: Throwable) {
+        session.exceptionHandler.handleExceptionOnElementAnalysis(element, throwable)
+    }
+}
 
 inline fun <R> withFileAnalysisExceptionWrapping(file: FirFile, block: () -> R): R {
-    return org.jetbrains.kotlin.util.withFileAnalysisExceptionWrapping(
-        file.sourceFile?.path,
-        file.source,
-        { file.sourceFileLinesMapping?.getLineAndColumnByOffset(it) },
-        block,
-    )
+    return try {
+        block()
+    } catch (throwable: Throwable) {
+        file.moduleData.session.exceptionHandler.handleExceptionOnFileAnalysis(file, throwable)
+    }
+}
+
+abstract class FirExceptionHandler : FirSessionComponent {
+    abstract fun handleExceptionOnElementAnalysis(element: FirElement, throwable: Throwable): Nothing
+    abstract fun handleExceptionOnFileAnalysis(file: FirFile, throwable: Throwable): Nothing
+}
+
+val FirSession.exceptionHandler: FirExceptionHandler by FirSession.sessionComponentAccessor()
+
+object FirCliExceptionHandler : FirExceptionHandler() {
+    override fun handleExceptionOnElementAnalysis(element: FirElement, throwable: Throwable): Nothing {
+        throw throwable.wrapIntoSourceCodeAnalysisExceptionIfNeeded(element.source)
+    }
+
+    override fun handleExceptionOnFileAnalysis(file: FirFile, throwable: Throwable): Nothing {
+        throw throwable.wrapIntoFileAnalysisExceptionIfNeeded(
+            file.sourceFile?.path,
+            file.source
+        ) { file.sourceFileLinesMapping?.getLineAndColumnByOffset(it) }
+    }
 }
 
 @JvmInline

compiler/frontend.common/src/org/jetbrains/kotlin/util/AnalysisExceptions.kt

@@ -17,16 +17,8 @@ class SourceCodeAnalysisException(val source: KtSourceElement, override val caus
     override val message get() = cause.classNameAndMessage
 }
 
-inline fun <R> whileAnalysing(element: KtSourceElement?, block: () -> R): R {
-    return try {
-        block()
-    } catch (throwable: Throwable) {
-        throw throwable.wrapIntoSourceCodeAnalysisExceptionIfNeeded(element)
-    }
-}
 
-@PublishedApi
-internal fun Throwable.wrapIntoSourceCodeAnalysisExceptionIfNeeded(element: KtSourceElement?) = when (this) {
+fun Throwable.wrapIntoSourceCodeAnalysisExceptionIfNeeded(element: KtSourceElement?) = when (this) {
     is SourceCodeAnalysisException -> this
     is IndexNotReadyException -> this
     is ControlFlowException -> this
@@ -49,21 +41,7 @@ class FileAnalysisException(
         }
 }
 
-inline fun <R> withFileAnalysisExceptionWrapping(
-    filePath: String?,
-    fileSource: AbstractKtSourceElement?,
-    crossinline linesMapping: (Int) -> Pair<Int, Int>?,
-    block: () -> R,
-): R {
-    return try {
-        block()
-    } catch (throwable: Throwable) {
-        throw throwable.wrapIntoFileAnalysisExceptionIfNeeded(filePath, fileSource) { linesMapping(it) }
-    }
-}
-
-@PublishedApi
-internal fun Throwable.wrapIntoFileAnalysisExceptionIfNeeded(
+fun Throwable.wrapIntoFileAnalysisExceptionIfNeeded(
     filePath: String?,
     fileSource: AbstractKtSourceElement?,
     linesMapping: (Int) -> Pair<Int, Int>?,