diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt index 9901f5a2481..cc71b2975c1 100644 --- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt +++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt @@ -1528,7 +1528,7 @@ class RawFirBuilder( val source = expression.toFirSourceElement(FirFakeSourceElementKind.ImplicitUnit) val result = expression.returnedExpression?.toFirExpression("Incorrect return expression") ?: buildUnitExpression { this.source = source } - return result.toReturn(source, expression.getTargetLabel()?.getReferencedName()) + return result.toReturn(source, expression.getTargetLabel()?.getReferencedName(), fromKtReturnExpression = true) } override fun visitTryExpression(expression: KtTryExpression, data: Unit): FirElement { diff --git a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt index be9bec961f1..23f169d1c74 100644 --- a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt +++ b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt @@ -148,7 +148,11 @@ abstract class BaseFirBuilder(val baseSession: FirSession, val context: Conte val ANONYMOUS_OBJECT_NAME = Name.special("") } - fun FirExpression.toReturn(baseSource: FirSourceElement? = source, labelName: String? = null): FirReturnExpression { + fun FirExpression.toReturn( + baseSource: FirSourceElement? = source, + labelName: String? = null, + fromKtReturnExpression: Boolean = false + ): FirReturnExpression { return buildReturnExpression { fun FirFunctionTarget.bindToErrorFunction(message: String, kind: DiagnosticKind) { bind( @@ -162,7 +166,9 @@ abstract class BaseFirBuilder(val baseSession: FirSession, val context: Conte ) } - source = baseSource?.fakeElement(FirFakeSourceElementKind.ImplicitReturn) + source = + if (fromKtReturnExpression) baseSource?.realElement() + else baseSource?.fakeElement(FirFakeSourceElementKind.ImplicitReturn) result = this@toReturn if (labelName == null) { target = context.firFunctionTargets.lastOrNull { !it.isLambda } ?: FirFunctionTarget(labelName, isLambda = false).apply { diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt index d0907c15124..cc05f1bb283 100644 --- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt +++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt @@ -264,6 +264,13 @@ fun FirSourceElement.fakeElement(newKind: FirFakeSourceElementKind): FirSourceEl } } +fun FirSourceElement.realElement(): FirSourceElement = when (this) { + is FirRealPsiSourceElement<*> -> this + is FirLightSourceElement -> FirLightSourceElement(lighterASTNode, startOffset, endOffset, treeStructure, FirRealSourceElementKind) + is FirPsiSourceElement<*> -> FirRealPsiSourceElement(psi) +} + + class FirLightSourceElement( override val lighterASTNode: LighterASTNode, override val startOffset: Int,