[FIR] Infer type of when expressions and try expressions like it is function call

compiler/fir/cones/src/org/jetbrains/kotlin/fir/symbols/SyntheticCallableId.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2019 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.fir.symbols
+
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+object SyntheticCallableId {
+    val WHEN = CallableId(
+        FqName("_synthetic"),
+        Name.identifier("WHEN_CALL")
+    )
+    val TRY = CallableId(
+        FqName("_synthetic"),
+        Name.identifier("TRY_CALL")
+    )
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/CopyUtils.kt

@@ -11,11 +11,10 @@ import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirAnonymousFunctionImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirTypeParameterImpl
-import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
-import org.jetbrains.kotlin.fir.expressions.FirBlock
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirFunctionCallImpl
+import org.jetbrains.kotlin.fir.expressions.impl.FirTryExpressionImpl
+import org.jetbrains.kotlin.fir.expressions.impl.FirWhenExpressionImpl
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
@@ -79,4 +78,21 @@ fun FirTypeParameter.copy(
         this.bounds += bounds
         this.annotations += annotations
     }
+}
+
+fun FirWhenExpression.copy(
+    resultType: FirTypeRef = this.typeRef,
+    calleeReference: FirReference = this.calleeReference
+): FirWhenExpressionImpl = FirWhenExpressionImpl(psi, subject, subjectVariable, calleeReference).apply {
+    this@apply.branches.addAll(this@copy.branches)
+    this.typeRef = resultType
+    this.calleeReference = calleeReference
+}
+
+fun FirTryExpression.copy(
+    resultType: FirTypeRef = this.typeRef,
+    calleeReference: FirReference = this.calleeReference
+): FirTryExpressionImpl = FirTryExpressionImpl(psi, tryBlock, finallyBlock, calleeReference).apply {
+    this@apply.catches.addAll(this@copy.catches)
+    this.typeRef = resultType
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -36,12 +36,10 @@ import org.jetbrains.kotlin.resolve.calls.results.TypeSpecificityComparator
 
 class FirCallResolver(
     private val transformer: FirBodyResolveTransformer,
-    private val inferenceComponents: InferenceComponents,
     private val scopes: List<FirScope>,
     private val localScopes: List<FirScope>,
     private val implicitReceiverStack: List<ImplicitReceiverValue>,
-    private val qualifiedResolver: FirQualifiedNameResolver,
-    private val resolutionStageRunner: ResolutionStageRunner
+    private val qualifiedResolver: FirQualifiedNameResolver
 ) : BodyResolveComponents by transformer {
 
     fun resolveCallAndSelectCandidate(functionCall: FirFunctionCall, expectedTypeRef: FirTypeRef?, file: FirFile): FirFunctionCall {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SessionHolder.kt

@@ -8,7 +8,10 @@ package org.jetbrains.kotlin.fir.resolve
 import com.google.common.collect.SetMultimap
 import org.jetbrains.kotlin.fir.FirSession
 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.resolve.calls.InferenceComponents
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.phasedFir
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
@@ -26,6 +29,10 @@ interface BodyResolveComponents : SessionHolder {
     val labels: SetMultimap<Name, ConeKotlinType>
     val noExpectedType: FirTypeRef
     val symbolProvider: FirSymbolProvider
+    val file: FirFile
+    val container: FirDeclaration
+    val inferenceComponents: InferenceComponents
+    val resolutionStageRunner: ResolutionStageRunner
 
     val <D> AbstractFirBasedSymbol<D>.phasedFir: D where D : FirDeclaration, D : FirSymbolOwner<D>
         get() = phasedFir(session, FirResolvePhase.DECLARATIONS)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -9,9 +9,10 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.resolve.inference.returnExpressions
+import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
 import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
 import org.jetbrains.kotlin.fir.resolve.withNullability
+import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.addSubtypeConstraintIfCompatible
@@ -38,7 +39,15 @@ fun resolveArgumentExpression(
     typeProvider: (FirExpression) -> FirTypeRef?
 ) {
     return when (argument) {
-        is FirFunctionCall -> resolveSubCallArgument(csBuilder, argument, expectedType, sink, isReceiver, isSafeCall, typeProvider)
+        is FirFunctionCall, is FirCallLikeControlFlowExpression -> resolveSubCallArgument(
+            csBuilder,
+            argument as FirResolvable,
+            expectedType,
+            sink,
+            isReceiver,
+            isSafeCall,
+            typeProvider
+        )
         is FirQualifiedAccessExpression -> resolvePlainExpressionArgument(
             csBuilder,
             argument,
@@ -65,10 +74,40 @@ fun resolveArgumentExpression(
             acceptLambdaAtoms,
             typeProvider
         )
+        is FirBlock -> resolveBlockArgument(csBuilder, argument, expectedType, expectedTypeRef, sink, isReceiver, isSafeCall, acceptLambdaAtoms, typeProvider)
         else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, isSafeCall, typeProvider)
     }
 }
 
+private fun resolveBlockArgument(
+    csBuilder: ConstraintSystemBuilder,
+    block: FirBlock,
+    expectedType: ConeKotlinType,
+    expectedTypeRef: FirTypeRef,
+    sink: CheckerSink,
+    isReceiver: Boolean,
+    isSafeCall: Boolean,
+    acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit,
+    typeProvider: (FirExpression) -> FirTypeRef?
+) {
+    val returnArguments = block.returnExpressions()
+    if (returnArguments.isEmpty()) {
+        checkApplicabilityForArgumentType(
+            csBuilder,
+            block.typeRef.coneTypeUnsafe(),
+            expectedType.type,
+            SimpleConstraintSystemConstraintPosition,
+            false,
+            expectedType.type.withNullability(ConeNullability.NULLABLE),
+            sink
+        )
+        return
+    }
+    for (argument in returnArguments) {
+        resolveArgumentExpression(csBuilder, argument, expectedType, expectedTypeRef, sink, isReceiver, isSafeCall, acceptLambdaAtoms, typeProvider)
+    }
+}
+
 fun resolveSubCallArgument(
     csBuilder: ConstraintSystemBuilder,
     argument: FirResolvable,
@@ -116,12 +155,23 @@ fun resolvePlainArgumentType(
         return
     }
 
+    checkApplicabilityForArgumentType(csBuilder, argumentType, expectedType, position, isReceiver, nullableExpectedType, sink)
+}
+
+private fun checkApplicabilityForArgumentType(
+    csBuilder: ConstraintSystemBuilder,
+    argumentType: ConeKotlinType,
+    expectedType: ConeKotlinType,
+    position: SimpleConstraintSystemConstraintPosition,
+    isReceiver: Boolean,
+    nullableExpectedType: ConeKotlinType,
+    sink: CheckerSink
+) {
     if (!csBuilder.addSubtypeConstraintIfCompatible(argumentType, expectedType, position)) {
         if (!isReceiver) {
             if (!csBuilder.addSubtypeConstraintIfCompatible(argumentType, nullableExpectedType, position)) {
                 csBuilder.addSubtypeConstraint(argumentType, expectedType, position)
             }
-
             return
         }
         if (csBuilder.addSubtypeConstraintIfCompatible(argumentType, nullableExpectedType, position)) {
@@ -130,7 +180,6 @@ fun resolvePlainArgumentType(
             csBuilder.addSubtypeConstraint(argumentType, expectedType, position)
             sink.reportApplicability(CandidateApplicability.WRONG_RECEIVER)
         }
-
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CandidateFactory.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzer
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
@@ -42,8 +43,9 @@ class CandidateFactory(
 
 fun PostponedArgumentsAnalyzer.Context.addSubsystemFromExpression(expression: FirExpression) {
     when (expression) {
-        is FirFunctionCall -> expression.candidate()?.let { addOtherSystem(it.system.asReadOnlyStorage()) }
+        is FirFunctionCall, is FirCallLikeControlFlowExpression -> (expression as FirResolvable).candidate()?.let { addOtherSystem(it.system.asReadOnlyStorage()) }
         is FirWrappedArgumentExpression -> addSubsystemFromExpression(expression.expression)
+        is FirBlock -> expression.returnExpressions().forEach { addSubsystemFromExpression(it) }
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/InferenceCompletion.kt

@@ -5,10 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls
 
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.expressions.FirWrappedArgumentExpression
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
@@ -182,6 +179,23 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
                     }
                     this.arguments.forEach { it.process(to) }
                 }
+                is FirWhenExpression -> {
+                    val candidate = (this.calleeReference as? FirNamedReferenceWithCandidate)?.candidate
+                    candidate?.postponedAtoms?.forEach {
+                        to.addIfNotNull(it.safeAs<PostponedResolvedAtomMarker>()?.takeUnless { it.analyzed })
+                    }
+                    this.branches.forEach { it.result.process(to) }
+                }
+
+                is FirTryExpression -> {
+                    val candidate = (this.calleeReference as? FirNamedReferenceWithCandidate)?.candidate
+                    candidate?.postponedAtoms?.forEach {
+                        to.addIfNotNull(it.safeAs<PostponedResolvedAtomMarker>()?.takeUnless { it.analyzed })
+                    }
+                    tryBlock.process(to)
+                    catches.forEach { it.block.process(to) }
+                }
+
                 is FirWrappedArgumentExpression -> this.expression.process(to)
                 // TOOD: WTF?
             }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Synthetics.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirNamedFunction
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyGetter
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertySetter
@@ -16,6 +17,7 @@ import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.symbols.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.name.Name
 
@@ -23,6 +25,11 @@ interface SyntheticSymbol : ConeSymbol
 
 class SyntheticPropertySymbol(callableId: CallableId) : FirPropertySymbol(callableId), SyntheticSymbol
 
+class FirSyntheticFunctionSymbol(
+    callableId: CallableId,
+    val file: FirFile
+) : FirNamedFunctionSymbol(callableId), SyntheticSymbol
+
 class FirSyntheticPropertiesScope(
     val session: FirSession,
     private val baseScope: FirScope,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -8,9 +8,8 @@ package org.jetbrains.kotlin.fir.resolve.inference
 import org.jetbrains.kotlin.fir.copy
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
-import org.jetbrains.kotlin.fir.expressions.FirBlock
-import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -18,6 +17,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.resolve.transformers.MapArguments
 import org.jetbrains.kotlin.fir.resolve.typeFromCallee
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
+import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.transformSingle
 import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -33,8 +33,7 @@ import org.jetbrains.kotlin.types.model.StubTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 class FirCallCompleter(
-    private val transformer: FirBodyResolveTransformer,
-    private val inferenceComponents: InferenceComponents
+    private val transformer: FirBodyResolveTransformer
 ) : BodyResolveComponents by transformer {
     fun <T : FirResolvable> completeCall(call: T, expectedTypeRef: FirTypeRef?): T {
         val typeRef = typeFromCallee(call)
@@ -125,4 +124,4 @@ class FirCallCompleter(
             return (newLambdaExpression.body?.returnExpressions() ?: emptyList()) to InferenceSession.default
         }
     }
-}
+}

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirTopLevelDeclaredMemberScope
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.symbols.*
+import org.jetbrains.kotlin.fir.symbols.StandardClassIds.Int
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
@@ -45,41 +46,42 @@ open class FirBodyResolveTransformer(
     val scopeSession: ScopeSession = ScopeSession()
 ) : FirAbstractPhaseTransformer<Any?>(phase), BodyResolveComponents {
     final override val returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorWithJump(session, scopeSession)
-    override val labels: SetMultimap<Name, ConeKotlinType> = LinkedHashMultimap.create()
-    override val noExpectedType = FirImplicitTypeRefImpl(null)
+    final override val labels: SetMultimap<Name, ConeKotlinType> = LinkedHashMultimap.create()
+    final override val noExpectedType = FirImplicitTypeRefImpl(null)
 
-    override val symbolProvider = session.service<FirSymbolProvider>()
+    final override val symbolProvider = session.service<FirSymbolProvider>()
     val scopes = mutableListOf<FirScope>()
 
     private var packageFqName = FqName.ROOT
-    private lateinit var file: FirFile
+    final override lateinit var file: FirFile
+        private set
 
     private var _container: FirDeclaration? = null
-    internal var container: FirDeclaration
+    final override var container: FirDeclaration
         get() = _container!!
-        set(value) {
+        private set(value) {
             _container = value
         }
 
     private val localScopes = mutableListOf<FirLocalScope>()
     private val implicitReceiverStack = mutableListOf<ImplicitReceiverValue>()
-    private val inferenceComponents = inferenceComponents(session, returnTypeCalculator, scopeSession)
+    final override val inferenceComponents = inferenceComponents(session, returnTypeCalculator, scopeSession)
 
     private var primaryConstructorParametersScope: FirLocalScope? = null
 
-    private val callCompleter: FirCallCompleter = FirCallCompleter(this, inferenceComponents)
+    private val callCompleter: FirCallCompleter = FirCallCompleter(this)
     private val qualifiedResolver: FirQualifiedNameResolver = FirQualifiedNameResolver(this)
-    private val resolutionStageRunner: ResolutionStageRunner = ResolutionStageRunner(inferenceComponents)
+    final override val resolutionStageRunner: ResolutionStageRunner = ResolutionStageRunner(inferenceComponents)
     private val callResolver: FirCallResolver = FirCallResolver(
         this,
-        inferenceComponents,
         scopes,
         localScopes,
         implicitReceiverStack,
-        qualifiedResolver,
-        resolutionStageRunner
+        qualifiedResolver
     )
 
+    private val syntheticCallGenerator: FirSyntheticCallGenerator = FirSyntheticCallGenerator(this)
+
     override val <D> AbstractFirBasedSymbol<D>.phasedFir: D where D : FirDeclaration, D : FirSymbolOwner<D>
         get() {
             val requiredPhase = transformerPhase.prev
@@ -356,21 +358,28 @@ open class FirBodyResolveTransformer(
 
     data class LambdaResolution(val expectedReturnTypeRef: FirResolvedTypeRef?)
 
-    override fun transformTryExpression(tryExpression: FirTryExpression, data: Any?): CompositeTransformResult<FirStatement> {
-        @Suppress("NAME_SHADOWING")
-        val tryExpression = tryExpression.transformChildren(this, data) as FirTryExpression
-        if (tryExpression.resultType !is FirResolvedTypeRef) {
-            val type = commonSuperType((listOf(tryExpression.tryBlock) + tryExpression.catches.map { it.block }).map {
-                val expression = it.statements.lastOrNull() as? FirExpression
-                if (expression != null) {
-                    (expression.resultType as? FirResolvedTypeRef) ?: FirErrorTypeRefImpl(null, "No type for when branch result")
-                } else {
-                    FirImplicitUnitTypeRef(null)
-                }
-            })
-            if (type != null) tryExpression.resultType = type
+
+    override fun transformCatch(catch: FirCatch, data: Any?): CompositeTransformResult<FirCatch> {
+        return withScopeCleanup(localScopes) {
+            localScopes += FirLocalScope()
+            catch.transformParameter(this, noExpectedType)
+            catch.transformBlock(this, null).compose()
         }
-        return tryExpression.compose()
+    }
+
+    override fun transformTryExpression(tryExpression: FirTryExpression, data: Any?): CompositeTransformResult<FirStatement> {
+        if (tryExpression.calleeReference is FirResolvedCallableReference && tryExpression.resultType !is FirImplicitTypeRef) {
+            return tryExpression.compose()
+        }
+        tryExpression.transformTryBlock(this, null)
+        tryExpression.transformCatches(this, null)
+
+        @Suppress("NAME_SHADOWING")
+        val tryExpression = syntheticCallGenerator.generateCalleeForTryExpression(tryExpression) ?: return tryExpression.compose()
+        val expectedTypeRef = data as FirTypeRef?
+        val result = callCompleter.completeCall(tryExpression, expectedTypeRef)
+
+        return result.transformFinallyBlock(this, noExpectedType).compose()
     }
 
     override fun transformFunctionCall(functionCall: FirFunctionCall, data: Any?): CompositeTransformResult<FirStatement> {
@@ -435,19 +444,27 @@ open class FirBodyResolveTransformer(
     }
 
     override fun transformWhenExpression(whenExpression: FirWhenExpression, data: Any?): CompositeTransformResult<FirStatement> {
-        whenExpression.transformChildren(this, data)
-        if (whenExpression.resultType !is FirResolvedTypeRef) {
-            val type = commonSuperType(whenExpression.branches.map {
-                val expression = it.result.statements.lastOrNull() as? FirExpression
-                if (expression != null) {
-                    (expression.resultType as? FirResolvedTypeRef) ?: FirErrorTypeRefImpl(null, "No type for when branch result")
-                } else {
-                    FirImplicitUnitTypeRef(null)
-                }
-            })
-            if (type != null) whenExpression.resultType = type
+        if (whenExpression.calleeReference is FirResolvedCallableReference && whenExpression.resultType !is FirImplicitTypeRef) {
+            return whenExpression.compose()
+        }
+
+        return withScopeCleanup(localScopes) with@{
+            if (whenExpression.subjectVariable != null) {
+                localScopes += FirLocalScope()
+            }
+            whenExpression.transformSubject(this, noExpectedType)
+            whenExpression.transformBranches(this, null)
+
+            @Suppress("NAME_SHADOWING")
+            val whenExpression = syntheticCallGenerator.generateCalleeForWhenExpression(whenExpression) ?: run {
+                // TODO: bodies will be unresolved. Maybe run usual transform without completer?
+                return@with whenExpression.compose()
+            }
+
+            val expectedTypeRef = data as FirTypeRef?
+            val result = callCompleter.completeCall(whenExpression, expectedTypeRef)
+            result.compose()
         }
-        return whenExpression.compose()
     }
 
     override fun transformWhenSubjectExpression(

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

@@ -5,19 +5,21 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers
 
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.copy
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirResolvedCallableReferenceImpl
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.candidate
 import org.jetbrains.kotlin.fir.resolve.constructFunctionalTypeRef
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substituteOrNull
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.FirTypeProjectionWithVariance
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.fir.types.impl.FirTypeProjectionWithVarianceImpl
@@ -138,4 +140,66 @@ class FirCallCompletionResultsWriterTransformer(
         return super.transformAnonymousFunction(anonymousFunction, data)
     }
 
+    override fun transformBlock(block: FirBlock, data: Nothing?): CompositeTransformResult<FirStatement> {
+        val initialType = block.resultType.coneTypeSafe<ConeKotlinType>()
+        if (initialType != null) {
+            val finalType = finalSubstitutor.substituteOrNull(initialType)
+            val resultType = block.resultType.withReplacedConeType(finalType)
+            block.replaceTypeRef(resultType)
+        }
+        return super.transformBlock(block, data)
+    }
+
+    override fun transformWhenExpression(whenExpression: FirWhenExpression, data: Nothing?): CompositeTransformResult<FirStatement> {
+        val calleeReference = whenExpression.calleeReference as? FirNamedReferenceWithCandidate ?: return whenExpression.compose()
+
+        val whenExpression = whenExpression.transformChildren(this, data) as FirWhenExpression
+
+        val declaration = whenExpression.candidate()?.symbol?.fir as? FirMemberFunction<*> ?: return whenExpression.compose()
+
+        val subCandidate = calleeReference.candidate
+
+        val typeRef = typeCalculator.tryCalculateReturnType(declaration)
+
+        val initialType = subCandidate.substitutor.substituteOrNull(typeRef.type)
+        val finalType = finalSubstitutor.substituteOrNull(initialType)
+
+        val resultType = typeRef.withReplacedConeType(finalType)
+
+        return whenExpression.copy(
+            resultType = resultType,
+            calleeReference = FirResolvedCallableReferenceImpl(
+                calleeReference.psi,
+                calleeReference.name,
+                calleeReference.candidateSymbol
+            )
+        ).compose()
+    }
+
+    override fun transformTryExpression(tryExpression: FirTryExpression, data: Nothing?): CompositeTransformResult<FirStatement> {
+        val calleeReference = tryExpression.calleeReference as? FirNamedReferenceWithCandidate ?: return tryExpression.compose()
+
+        val tryExpression = tryExpression.transformChildren(this, data) as FirTryExpression
+
+        val declaration = tryExpression.candidate()?.symbol?.fir as? FirMemberFunction<*> ?: return tryExpression.compose()
+
+        val subCandidate = calleeReference.candidate
+
+        val typeRef = typeCalculator.tryCalculateReturnType(declaration)
+
+        val initialType = subCandidate.substitutor.substituteOrNull(typeRef.type)
+        val finalType = finalSubstitutor.substituteOrNull(initialType)
+
+        val resultType = typeRef.withReplacedConeType(finalType)
+
+        return tryExpression.copy(
+            resultType = resultType,
+            calleeReference = FirResolvedCallableReferenceImpl(
+                calleeReference.psi,
+                calleeReference.name,
+                calleeReference.candidateSymbol
+            )
+        ).compose()
+
+    }
 }

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

@@ -0,0 +1,174 @@
+/*
+ * Copyright 2010-2019 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.fir.resolve.transformers
+
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.copy
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.impl.FirMemberFunctionImpl
+import org.jetbrains.kotlin.fir.declarations.impl.FirTypeParameterImpl
+import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirTryExpression
+import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
+import org.jetbrains.kotlin.fir.references.FirStubReference
+import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.calls.*
+import org.jetbrains.kotlin.fir.symbols.SyntheticCallableId
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
+import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
+import org.jetbrains.kotlin.fir.types.impl.FirTypeProjectionWithVarianceImpl
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.types.Variance
+
+class FirSyntheticCallGenerator(private val transformer: FirBodyResolveTransformer) : BodyResolveComponents by transformer {
+    fun generateCalleeForWhenExpression(whenExpression: FirWhenExpression): FirWhenExpression? {
+        val stubReference = whenExpression.calleeReference
+        assert(stubReference is FirStubReference)
+
+        val function = generateWhenCallDeclaration()
+
+        val arguments = whenExpression.branches.map { it.result }
+        val reference = generateCalleeReferenceWithCandidate(function, arguments, SyntheticCallableId.WHEN.callableName) ?: return null // TODO
+
+        return whenExpression.copy(calleeReference = reference)
+    }
+
+    fun generateCalleeForTryExpression(tryExpression: FirTryExpression): FirTryExpression? {
+        val stubReference = tryExpression.calleeReference
+        assert(stubReference is FirStubReference)
+
+        val function = generateTryCallDeclaration(tryExpression)
+        val arguments = mutableListOf<FirExpression>()
+
+        with(tryExpression) {
+            arguments += tryBlock
+            catches.map {
+                arguments += it.block
+            }
+        }
+
+        val reference = generateCalleeReferenceWithCandidate(function, arguments, SyntheticCallableId.TRY.callableName) ?: return null // TODO
+
+        return tryExpression.copy(calleeReference = reference)
+    }
+
+    private fun generateTryCallDeclaration(tryExpression: FirTryExpression): FirMemberFunctionImpl {
+        val functionSymbol = FirSyntheticFunctionSymbol(SyntheticCallableId.TRY, file)
+        val typeParameterSymbol = FirTypeParameterSymbol()
+        val typeParameter = FirTypeParameterImpl(session, null, typeParameterSymbol, Name.identifier("K"), Variance.INVARIANT, false)
+
+        val returnType = FirResolvedTypeRefImpl(null, ConeTypeParameterTypeImpl(typeParameterSymbol.toLookupTag(), false))
+
+        val tryType = FirResolvedTypeRefImpl(null, returnType.coneTypeUnsafe())
+        val catchTypes = tryExpression.catches.map {
+            FirResolvedTypeRefImpl(null, returnType.coneTypeUnsafe())
+        }
+
+        val typeArgument = FirTypeProjectionWithVarianceImpl(null, Variance.INVARIANT, returnType)
+
+        return generateMemberFunction(session, functionSymbol, SyntheticCallableId.TRY.callableName, typeArgument.typeRef).apply {
+            typeParameters += typeParameter
+
+            valueParameters += tryType.toValueParameter(session, "tryBlock")
+            catchTypes.forEachIndexed { i, type ->
+                valueParameters += type.toValueParameter(session, "catchBlock_$i")
+            }
+        }
+    }
+
+    private fun generateWhenCallDeclaration(): FirMemberFunctionImpl {
+        val functionSymbol = FirSyntheticFunctionSymbol(SyntheticCallableId.WHEN, file)
+        val typeParameterSymbol = FirTypeParameterSymbol()
+        val typeParameter = FirTypeParameterImpl(session, null, typeParameterSymbol, Name.identifier("K"), Variance.INVARIANT, false)
+
+        val returnType = FirResolvedTypeRefImpl(null, ConeTypeParameterTypeImpl(typeParameterSymbol.toLookupTag(), false))
+        val branchType = FirResolvedTypeRefImpl(null, returnType.coneTypeUnsafe<ConeKotlinType>().createArrayOf(session))
+
+        val typeArgument = FirTypeProjectionWithVarianceImpl(null, Variance.INVARIANT, returnType)
+
+        return generateMemberFunction(session, functionSymbol, SyntheticCallableId.WHEN.callableName, typeArgument.typeRef).apply {
+            typeParameters += typeParameter
+            valueParameters += branchType.toValueParameter(session, "branches", true)
+        }
+    }
+
+    private fun generateCalleeReferenceWithCandidate(
+        function: FirMemberFunctionImpl,
+        arguments: List<FirExpression>,
+        name: Name
+    ): FirNamedReferenceWithCandidate? {
+        val callInfo = generateCallInfo(arguments)
+        val candidate = generateCandidate(callInfo, function)
+        val applicability = resolutionStageRunner.processCandidate(candidate)
+        if (applicability <= CandidateApplicability.INAPPLICABLE) {
+            return null
+        }
+
+        return FirNamedReferenceWithCandidate(null, name, candidate)
+    }
+
+    private fun generateCandidate(callInfo: CallInfo, function: FirMemberFunctionImpl): Candidate =
+        CandidateFactory(inferenceComponents, callInfo).createCandidate(
+            symbol = function.symbol,
+            dispatchReceiverValue = null,
+            implicitExtensionReceiverValue = null,
+            explicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER
+        )
+
+    private fun generateCallInfo(arguments: List<FirExpression>) = CallInfo(
+        callKind = CallKind.Function,
+        explicitReceiver = null,
+        arguments = arguments,
+        isSafeCall = false,
+        typeArguments = emptyList(),
+        session = session,
+        containingFile = file,
+        container = container
+    ) { it.resultType }
+
+    private fun generateMemberFunction(session: FirSession, symbol: FirNamedFunctionSymbol, name: Name, returnType: FirTypeRef) =
+        FirMemberFunctionImpl(
+            session = session,
+            psi = null,
+            symbol = symbol,
+            name = name,
+            visibility = Visibilities.PUBLIC,
+            modality = Modality.FINAL,
+            isExpect = false,
+            isActual = false,
+            isOverride = false,
+            isOperator = false,
+            isInfix = false,
+            isInline = false,
+            isTailRec = false,
+            isExternal = false,
+            isSuspend = false,
+            receiverTypeRef = null,
+            returnTypeRef = returnType
+        ).apply {
+            this.resolvePhase = FirResolvePhase.BODY_RESOLVE
+        }
+
+    private fun FirResolvedTypeRef.toValueParameter(session: FirSession, name: String, isVararg: Boolean = false) = FirValueParameterImpl(
+        session = session,
+        psi = null,
+        name = Name.identifier(name),
+        returnTypeRef = this,
+        defaultValue = null,
+        isCrossinline = false,
+        isNoinline = false,
+        isVararg = isVararg
+    ).apply {
+        this.resolvePhase = FirResolvePhase.BODY_RESOLVE
+    }
+}

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.FirProvider
+import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeSymbol
@@ -24,6 +25,7 @@ fun <D> AbstractFirBasedSymbol<D>.phasedFir(
     if (availablePhase < requiredPhase) {
         val provider = FirProvider.getInstance(session)
         val containingFile = when (this) {
+            is FirSyntheticFunctionSymbol -> file
             is ConeCallableSymbol -> provider.getFirCallableContainerFile(this)
             is ConeClassLikeSymbol -> provider.getFirClassifierContainerFile(this)
             else -> null

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirCallLikeControlFlowExpression.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2019 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.fir.expressions
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.VisitedSupertype
+import org.jetbrains.kotlin.fir.expressions.impl.FirUnknownTypeExpression
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+
+abstract class FirCallLikeControlFlowExpression(psi: PsiElement?) : @VisitedSupertype FirUnknownTypeExpression(psi), FirResolvable {
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        super<FirUnknownTypeExpression>.acceptChildren(visitor, data)
+    }
+
+    override fun accept(visitor: FirVisitorVoid) {
+        super<FirUnknownTypeExpression>.accept(visitor)
+    }
+
+    override fun acceptChildren(visitor: FirVisitorVoid) {
+        super<FirUnknownTypeExpression>.acceptChildren(visitor)
+    }
+
+    override fun <E : FirElement, D> transform(visitor: FirTransformer<D>, data: D): CompositeTransformResult<E> {
+        return super<FirUnknownTypeExpression>.transform(visitor, data)
+    }
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
+        return super<FirUnknownTypeExpression>.transformChildren(transformer, data)
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirCatch.kt

@@ -5,10 +5,13 @@
 
 package org.jetbrains.kotlin.fir.expressions
 
+import org.jetbrains.kotlin.fir.BaseTransformedType
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
+@BaseTransformedType
 interface FirCatch : FirElement {
     val parameter: FirValueParameter
 
@@ -21,4 +24,8 @@ interface FirCatch : FirElement {
         parameter.accept(visitor, data)
         block.accept(visitor, data)
     }
+
+    fun <D> transformParameter(transformer: FirTransformer<D>, data: D): FirCatch
+
+    fun <D> transformBlock(transformer: FirTransformer<D>, data: D): FirCatch
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirTryExpression.kt

@@ -6,12 +6,12 @@
 package org.jetbrains.kotlin.fir.expressions
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.expressions.impl.FirUnknownTypeExpression
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
 abstract class FirTryExpression(
     psi: PsiElement?
-) : FirUnknownTypeExpression(psi) {
+) : FirCallLikeControlFlowExpression(psi) {
     abstract val tryBlock: FirBlock
 
     abstract val catches: List<FirCatch>
@@ -22,9 +22,14 @@ abstract class FirTryExpression(
         visitor.visitTryExpression(this, data)
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        calleeReference.accept(visitor, data)
         tryBlock.accept(visitor, data)
         catches.forEach { it.accept(visitor, data) }
         finallyBlock?.accept(visitor, data)
         super.acceptChildren(visitor, data)
     }
+
+    abstract fun <D> transformTryBlock(transformer: FirTransformer<D>, data: D): FirTryExpression
+    abstract fun <D> transformCatches(transformer: FirTransformer<D>, data: D): FirTryExpression
+    abstract fun <D> transformFinallyBlock(transformer: FirTransformer<D>, data: D): FirTryExpression
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirWhenBranch.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.expressions
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
 interface FirWhenBranch : FirElement {
@@ -24,4 +25,6 @@ interface FirWhenBranch : FirElement {
         condition.accept(visitor, data)
         result.accept(visitor, data)
     }
+
+    fun <D> transformCondition(transformer: FirTransformer<D>, data: D): FirWhenBranch
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirWhenExpression.kt

@@ -6,10 +6,10 @@
 package org.jetbrains.kotlin.fir.expressions
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.expressions.impl.FirUnknownTypeExpression
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
-abstract class FirWhenExpression(psi: PsiElement?) : FirUnknownTypeExpression(psi) {
+abstract class FirWhenExpression(psi: PsiElement?) : FirCallLikeControlFlowExpression(psi) {
     abstract val subject: FirExpression?
 
     // when (val subjectVariable = subject()) { ... }
@@ -21,10 +21,15 @@ abstract class FirWhenExpression(psi: PsiElement?) : FirUnknownTypeExpression(ps
         visitor.visitWhenExpression(this, data)
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        calleeReference.accept(visitor, data)
         subjectVariable?.accept(visitor, data) ?: subject?.accept(visitor, data)
         for (branch in branches) {
             branch.accept(visitor, data)
         }
         super.acceptChildren(visitor, data)
     }
+
+    abstract fun <D> transformBranches(transformer: FirTransformer<D>, data: D): FirWhenExpression
+
+    abstract fun <D> transformSubject(transformer: FirTransformer<D>, data: D): FirWhenExpression
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirCatchImpl.kt

@@ -20,9 +20,18 @@ class FirCatchImpl(
     override var block: FirBlock
 ) : FirAbstractElement(psi), FirCatch {
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
-        parameter = parameter.transformSingle(transformer, data)
-        block = block.transformSingle(transformer, data)
+        transformParameter(transformer, data)
+        transformBlock(transformer, data)
         return this
     }
 
+    override fun <D> transformParameter(transformer: FirTransformer<D>, data: D): FirCatch {
+        parameter = parameter.transformSingle(transformer, data)
+        return this
+    }
+
+    override fun <D> transformBlock(transformer: FirTransformer<D>, data: D): FirCatch {
+        block = block.transformSingle(transformer, data)
+        return this
+    }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirTryExpressionImpl.kt

@@ -7,9 +7,11 @@ package org.jetbrains.kotlin.fir.expressions.impl
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirReference
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirCatch
 import org.jetbrains.kotlin.fir.expressions.FirTryExpression
+import org.jetbrains.kotlin.fir.references.FirStubReference
 import org.jetbrains.kotlin.fir.transformInplace
 import org.jetbrains.kotlin.fir.transformSingle
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
@@ -17,14 +19,31 @@ import org.jetbrains.kotlin.fir.visitors.FirTransformer
 class FirTryExpressionImpl(
     psi: PsiElement?,
     override var tryBlock: FirBlock,
-    override var finallyBlock: FirBlock?
+    override var finallyBlock: FirBlock?,
+    override var calleeReference: FirReference = FirStubReference()
 ) : FirTryExpression(psi) {
     override val catches = mutableListOf<FirCatch>()
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
-        tryBlock = tryBlock.transformSingle(transformer, data)
-        finallyBlock = finallyBlock?.transformSingle(transformer, data)
-        catches.transformInplace(transformer, data)
+        calleeReference = calleeReference.transformSingle(transformer, data)
+        transformTryBlock(transformer, data)
+        transformCatches(transformer, data)
+        transformFinallyBlock(transformer, data)
         return super.transformChildren(transformer, data)
     }
+
+    override fun <D> transformTryBlock(transformer: FirTransformer<D>, data: D): FirTryExpression {
+        tryBlock = tryBlock.transformSingle(transformer, data)
+        return this
+    }
+
+    override fun <D> transformCatches(transformer: FirTransformer<D>, data: D): FirTryExpression {
+        catches.transformInplace(transformer, data)
+        return this
+    }
+
+    override fun <D> transformFinallyBlock(transformer: FirTransformer<D>, data: D): FirTryExpression {
+        finallyBlock = finallyBlock?.transformSingle(transformer, data)
+        return this
+    }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirWhenBranchImpl.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.expressions.impl
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhenBranch
@@ -21,8 +20,13 @@ class FirWhenBranchImpl(
     override var result: FirBlock
 ) : FirAbstractElement(psi), FirWhenBranch {
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
-        condition = condition.transformSingle(transformer, data)
+        transformCondition(transformer, data)
         result = result.transformSingle(transformer, data)
         return this
     }
+
+    override fun <D> transformCondition(transformer: FirTransformer<D>, data: D): FirWhenBranch {
+        condition = condition.transformSingle(transformer, data)
+        return this
+    }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirWhenExpressionImpl.kt

@@ -6,27 +6,44 @@
 package org.jetbrains.kotlin.fir.expressions.impl
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirReference
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirVariable
 import org.jetbrains.kotlin.fir.expressions.FirWhenBranch
 import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
+import org.jetbrains.kotlin.fir.references.FirStubReference
+import org.jetbrains.kotlin.fir.transformInplace
+import org.jetbrains.kotlin.fir.transformSingle
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
 
 class FirWhenExpressionImpl(
     psiElement: PsiElement?,
     override var subject: FirExpression? = null,
-    override var subjectVariable: FirVariable<*>? = null
+    override var subjectVariable: FirVariable<*>? = null,
+    override var calleeReference: FirReference = FirStubReference()
 ) : FirWhenExpression(psiElement) {
     override val branches = mutableListOf<FirWhenBranch>()
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
+        calleeReference = calleeReference.transformSingle(transformer, data)
+        transformSubject(transformer, data)
+        branches.transformInplace(transformer, data)
+        return super.transformChildren(transformer, data)
+    }
+
+    override fun <D> transformBranches(transformer: FirTransformer<D>, data: D): FirWhenExpression {
+        branches.transformInplace(transformer, data)
+        return this
+    }
+
+    override fun <D> transformSubject(transformer: FirTransformer<D>, data: D): FirWhenExpression {
         if (subjectVariable != null) {
             subjectVariable = subjectVariable?.transformSingle(transformer, data)
         } else {
             subject = subject?.transformSingle(transformer, data)
         }
-        branches.transformInplace(transformer, data)
-        return super.transformChildren(transformer, data)
+        return this
     }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/references/FirStubReference.kt

@@ -0,0 +1,14 @@
+/*
+ * Copyright 2010-2019 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.fir.references
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirPureAbstractElement
+import org.jetbrains.kotlin.fir.FirReference
+
+class FirStubReference : FirPureAbstractElement(), FirReference {
+    override val psi: PsiElement? get() = null
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/symbols/impl/FirFunctionSymbol.kt

@@ -18,7 +18,7 @@ sealed class FirFunctionSymbol<D : FirMemberFunction<D>>(
         get() = emptyList()
 }
 
-class FirNamedFunctionSymbol(
+open class FirNamedFunctionSymbol(
     callableId: CallableId,
     val isFakeOverride: Boolean = false,
     // Actual for fake override only

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirTransformerGenerated.kt

@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.fir.types.*
 abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirElement>, D>() {
     abstract fun <E : FirElement> transformElement(element: E, data: D): CompositeTransformResult<E>
 
-    open fun <E : FirElement> transformCatch(catch: E, data: D): CompositeTransformResult<E> {
+    open fun transformCatch(catch: FirCatch, data: D): CompositeTransformResult<FirCatch> {
         return transformElement(catch, data)
     }
 
@@ -284,6 +284,18 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformUnknownTypeExpression(block, data)
     }
 
+    open fun transformCallLikeControlFlowExpression(callLikeControlFlowExpression: FirCallLikeControlFlowExpression, data: D): CompositeTransformResult<FirStatement> {
+        return transformUnknownTypeExpression(callLikeControlFlowExpression, data)
+    }
+
+    open fun transformTryExpression(tryExpression: FirTryExpression, data: D): CompositeTransformResult<FirStatement> {
+        return transformCallLikeControlFlowExpression(tryExpression, data)
+    }
+
+    open fun transformWhenExpression(whenExpression: FirWhenExpression, data: D): CompositeTransformResult<FirStatement> {
+        return transformCallLikeControlFlowExpression(whenExpression, data)
+    }
+
     open fun transformClassReferenceExpression(classReferenceExpression: FirClassReferenceExpression, data: D): CompositeTransformResult<FirStatement> {
         return transformUnknownTypeExpression(classReferenceExpression, data)
     }
@@ -308,14 +320,6 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformUnknownTypeExpression(resolvedQualifier, data)
     }
 
-    open fun transformTryExpression(tryExpression: FirTryExpression, data: D): CompositeTransformResult<FirStatement> {
-        return transformUnknownTypeExpression(tryExpression, data)
-    }
-
-    open fun transformWhenExpression(whenExpression: FirWhenExpression, data: D): CompositeTransformResult<FirStatement> {
-        return transformUnknownTypeExpression(whenExpression, data)
-    }
-
     open fun transformWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression, data: D): CompositeTransformResult<FirStatement> {
         return transformUnknownTypeExpression(whenSubjectExpression, data)
     }
@@ -372,8 +376,12 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformLoop(whileLoop, data)
     }
 
+    open fun transformResolvable(resolvable: FirResolvable, data: D): CompositeTransformResult<FirStatement> {
+        return transformStatement(resolvable, data)
+    }
+
     open fun transformQualifiedAccess(qualifiedAccess: FirQualifiedAccess, data: D): CompositeTransformResult<FirStatement> {
-        return transformStatement(qualifiedAccess, data)
+        return transformResolvable(qualifiedAccess, data)
     }
 
     open fun transformAssignment(assignment: FirAssignment, data: D): CompositeTransformResult<FirStatement> {
@@ -500,6 +508,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformCall(call, data)
     }
 
+    final override fun visitCallLikeControlFlowExpression(callLikeControlFlowExpression: FirCallLikeControlFlowExpression, data: D): CompositeTransformResult<FirElement> {
+        return transformCallLikeControlFlowExpression(callLikeControlFlowExpression, data)
+    }
+
     final override fun visitCallWithArgumentList(callWithArgumentList: FirCallWithArgumentList, data: D): CompositeTransformResult<FirElement> {
         return transformCallWithArgumentList(callWithArgumentList, data)
     }
@@ -736,6 +748,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformRegularClass(regularClass, data)
     }
 
+    final override fun visitResolvable(resolvable: FirResolvable, data: D): CompositeTransformResult<FirElement> {
+        return transformResolvable(resolvable, data)
+    }
+
     final override fun visitResolvedCallableReference(resolvedCallableReference: FirResolvedCallableReference, data: D): CompositeTransformResult<FirElement> {
         return transformResolvedCallableReference(resolvedCallableReference, data)
     }

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorGenerated.kt

@@ -284,6 +284,18 @@ abstract class FirVisitor<out R, in D> {
         return visitUnknownTypeExpression(block, data)
     }
 
+    open fun visitCallLikeControlFlowExpression(callLikeControlFlowExpression: FirCallLikeControlFlowExpression, data: D): R {
+        return visitUnknownTypeExpression(callLikeControlFlowExpression, data)
+    }
+
+    open fun visitTryExpression(tryExpression: FirTryExpression, data: D): R {
+        return visitCallLikeControlFlowExpression(tryExpression, data)
+    }
+
+    open fun visitWhenExpression(whenExpression: FirWhenExpression, data: D): R {
+        return visitCallLikeControlFlowExpression(whenExpression, data)
+    }
+
     open fun visitClassReferenceExpression(classReferenceExpression: FirClassReferenceExpression, data: D): R {
         return visitUnknownTypeExpression(classReferenceExpression, data)
     }
@@ -308,14 +320,6 @@ abstract class FirVisitor<out R, in D> {
         return visitUnknownTypeExpression(resolvedQualifier, data)
     }
 
-    open fun visitTryExpression(tryExpression: FirTryExpression, data: D): R {
-        return visitUnknownTypeExpression(tryExpression, data)
-    }
-
-    open fun visitWhenExpression(whenExpression: FirWhenExpression, data: D): R {
-        return visitUnknownTypeExpression(whenExpression, data)
-    }
-
     open fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression, data: D): R {
         return visitUnknownTypeExpression(whenSubjectExpression, data)
     }
@@ -372,8 +376,12 @@ abstract class FirVisitor<out R, in D> {
         return visitLoop(whileLoop, data)
     }
 
+    open fun visitResolvable(resolvable: FirResolvable, data: D): R {
+        return visitStatement(resolvable, data)
+    }
+
     open fun visitQualifiedAccess(qualifiedAccess: FirQualifiedAccess, data: D): R {
-        return visitStatement(qualifiedAccess, data)
+        return visitResolvable(qualifiedAccess, data)
     }
 
     open fun visitAssignment(assignment: FirAssignment, data: D): R {

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorVoidGenerated.kt

@@ -284,6 +284,18 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitUnknownTypeExpression(block, null)
     }
 
+    open fun visitCallLikeControlFlowExpression(callLikeControlFlowExpression: FirCallLikeControlFlowExpression) {
+        visitUnknownTypeExpression(callLikeControlFlowExpression, null)
+    }
+
+    open fun visitTryExpression(tryExpression: FirTryExpression) {
+        visitCallLikeControlFlowExpression(tryExpression, null)
+    }
+
+    open fun visitWhenExpression(whenExpression: FirWhenExpression) {
+        visitCallLikeControlFlowExpression(whenExpression, null)
+    }
+
     open fun visitClassReferenceExpression(classReferenceExpression: FirClassReferenceExpression) {
         visitUnknownTypeExpression(classReferenceExpression, null)
     }
@@ -308,14 +320,6 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitUnknownTypeExpression(resolvedQualifier, null)
     }
 
-    open fun visitTryExpression(tryExpression: FirTryExpression) {
-        visitUnknownTypeExpression(tryExpression, null)
-    }
-
-    open fun visitWhenExpression(whenExpression: FirWhenExpression) {
-        visitUnknownTypeExpression(whenExpression, null)
-    }
-
     open fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression) {
         visitUnknownTypeExpression(whenSubjectExpression, null)
     }
@@ -372,8 +376,12 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitLoop(whileLoop, null)
     }
 
+    open fun visitResolvable(resolvable: FirResolvable) {
+        visitStatement(resolvable, null)
+    }
+
     open fun visitQualifiedAccess(qualifiedAccess: FirQualifiedAccess) {
-        visitStatement(qualifiedAccess, null)
+        visitResolvable(qualifiedAccess, null)
     }
 
     open fun visitAssignment(assignment: FirAssignment) {
@@ -500,6 +508,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitCall(call)
     }
 
+    final override fun visitCallLikeControlFlowExpression(callLikeControlFlowExpression: FirCallLikeControlFlowExpression, data: Nothing?) {
+        visitCallLikeControlFlowExpression(callLikeControlFlowExpression)
+    }
+
     final override fun visitCallWithArgumentList(callWithArgumentList: FirCallWithArgumentList, data: Nothing?) {
         visitCallWithArgumentList(callWithArgumentList)
     }
@@ -736,6 +748,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitRegularClass(regularClass)
     }
 
+    final override fun visitResolvable(resolvable: FirResolvable, data: Nothing?) {
+        visitResolvable(resolvable)
+    }
+
     final override fun visitResolvedCallableReference(resolvedCallableReference: FirResolvedCallableReference, data: Nothing?) {
         visitResolvedCallableReference(resolvedCallableReference)
     }

compiler/testData/ir/irText/expressions/kt27933.fir.txt

@@ -3,7 +3,7 @@ FILE fqName:<root> fileName:/kt27933.kt
     BLOCK_BODY
       VAR name:r type:kotlin.String [var]
         CONST String type=kotlin.String value=""
-      WHEN type=kotlin.String origin=IF
+      WHEN type=kotlin.Unit origin=IF
         BRANCH
           if: CALL 'public final fun not (): kotlin.Boolean declared in kotlin.Boolean' type=kotlin.Boolean origin=EXCLEQ
             $this: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EXCLEQ
@@ -15,7 +15,7 @@ FILE fqName:<root> fileName:/kt27933.kt
           then: BLOCK type=kotlin.Unit origin=null
             SET_VAR 'var r: kotlin.String [var] declared in <root>.box' type=kotlin.String origin=null
               CONST String type=kotlin.String value="O"
-      WHEN type=kotlin.String origin=IF
+      WHEN type=kotlin.Unit origin=IF
         BRANCH
           if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
             arg0: GET_VAR 'var r: kotlin.String [var] declared in <root>.box' type=kotlin.String origin=null

compiler/testData/ir/irText/expressions/smartCasts.fir.txt

@@ -36,7 +36,7 @@ FILE fqName:<root> fileName:/smartCasts.kt
   FUN name:test2 visibility:public modality:FINAL <> (x:kotlin.Any) returnType:kotlin.String
     VALUE_PARAMETER name:x index:0 type:kotlin.Any
     BLOCK_BODY
-      WHEN type=kotlin.String origin=IF
+      WHEN type=kotlin.Unit origin=IF
         BRANCH
           if: TYPE_OP type=kotlin.Boolean origin=NOT_INSTANCEOF typeOperand=kotlin.String
             GET_VAR 'x: kotlin.Any declared in <root>.test2' type=kotlin.Any origin=null
@@ -48,7 +48,7 @@ FILE fqName:<root> fileName:/smartCasts.kt
   FUN name:test3 visibility:public modality:FINAL <> (x:kotlin.Any) returnType:kotlin.String
     VALUE_PARAMETER name:x index:0 type:kotlin.Any
     BLOCK_BODY
-      WHEN type=kotlin.String origin=IF
+      WHEN type=kotlin.Unit origin=IF
         BRANCH
           if: TYPE_OP type=kotlin.Boolean origin=NOT_INSTANCEOF typeOperand=kotlin.String
             GET_VAR 'x: kotlin.Any declared in <root>.test3' type=kotlin.Any origin=null

compiler/testData/ir/irText/expressions/varargWithImplicitCast.fir.txt

@@ -2,7 +2,7 @@ FILE fqName:<root> fileName:/varargWithImplicitCast.kt
   FUN name:testScalar visibility:public modality:FINAL <> (a:kotlin.Any) returnType:kotlin.IntArray
     VALUE_PARAMETER name:a index:0 type:kotlin.Any
     BLOCK_BODY
-      WHEN type=kotlin.IntArray origin=IF
+      WHEN type=kotlin.Unit origin=IF
         BRANCH
           if: TYPE_OP type=kotlin.Boolean origin=NOT_INSTANCEOF typeOperand=kotlin.Int
             GET_VAR 'a: kotlin.Any declared in <root>.testScalar' type=kotlin.Any origin=null
@@ -14,7 +14,7 @@ FILE fqName:<root> fileName:/varargWithImplicitCast.kt
   FUN name:testSpread visibility:public modality:FINAL <> (a:kotlin.Any) returnType:kotlin.IntArray
     VALUE_PARAMETER name:a index:0 type:kotlin.Any
     BLOCK_BODY
-      WHEN type=kotlin.IntArray origin=IF
+      WHEN type=kotlin.Unit origin=IF
         BRANCH
           if: TYPE_OP type=kotlin.Boolean origin=NOT_INSTANCEOF typeOperand=kotlin.IntArray
             GET_VAR 'a: kotlin.Any declared in <root>.testSpread' type=kotlin.Any origin=null

compiler/testData/ir/irText/expressions/whenWithSubjectVariable.fir.txt

@@ -13,8 +13,8 @@ FILE fqName:<root> fileName:/whenWithSubjectVariable.kt
             BRANCH
               if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
                 arg0: GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
-                arg1: CONST Int type=IrErrorType value=42
-              then: CONST Int type=IrErrorType value=1
+                arg1: CONST Int type=kotlin.Int value=42
+              then: CONST Int type=kotlin.Int value=1
             BRANCH
               if: TYPE_OP type=kotlin.Boolean origin=INSTANCEOF typeOperand=kotlin.String
                 GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
@@ -22,21 +22,15 @@ FILE fqName:<root> fileName:/whenWithSubjectVariable.kt
             BRANCH
               if: TYPE_OP type=kotlin.Boolean origin=NOT_INSTANCEOF typeOperand=kotlin.Int
                 GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
-              then: CONST Int type=IrErrorType value=2
+              then: CONST Int type=kotlin.Int value=2
             BRANCH
-              if: CALL 'public open fun contains (value: kotlin.Int): kotlin.Boolean declared in kotlin.ranges.IntRange' type=kotlin.Boolean origin=null
-                $this: CALL 'public final fun rangeTo (other: kotlin.Int): kotlin.ranges.IntRange declared in kotlin.Int' type=kotlin.ranges.IntRange origin=null
-                  $this: CONST Int type=kotlin.Int value=0
-                  other: CONST Int type=kotlin.Int value=10
-                value: GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
-              then: CONST Int type=IrErrorType value=3
+              if: ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [kotlin/ranges/IntRange.contains]>#' type=IrErrorType
+                GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=3
             BRANCH
-              if: CALL 'public open fun contains (value: kotlin.Int): kotlin.Boolean declared in kotlin.ranges.IntRange' type=kotlin.Boolean origin=null
-                $this: CALL 'public final fun rangeTo (other: kotlin.Int): kotlin.ranges.IntRange declared in kotlin.Int' type=kotlin.ranges.IntRange origin=null
-                  $this: CONST Int type=kotlin.Int value=10
-                  other: CONST Int type=kotlin.Int value=20
-                value: GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
-              then: CONST Int type=IrErrorType value=4
+              if: ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [kotlin/ranges/IntRange.contains]>#' type=IrErrorType
+                GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=4
             BRANCH
               if: CONST Boolean type=kotlin.Boolean value=true
               then: CALL 'public final fun unaryMinus (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null

compiler/testData/ir/irText/regressions/coercionInLoop.fir.txt

@@ -13,7 +13,7 @@ FILE fqName:<root> fileName:/coercionInLoop.kt
         condition: CALL 'public abstract fun hasNext (): kotlin.Boolean declared in kotlin.collections.Iterator' type=kotlin.Boolean origin=null
           $this: GET_VAR 'val x: kotlin.collections.DoubleIterator [val] declared in <root>.box' type=kotlin.collections.DoubleIterator origin=null
         body: BLOCK type=kotlin.Int origin=null
-          WHEN type=kotlin.String origin=IF
+          WHEN type=kotlin.Unit origin=IF
             BRANCH
               if: CALL 'public final fun not (): kotlin.Boolean declared in kotlin.Boolean' type=kotlin.Boolean origin=EXCLEQ
                 $this: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EXCLEQ