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

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)
     }