[FIR] KT-56846: Unwrap SourceCodeAnalysisException on transformFile

^KT-56846 Fixed

Merge-request: KT-MR-9065
Merged-by: Nikolay Lunyak <Nikolay.Lunyak@jetbrains.com>

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirAbstractPhaseTransformer.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 
 abstract class FirAbstractPhaseTransformer<D>(
     protected val baseTransformerPhase: FirResolvePhase
@@ -27,7 +28,9 @@ abstract class FirAbstractPhaseTransformer<D>(
 
     override fun transformFile(file: FirFile, data: D): FirFile {
         checkSessionConsistency(file)
-        return super.transformFile(file, data)
+        return withFileAnalysisExceptionWrapping(file) {
+            super.transformFile(file, data)
+        }
     }
 
     protected fun checkSessionConsistency(file: FirFile) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirImportResolveTransformer.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedParentInImport
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 
@@ -46,15 +47,16 @@ open class FirImportResolveTransformer protected constructor(
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
         checkSessionConsistency(file)
-        return file.also {
+        withFileAnalysisExceptionWrapping(file) {
             val prevValue = currentFile
             currentFile = file
             try {
-                it.transformChildren(this, null)
+                file.transformChildren(this, null)
             } finally {
                 currentFile = prevValue
             }
         }
+        return file
     }
 
     override fun transformImport(import: FirImport, data: Any?): FirImport {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSealedClassInheritorsProcessor.kt

@@ -96,7 +96,9 @@ class FirSealedClassInheritorsProcessor(
         }
 
         override fun transformFile(file: FirFile, data: Any?): FirFile {
-            return (file.transformChildren(this, data) as FirFile)
+            return withFileAnalysisExceptionWrapping(file) {
+                file.transformChildren(this, data) as FirFile
+            }
         }
 
         override fun transformRegularClass(regularClass: FirRegularClass, data: Any?): FirStatement {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.toSymbol
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.fir.whileAnalysing
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 
 class FirStatusResolveProcessor(
     session: FirSession,
@@ -199,7 +200,9 @@ abstract class AbstractFirStatusResolveTransformer(
     protected abstract fun FirDeclaration.needResolveNestedClassifiers(): Boolean
 
     override fun transformFile(file: FirFile, data: FirResolvedDeclarationStatus?): FirFile {
-        transformDeclarationContent(file, data)
+        withFileAnalysisExceptionWrapping(file) {
+            transformDeclarationContent(file, data)
+        }
         return file
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -65,9 +65,11 @@ open class FirSupertypeResolverTransformer(
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
         checkSessionConsistency(file)
-        file.accept(supertypeResolverVisitor, null)
-        supertypeComputationSession.breakLoops(session)
-        return file.transform(applySupertypesTransformer, null)
+        return withFileAnalysisExceptionWrapping(file) {
+            file.accept(supertypeResolverVisitor, null)
+            supertypeComputationSession.breakLoops(session)
+            file.transform(applySupertypesTransformer, null)
+        }
     }
 }
 
@@ -110,7 +112,9 @@ open class FirApplySupertypesTransformer(
     }
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
-        return transformDeclarationContent(file, null) as FirFile
+        return withFileAnalysisExceptionWrapping(file) {
+            transformDeclarationContent(file, null) as FirFile
+        }
     }
 
     override fun transformRegularClass(regularClass: FirRegularClass, data: Any?): FirStatement {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformerDispatcher.kt

@@ -39,10 +39,12 @@ abstract class FirAbstractBodyResolveTransformerDispatcher(
     override fun transformFile(file: FirFile, data: ResolutionMode): FirFile {
         checkSessionConsistency(file)
         return context.withFile(file, components) {
-            firTowerDataContextCollector?.addFileContext(file, context.towerDataContext)
+            withFileAnalysisExceptionWrapping(file) {
+                firTowerDataContextCollector?.addFileContext(file, context.towerDataContext)
 
-            @Suppress("UNCHECKED_CAST")
-            transformDeclarationContent(file, data) as FirFile
+                @Suppress("UNCHECKED_CAST")
+                transformDeclarationContent(file, data) as FirFile
+            }
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformerAdapters.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.AdapterForResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 
 @OptIn(AdapterForResolveProcessor::class)
 class FirBodyResolveProcessor(session: FirSession, scopeSession: ScopeSession) : FirTransformerBasedResolveProcessor(
@@ -36,6 +37,8 @@ class FirBodyResolveTransformerAdapter(session: FirSession, scopeSession: ScopeS
     }
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
-        return file.transform(transformer, ResolutionMode.ContextIndependent)
+        return withFileAnalysisExceptionWrapping(file) {
+            file.transform(transformer, ResolutionMode.ContextIndependent)
+        }
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt

@@ -60,7 +60,9 @@ class FirImplicitTypeBodyResolveTransformerAdapter(session: FirSession, scopeSes
     }
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
-        return file.transform(transformer, ResolutionMode.ContextIndependent)
+        return withFileAnalysisExceptionWrapping(file) {
+            file.transform(transformer, ResolutionMode.ContextIndependent)
+        }
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformerAdapter.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFull
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 
 @OptIn(AdapterForResolveProcessor::class)
 class FirContractResolveProcessor(session: FirSession, scopeSession: ScopeSession) : FirTransformerBasedResolveProcessor(
@@ -35,7 +36,9 @@ class FirContractResolveTransformerAdapter(session: FirSession, scopeSession: Sc
     }
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
-        return file.transform(transformer, ResolutionMode.ContextIndependent)
+        return withFileAnalysisExceptionWrapping(file) {
+            file.transform(transformer, ResolutionMode.ContextIndependent)
+        }
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/AbstractFirSpecificAnnotationResolveTransformer.kt

@@ -326,7 +326,9 @@ internal abstract class AbstractFirSpecificAnnotationResolveTransformer(
     inline fun <T> withFile(file: FirFile, f: () -> T): T {
         typeResolverTransformer.withFile(file) {
             argumentsTransformer.context.withFile(file, argumentsTransformer.components) {
-                return f()
+                return withFileAnalysisExceptionWrapping(file) {
+                    f()
+                }
             }
         }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsMappingProcessor.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.AdapterForResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 
 @OptIn(AdapterForResolveProcessor::class)
 class FirAnnotationArgumentsMappingProcessor(
@@ -32,6 +33,8 @@ class FirAnnotationArgumentsMappingTransformerAdapter(session: FirSession, scope
     }
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
-        return file.transform(transformer, ResolutionMode.ContextIndependent)
+        return withFileAnalysisExceptionWrapping(file) {
+            file.transform(transformer, ResolutionMode.ContextIndependent)
+        }
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsResolveProcessor.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.AdapterForResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 
 @OptIn(AdapterForResolveProcessor::class)
 class FirAnnotationArgumentsResolveProcessor(
@@ -32,6 +33,8 @@ class FirAnnotationArgumentsResolveTransformerAdapter(session: FirSession, scope
     }
 
     override fun transformFile(file: FirFile, data: Any?): FirFile {
-        return file.transform(transformer, ResolutionMode.ContextIndependent)
+        return withFileAnalysisExceptionWrapping(file) {
+            file.transform(transformer, ResolutionMode.ContextIndependent)
+        }
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompanionGenerationProcessor.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.name.SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
 
@@ -42,7 +43,9 @@ class FirCompanionGenerationTransformer(val session: FirSession) : FirTransforme
     override fun transformFile(file: FirFile, data: Nothing?): FirFile {
         // I don't want to use laziness here to prevent possible multi-threading problems
         if (generatedDeclarationProvider == null) return file
-        return file.transformDeclarations(this, data)
+        return withFileAnalysisExceptionWrapping(file) {
+            file.transformDeclarations(this, data)
+        }
     }
 
     override fun transformRegularClass(regularClass: FirRegularClass, data: Nothing?): FirStatement {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompilerRequiredAnnotationsResolveTransformer.kt

@@ -76,8 +76,10 @@ abstract class AbstractFirCompilerRequiredAnnotationsResolveTransformer(
     }
 
     override fun transformFile(file: FirFile, data: Nothing?): FirFile {
-        checkSessionConsistency(file)
-        file.resolveAnnotations()
+        withFileAnalysisExceptionWrapping(file) {
+            checkSessionConsistency(file)
+            file.resolveAnnotations()
+        }
         return file
     }
 

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

@@ -172,7 +172,7 @@ object FirCliExceptionHandler : FirExceptionHandler() {
     override fun handleExceptionOnFileAnalysis(file: FirFile, throwable: Throwable): Nothing {
         throw throwable.wrapIntoFileAnalysisExceptionIfNeeded(
             file.sourceFile?.path,
-            file.source
+            file.source,
         ) { file.sourceFileLinesMapping?.getLineAndColumnByOffset(it) }
     }
 }

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

@@ -5,11 +5,11 @@
 
 package org.jetbrains.kotlin.util
 
+import com.intellij.lang.LighterASTNode
 import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.project.IndexNotReadyException
-import org.jetbrains.kotlin.AbstractKtSourceElement
-import org.jetbrains.kotlin.KtRealSourceElementKind
-import org.jetbrains.kotlin.KtSourceElement
+import com.intellij.util.diff.FlyweightCapableTreeStructure
+import org.jetbrains.kotlin.*
 
 val Throwable.classNameAndMessage get() = "${this::class.qualifiedName}: $message"
 
@@ -43,22 +43,60 @@ class FileAnalysisException(
 
 fun Throwable.wrapIntoFileAnalysisExceptionIfNeeded(
     filePath: String?,
-    fileSource: AbstractKtSourceElement?,
+    fileSource: KtSourceElement?,
     linesMapping: (Int) -> Pair<Int, Int>?,
 ) = when {
-    filePath == null -> this
-
-    this is SourceCodeAnalysisException -> when (fileSource) {
-        source -> FileAnalysisException(filePath, cause)
+    filePath == null || fileSource == null -> when (this) {
+        is SourceCodeAnalysisException -> error("Sourceless FirFile contains a FirElement with a real source element")
+        else -> this
+    }
+    this is SourceCodeAnalysisException -> when {
+        fileSource == source -> FileAnalysisException(filePath, cause)
+        source.isDefinitelyNotInsideFile(fileSource) -> reportFileMismatch(source, fileSource, cause)
         else -> FileAnalysisException(filePath, cause, linesMapping(source.startOffset))
     }
-
     this is IndexNotReadyException -> this
     this is ControlFlowException -> this
     this is VirtualMachineError -> this
     else -> FileAnalysisException(filePath, this)
 }
 
+private fun KtSourceElement.isDefinitelyNotInsideFile(fileSource: KtSourceElement): Boolean {
+    val thisPsi = psi
+    val otherPsi = fileSource.psi
+
+    return when {
+        thisPsi != null && otherPsi != null -> thisPsi.containingFile != otherPsi
+        else -> !lighterASTNode.isInside(fileSource.lighterASTNode, treeStructure)
+    }
+}
+
+private fun reportFileMismatch(source: KtSourceElement, fileSource: KtSourceElement, cause: Throwable): Throwable {
+    val thisPsi = source.psi
+    val otherPsi = fileSource.psi
+    val comparison = "This:\n\n${source.text?.asQuote}\n\n...is not present in"
+
+    val expectedFileMessage = if (thisPsi != null && otherPsi != null) {
+        val actualPath = thisPsi.containingFile.virtualFile.path
+        val expectedPath = otherPsi.containingFile.virtualFile.path
+        "$expectedPath, but rather in $actualPath"
+    } else {
+        "...${fileSource.text?.asQuote}"
+    }
+
+    return IllegalStateException(
+        "KtSourceElement inside a SourceCodeAnalysisException was matched against the wrong FirFile source. $comparison$expectedFileMessage",
+        cause,
+    )
+}
+
+private val CharSequence.asQuote: String
+    get() = split("\n").joinToString("\n") { "> $it" }
+
+private fun LighterASTNode.isInside(other: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): Boolean {
+    return generateSequence(tree.getParent(this)) { tree.getParent(it) }.any { it == other }
+}
+
 inline fun <R> withSourceCodeAnalysisExceptionUnwrapping(block: () -> R): R {
     return try {
         block()