FIR body resolve: support subject type calculation

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
-import org.jetbrains.kotlin.fir.expressions.impl.FirWhenSubjectExpression
 import org.jetbrains.kotlin.fir.references.FirPropertyFromParameterCallableReference
 import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.buildUseSiteScope
@@ -798,12 +797,6 @@ internal class Fir2IrVisitor(
     private fun FirExpression.toIrExpression(): IrExpression {
         return when (this) {
             is FirBlock -> convertToIrExpressionOrBlock()
-            is FirWhenSubjectExpression -> {
-                val lastSubjectVariable = subjectVariableStack.last()
-                convertWithOffsets { startOffset, endOffset ->
-                    IrGetValueImpl(startOffset, endOffset, lastSubjectVariable.type, lastSubjectVariable.symbol)
-                }
-            }
             is FirUnitExpression -> convertWithOffsets { startOffset, endOffset ->
                 IrGetObjectValueImpl(
                     startOffset, endOffset, unitType,
@@ -917,6 +910,13 @@ internal class Fir2IrVisitor(
         }
     }
 
+    override fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression, data: Any?): IrElement {
+        val lastSubjectVariable = subjectVariableStack.last()
+        return whenSubjectExpression.convertWithOffsets { startOffset, endOffset ->
+            IrGetValueImpl(startOffset, endOffset, lastSubjectVariable.type, lastSubjectVariable.symbol)
+        }
+    }
+
     private val loopMap = mutableMapOf<FirLoop, IrLoop>()
 
     override fun visitDoWhileLoop(doWhileLoop: FirDoWhileLoop, data: Any?): IrElement {

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -10,6 +10,7 @@ import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.fir.FirReference
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirWhenSubject
 import org.jetbrains.kotlin.fir.declarations.impl.FirVariableImpl
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.*
@@ -179,10 +180,12 @@ internal fun FirExpression.generateNotNullOrOther(
 ): FirWhenExpression {
     val subjectName = Name.special("<$caseId>")
     val subjectVariable = generateTemporaryVariable(session, psi, subjectName, this)
-    val subjectExpression = FirWhenSubjectExpression(session, psi)
+    val subject = FirWhenSubject()
+    val subjectExpression = FirWhenSubjectExpressionImpl(session, psi, subject)
     return FirWhenExpressionImpl(
         session, basePsi, this, subjectVariable
     ).apply {
+        subject.bind(this)
         branches += FirWhenBranchImpl(
             session, psi,
             FirOperatorCallImpl(session, psi, FirOperation.EQ).apply {
@@ -219,10 +222,11 @@ internal fun FirExpression.generateLazyLogicalOperation(
 
 internal fun KtWhenCondition.toFirWhenCondition(
     session: FirSession,
+    subject: FirWhenSubject,
     convert: KtExpression?.(String) -> FirExpression,
     toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef
 ): FirExpression {
-    val firSubjectExpression = FirWhenSubjectExpression(session, this)
+    val firSubjectExpression = FirWhenSubjectExpressionImpl(session, this, subject)
     return when (this) {
         is KtWhenConditionWithExpression -> {
             FirOperatorCallImpl(
@@ -255,12 +259,13 @@ internal fun KtWhenCondition.toFirWhenCondition(
 internal fun Array<KtWhenCondition>.toFirWhenCondition(
     session: FirSession,
     basePsi: KtElement,
+    subject: FirWhenSubject,
     convert: KtExpression?.(String) -> FirExpression,
     toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef
 ): FirExpression {
     var firCondition: FirExpression? = null
     for (condition in this) {
-        val firConditionElement = condition.toFirWhenCondition(session, convert, toFirOrErrorTypeRef)
+        val firConditionElement = condition.toFirWhenCondition(session, subject, convert, toFirOrErrorTypeRef)
         firCondition = when (firCondition) {
             null -> firConditionElement
             else -> firCondition.generateLazyLogicalOperation(

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -972,26 +972,30 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
         }
 
         override fun visitWhenExpression(expression: KtWhenExpression, data: Unit): FirElement {
-            val subjectExpression = expression.subjectExpression
+            val ktSubjectExpression = expression.subjectExpression
-            val subject = when (subjectExpression) {
+            val subjectExpression = when (ktSubjectExpression) {
-                is KtVariableDeclaration -> subjectExpression.initializer
+                is KtVariableDeclaration -> ktSubjectExpression.initializer
-                else -> subjectExpression
+                else -> ktSubjectExpression
-            }?.toFirExpression("Incorrect when subject expression: ${subjectExpression?.text}")
+            }?.toFirExpression("Incorrect when subject expression: ${ktSubjectExpression?.text}")
-            val subjectVariable = when (subjectExpression) {
+            val subjectVariable = when (ktSubjectExpression) {
                 is KtVariableDeclaration -> FirVariableImpl(
-                    session, subjectExpression, subjectExpression.nameAsSafeName,
+                    session, ktSubjectExpression, ktSubjectExpression.nameAsSafeName,
-                    subjectExpression.typeReference.toFirOrImplicitType(),
+                    ktSubjectExpression.typeReference.toFirOrImplicitType(),
-                    isVar = false, initializer = subject
+                    isVar = false, initializer = subjectExpression
                 )
                 else -> null
             }
-            val hasSubject = subject != null
+            val hasSubject = subjectExpression != null
+            val subject = FirWhenSubject()
             return FirWhenExpressionImpl(
                 session,
                 expression,
-                subject,
+                subjectExpression,
                 subjectVariable
             ).apply {
+                if (hasSubject) {
+                    subject.bind(this)
+                }
                 for (entry in expression.entries) {
                     val branch = entry.expression.toFirBlock()
                     branches += if (!entry.isElse) {
@@ -999,6 +1003,7 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
                             val firCondition = entry.conditions.toFirWhenCondition(
                                 this@RawFirBuilder.session,
                                 entry,
+                                subject,
                                 { toFirExpression(it) },
                                 { toFirOrErrorType() }
                             )

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

@@ -425,6 +425,18 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
         return whenExpression.compose()
     }
 
+    override fun transformWhenSubjectExpression(
+        whenSubjectExpression: FirWhenSubjectExpression,
+        data: Any?
+    ): CompositeTransformResult<FirStatement> {
+        val parentWhen = whenSubjectExpression.whenSubject.whenExpression
+        val subjectType = parentWhen.subject?.resultType ?: parentWhen.subjectVariable?.returnTypeRef
+        if (subjectType != null) {
+            whenSubjectExpression.resultType = subjectType
+        }
+        return whenSubjectExpression.compose()
+    }
+
     override fun <T> transformConstExpression(constExpression: FirConstExpression<T>, data: Any?): CompositeTransformResult<FirStatement> {
         val expectedType = data as FirTypeRef?
 

compiler/fir/resolve/testData/loadCompiledKotlin/coroutines/Basic.txt

@@ -1,4 +1,4 @@
-public final fun builder(c: R|error: createSuspendFunctionType not supported|): R|kotlin/Unit|
+public final fun builder(c: R|class error: createSuspendFunctionType not supported|): R|kotlin/Unit|
 
 public final class Controller : R|kotlin/Any| {
     public final suspend fun suspendFun(): R|kotlin/Unit|

compiler/fir/resolve/testData/loadCompiledKotlin/type/SuspendFunction.txt

@@ -1,7 +1,8 @@
-public final fun test1(): R|error: createSuspendFunctionType not supported|
+public final fun test1(): R|class error: createSuspendFunctionType not supported|
 
-public final fun test2(): R|error: createSuspendFunctionType not supported|
+public final fun test2(): R|class error: createSuspendFunctionType not supported|
 
-public final fun test3(): R|kotlin/collections/List<error: createSuspendFunctionType not supported>|
+public final fun test3(): R|kotlin/collections/List<class error: createSuspendFunctionType not supported>|
+
+public final fun test4(): R|class error: createSuspendFunctionType not supported|
 
-public final fun test4(): R|error: createSuspendFunctionType not supported|

compiler/fir/resolve/testData/resolve/expresssions/when.kt

@@ -1 +1,6 @@
-fun foo() = if (true) 1 else 0
+fun foo() = if (true) 1 else 0
+
+fun bar(arg: Any?) = when (arg) {
+    is Int -> arg as Int
+    else -> 42
+}

compiler/fir/resolve/testData/resolve/expresssions/when.txt

@@ -10,3 +10,14 @@ FILE: when.kt
         }
 
     }
+    public final fun bar(arg: R|kotlin/Any|?): R|kotlin/Int| {
+        ^bar when (R|<local>/arg|) {
+            ($subj$ is R|kotlin/Int|) ->  {
+                (R|<local>/arg| as R|kotlin/Int|)
+            }
+            else ->  {
+                Int(42)
+            }
+        }
+
+    }

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

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirExpressionStub
 import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
-import org.jetbrains.kotlin.fir.expressions.impl.FirWhenSubjectExpression
 import org.jetbrains.kotlin.fir.symbols.ConeCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
@@ -38,7 +37,6 @@ fun ConeKotlinType.render(): String {
     return when (this) {
         is ConeTypeVariableType -> "TypeVariable(${this.lookupTag.name})"
         is ConeDefinitelyNotNullType -> "${original.render()}!"
-        is ConeKotlinErrorType -> "error: $reason"
         is ConeClassErrorType -> "class error: $reason"
         is ConeCapturedType -> "captured type: lowerType = ${lowerType?.render()}"
         is ConeClassLikeType -> {
@@ -540,6 +538,10 @@ class FirRenderer(builder: StringBuilder) : FirVisitorVoid() {
         println("}")
     }
 
+    override fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression) {
+        print("\$subj\$")
+    }
+
     override fun visitTryExpression(tryExpression: FirTryExpression) {
         tryExpression.annotations.renderAnnotations()
         print("try")
@@ -604,7 +606,6 @@ class FirRenderer(builder: StringBuilder) : FirVisitorVoid() {
             when (expression) {
                 is FirExpressionStub -> "STUB"
                 is FirUnitExpression -> "Unit"
-                is FirWhenSubjectExpression -> "\$subj\$"
                 is FirElseIfTrueCondition -> "else"
                 else -> "??? ${expression.javaClass}"
             }

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

@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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
+
+import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
+
+class FirWhenSubject {
+    lateinit var whenExpression: FirWhenExpression
+
+    fun bind(expression: FirWhenExpression) {
+        whenExpression = expression
+    }
+}

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

@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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 org.jetbrains.kotlin.fir.FirWhenSubject
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+
+interface FirWhenSubjectExpression : FirExpression {
+    val whenSubject: FirWhenSubject
+
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
+        visitor.visitWhenSubjectExpression(this, data)
+}

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

@@ -6,11 +6,12 @@
 package org.jetbrains.kotlin.fir.expressions.impl
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.FirWhenSubject
+import org.jetbrains.kotlin.fir.expressions.FirWhenSubjectExpression
 
-class FirWhenSubjectExpression(
+class FirWhenSubjectExpressionImpl(
     session: FirSession,
-    psi: PsiElement?
+    psi: PsiElement?,
-): FirAbstractExpression(session, psi)
+    override val whenSubject: FirWhenSubject
+) : FirAbstractExpression(session, psi), FirWhenSubjectExpression

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

@@ -276,6 +276,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformExpression(whenExpression, data)
     }
 
+    open fun transformWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression, data: D): CompositeTransformResult<FirStatement> {
+        return transformExpression(whenSubjectExpression, data)
+    }
+
     open fun transformWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression, data: D): CompositeTransformResult<FirStatement> {
         return transformExpression(wrappedArgumentExpression, data)
     }
@@ -752,6 +756,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformWhenExpression(whenExpression, data)
     }
 
+    final override fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression, data: D): CompositeTransformResult<FirElement> {
+        return transformWhenSubjectExpression(whenSubjectExpression, data)
+    }
+
     final override fun visitWhileLoop(whileLoop: FirWhileLoop, data: D): CompositeTransformResult<FirElement> {
         return transformWhileLoop(whileLoop, data)
     }

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

@@ -276,6 +276,10 @@ abstract class FirVisitor<out R, in D> {
         return visitExpression(whenExpression, data)
     }
 
+    open fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression, data: D): R {
+        return visitExpression(whenSubjectExpression, data)
+    }
+
     open fun visitWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression, data: D): R {
         return visitExpression(wrappedArgumentExpression, data)
     }

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

@@ -276,6 +276,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitExpression(whenExpression, null)
     }
 
+    open fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression) {
+        visitExpression(whenSubjectExpression, null)
+    }
+
     open fun visitWrappedArgumentExpression(wrappedArgumentExpression: FirWrappedArgumentExpression) {
         visitExpression(wrappedArgumentExpression, null)
     }
@@ -752,6 +756,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitWhenExpression(whenExpression)
     }
 
+    final override fun visitWhenSubjectExpression(whenSubjectExpression: FirWhenSubjectExpression, data: Nothing?) {
+        visitWhenSubjectExpression(whenSubjectExpression)
+    }
+
     final override fun visitWhileLoop(whileLoop: FirWhileLoop, data: Nothing?) {
         visitWhileLoop(whileLoop)
     }