[FIR] Add abstract data flow analyzer with callbacks for body resolve transformer

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/DummyFirDataFlowAnalyzer.kt

@@ -0,0 +1,112 @@
+/*
+ * 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.dfa
+
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
+
+class DummyFirDataFlowAnalyzer : FirDataFlowAnalyzer() {
+    companion object {
+        private const val DUMMY = "<DUMMY>"
+    }
+
+    override fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): ConeKotlinType? {
+        return qualifiedAccessExpression.typeRef.coneTypeUnsafe()
+    }
+
+    override fun enterFunction(function: FirFunction<*>) {}
+
+    override fun exitFunction(function: FirFunction<*>): ControlFlowGraph = ControlFlowGraph(DUMMY)
+
+    override fun enterBlock(block: FirBlock) {}
+
+    override fun exitBlock(block: FirBlock) {}
+
+    override fun enterProperty(property: FirProperty) {}
+
+    override fun exitProperty(property: FirProperty) = ControlFlowGraph(DUMMY)
+
+    override fun exitTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall) {}
+
+    override fun exitOperatorCall(operatorCall: FirOperatorCall) {}
+
+    override fun exitJump(jump: FirJump<*>) {}
+
+    override fun enterWhenExpression(whenExpression: FirWhenExpression) {}
+
+    override fun enterWhenBranchCondition(whenBranch: FirWhenBranch) {}
+
+    override fun exitWhenBranchCondition(whenBranch: FirWhenBranch) {}
+
+    override fun exitWhenBranchResult(whenBranch: FirWhenBranch) {}
+
+    override fun exitWhenExpression(whenExpression: FirWhenExpression) {}
+
+    override fun enterWhileLoop(loop: FirLoop) {}
+
+    override fun exitWhileLoopCondition(loop: FirLoop) {}
+
+    override fun exitWhileLoop(loop: FirLoop) {}
+
+    override fun enterDoWhileLoop(loop: FirLoop) {}
+
+    override fun enterDoWhileLoopCondition(loop: FirLoop) {}
+
+    override fun exitDoWhileLoop(loop: FirLoop) {}
+
+    override fun enterTryExpression(tryExpression: FirTryExpression) {}
+
+    override fun exitTryMainBlock(tryExpression: FirTryExpression) {}
+
+    override fun enterCatchClause(catch: FirCatch) {}
+
+    override fun exitCatchClause(catch: FirCatch) {}
+
+    override fun enterFinallyBlock(tryExpression: FirTryExpression) {}
+
+    override fun exitFinallyBlock(tryExpression: FirTryExpression) {}
+
+    override fun exitTryExpression(tryExpression: FirTryExpression) {}
+
+    override fun exitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression) {}
+
+    override fun enterFunctionCall(functionCall: FirFunctionCall) {}
+
+    override fun exitFunctionCall(functionCall: FirFunctionCall) {}
+
+    override fun exitConstExpresion(constExpression: FirConstExpression<*>) {}
+
+    override fun exitVariableDeclaration(variable: FirVariable<*>) {}
+
+    override fun exitVariableAssignment(assignment: FirVariableAssignment) {}
+
+    override fun exitThrowExceptionNode(throwExpression: FirThrowExpression) {}
+
+    override fun enterBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression) {}
+
+    override fun exitLeftBinaryAndArgument(binaryLogicExpression: FirBinaryLogicExpression) {}
+
+    override fun exitBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression) {}
+
+    override fun enterBinaryOr(binaryLogicExpression: FirBinaryLogicExpression) {}
+
+    override fun exitLeftBinaryOrArgument(binaryLogicExpression: FirBinaryLogicExpression) {}
+
+    override fun exitBinaryOr(binaryLogicExpression: FirBinaryLogicExpression) {}
+
+    override fun enterAnnotationCall(annotationCall: FirAnnotationCall) {}
+
+    override fun exitAnnotationCall(annotationCall: FirAnnotationCall) {}
+
+    override fun enterInitBlock(initBlock: FirAnonymousInitializer) {}
+
+    override fun exitInitBlock(initBlock: FirAnonymousInitializer) {}
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -0,0 +1,101 @@
+/*
+ * 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.dfa
+
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+
+abstract class FirDataFlowAnalyzer {
+    abstract fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): ConeKotlinType?
+
+    // ----------------------------------- Named function -----------------------------------
+
+    abstract fun enterFunction(function: FirFunction<*>)
+    abstract fun exitFunction(function: FirFunction<*>): ControlFlowGraph
+
+    // ----------------------------------- Property -----------------------------------
+
+    abstract fun enterProperty(property: FirProperty)
+    abstract fun exitProperty(property: FirProperty): ControlFlowGraph
+
+    // ----------------------------------- Block -----------------------------------
+
+    abstract fun enterBlock(block: FirBlock)
+    abstract fun exitBlock(block: FirBlock)
+
+    // ----------------------------------- Operator call -----------------------------------
+
+    abstract fun exitTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall)
+    abstract fun exitOperatorCall(operatorCall: FirOperatorCall)
+
+    // ----------------------------------- Jump -----------------------------------
+
+    abstract fun exitJump(jump: FirJump<*>)
+
+    // ----------------------------------- When -----------------------------------
+
+    abstract fun enterWhenExpression(whenExpression: FirWhenExpression)
+    abstract fun enterWhenBranchCondition(whenBranch: FirWhenBranch)
+    abstract fun exitWhenBranchCondition(whenBranch: FirWhenBranch)
+    abstract fun exitWhenBranchResult(whenBranch: FirWhenBranch)
+    abstract fun exitWhenExpression(whenExpression: FirWhenExpression)
+
+    // ----------------------------------- While Loop -----------------------------------
+
+    abstract fun enterWhileLoop(loop: FirLoop)
+    abstract fun exitWhileLoopCondition(loop: FirLoop)
+    abstract fun exitWhileLoop(loop: FirLoop)
+
+    // ----------------------------------- Do while Loop -----------------------------------
+
+    abstract fun enterDoWhileLoop(loop: FirLoop)
+    abstract fun enterDoWhileLoopCondition(loop: FirLoop)
+    abstract fun exitDoWhileLoop(loop: FirLoop)
+
+    // ----------------------------------- Try-catch-finally -----------------------------------
+
+    abstract fun enterTryExpression(tryExpression: FirTryExpression)
+    abstract fun exitTryMainBlock(tryExpression: FirTryExpression)
+    abstract fun enterCatchClause(catch: FirCatch)
+    abstract fun exitCatchClause(catch: FirCatch)
+    abstract fun enterFinallyBlock(tryExpression: FirTryExpression)
+    abstract fun exitFinallyBlock(tryExpression: FirTryExpression)
+    abstract fun exitTryExpression(tryExpression: FirTryExpression)
+
+    // ----------------------------------- Resolvable call -----------------------------------
+
+    abstract fun exitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression)
+    abstract fun enterFunctionCall(functionCall: FirFunctionCall)
+    abstract fun exitFunctionCall(functionCall: FirFunctionCall)
+    abstract fun exitConstExpresion(constExpression: FirConstExpression<*>)
+    abstract fun exitVariableDeclaration(variable: FirVariable<*>)
+    abstract fun exitVariableAssignment(assignment: FirVariableAssignment)
+    abstract fun exitThrowExceptionNode(throwExpression: FirThrowExpression)
+
+    // ----------------------------------- Boolean operators -----------------------------------
+
+    abstract fun enterBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression)
+    abstract fun exitLeftBinaryAndArgument(binaryLogicExpression: FirBinaryLogicExpression)
+    abstract fun exitBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression)
+    abstract fun enterBinaryOr(binaryLogicExpression: FirBinaryLogicExpression)
+    abstract fun exitLeftBinaryOrArgument(binaryLogicExpression: FirBinaryLogicExpression)
+    abstract fun exitBinaryOr(binaryLogicExpression: FirBinaryLogicExpression)
+
+    // ----------------------------------- Annotations -----------------------------------
+
+    abstract fun enterAnnotationCall(annotationCall: FirAnnotationCall)
+    abstract fun exitAnnotationCall(annotationCall: FirAnnotationCall)
+
+    // ----------------------------------- Init block -----------------------------------
+
+    abstract fun enterInitBlock(initBlock: FirAnonymousInitializer)
+    abstract fun exitInitBlock(initBlock: FirAnonymousInitializer)
+}
+

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

@@ -14,6 +14,8 @@ import org.jetbrains.kotlin.fir.references.FirExplicitThisReference
 import org.jetbrains.kotlin.fir.references.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.*
+import org.jetbrains.kotlin.fir.resolve.dfa.DummyFirDataFlowAnalyzer
+import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.addImportingScopes
@@ -30,7 +32,6 @@ import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.SimpleTypeMarker
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
@@ -82,6 +83,7 @@ open class FirBodyResolveTransformer(
     )
 
     private val syntheticCallGenerator: FirSyntheticCallGenerator = FirSyntheticCallGenerator(this)
+    private val dataFlowAnalyzer: FirDataFlowAnalyzer = DummyFirDataFlowAnalyzer()
 
     override val <D> AbstractFirBasedSymbol<D>.phasedFir: D where D : FirDeclaration, D : FirSymbolOwner<D>
         get() {
@@ -115,8 +117,11 @@ open class FirBodyResolveTransformer(
     ): CompositeTransformResult<FirDeclaration> {
         if (implicitTypeOnly) return anonymousInitializer.compose()
         return withScopeCleanup(localScopes) {
+            dataFlowAnalyzer.enterInitBlock(anonymousInitializer)
             localScopes.addIfNotNull(primaryConstructorParametersScope)
-            super.transformAnonymousInitializer(anonymousInitializer, data)
+            super.transformAnonymousInitializer(anonymousInitializer, data).also {
+                dataFlowAnalyzer.exitInitBlock(it.single as FirAnonymousInitializer)
+            }
         }
     }
 
@@ -166,7 +171,7 @@ open class FirBodyResolveTransformer(
 
     override fun transformTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall, data: Any?): CompositeTransformResult<FirStatement> {
         val symbolProvider = session.service<FirSymbolProvider>()
-        val resolved = super.transformTypeOperatorCall(typeOperatorCall, data).single
+        val resolved = transformExpression(typeOperatorCall, data).single
         when ((resolved as FirTypeOperatorCall).operation) {
             FirOperation.IS, FirOperation.NOT_IS -> {
                 resolved.resultType = FirResolvedTypeRefImpl(
@@ -186,6 +191,7 @@ open class FirBodyResolveTransformer(
             }
             else -> error("Unknown type operator")
         }
+        dataFlowAnalyzer.exitTypeOperatorCall(typeOperatorCall)
         return resolved.compose()
     }
 
@@ -194,7 +200,7 @@ open class FirBodyResolveTransformer(
         data: Any?
     ): CompositeTransformResult<FirStatement> {
 
-        when (val callee = qualifiedAccessExpression.calleeReference) {
+        val result = when (val callee = qualifiedAccessExpression.calleeReference) {
             is FirExplicitThisReference -> {
                 val labelName = callee.labelName
                 val implicitReceiver = implicitReceiverStack[labelName]
@@ -203,7 +209,7 @@ open class FirBodyResolveTransformer(
                     null, implicitReceiver?.type ?: ConeKotlinErrorType("Unresolved this@$labelName"),
                     emptyList()
                 )
-                return qualifiedAccessExpression.compose()
+                qualifiedAccessExpression
             }
             is FirSuperReference -> {
                 if (callee.superTypeRef is FirResolvedTypeRef) {
@@ -215,29 +221,40 @@ open class FirBodyResolveTransformer(
                     qualifiedAccessExpression.resultType = superTypeRef
                     callee.replaceSuperTypeRef(superTypeRef)
                 }
-                return qualifiedAccessExpression.compose()
+                qualifiedAccessExpression
             }
             is FirDelegateFieldReference -> {
                 val delegateFieldSymbol = callee.coneSymbol
                 qualifiedAccessExpression.resultType = delegateFieldSymbol.delegate.typeRef
-                return qualifiedAccessExpression.compose()
+                qualifiedAccessExpression
             }
             is FirResolvedCallableReference -> {
                 if (qualifiedAccessExpression.typeRef !is FirResolvedTypeRef) {
                     storeTypeFromCallee(qualifiedAccessExpression)
                 }
-                return qualifiedAccessExpression.compose()
+                qualifiedAccessExpression
+            }
+            else -> {
+                val transformedCallee = callResolver.resolveVariableAccessAndSelectCandidate(qualifiedAccessExpression, file)
+                // NB: here we can get raw expression because of dropped qualifiers (see transform callee),
+                // so candidate existence must be checked before calling completion
+                val result = if (transformedCallee is FirQualifiedAccessExpression && transformedCallee.candidate() != null) {
+                    callCompleter.completeCall(transformedCallee, data as? FirTypeRef)
+                } else {
+                    transformedCallee
+                }
+                if (result is FirQualifiedAccessExpression) {
+                    dataFlowAnalyzer.getTypeUsingSmartcastInfo(result)?.let { typeWithSmartCasts ->
+                        result.resultType = FirResolvedTypeRefImpl(result.psi, typeWithSmartCasts, result.resultType.annotations)
+                    }
+                }
+                result
             }
         }
-
-        val transformedCallee = callResolver.resolveVariableAccessAndSelectCandidate(qualifiedAccessExpression, file)
-        // NB: here we can get raw expression because of dropped qualifiers (see transform callee),
-        // so candidate existence must be checked before calling completion
-        return if (transformedCallee is FirQualifiedAccessExpression && transformedCallee.candidate() != null) {
-            callCompleter.completeCall(transformedCallee, data as? FirTypeRef).compose()
-        } else {
-            transformedCallee.compose()
+        if (result is FirQualifiedAccessExpression) {
+            dataFlowAnalyzer.exitQualifiedAccessExpression(result)
         }
+        return result.compose()
     }
 
     override fun transformVariableAssignment(
@@ -246,14 +263,17 @@ open class FirBodyResolveTransformer(
     ): CompositeTransformResult<FirStatement> {
         // val resolvedAssignment = transformCallee(variableAssignment)
         val resolvedAssignment = callResolver.resolveVariableAccessAndSelectCandidate(variableAssignment, file)
-        return if (resolvedAssignment is FirVariableAssignment) {
+        val result = if (resolvedAssignment is FirVariableAssignment) {
             val completeAssignment = callCompleter.completeCall(resolvedAssignment, noExpectedType)
             val expectedType = typeFromCallee(completeAssignment)
-            completeAssignment.transformRValue(this, expectedType).compose()
+            completeAssignment.transformRValue(this, expectedType)
         } else {
             // This can happen in erroneous code only
-            resolvedAssignment.compose()
+            resolvedAssignment
         }
+        // TODO: maybe replace with FirAbstractAssignment for performance?
+        (result as? FirVariableAssignment)?.let { dataFlowAnalyzer.exitVariableAssignment(it) }
+        return result.compose()
     }
 
     override fun transformAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Any?): CompositeTransformResult<FirDeclaration> {
@@ -379,32 +399,46 @@ open class FirBodyResolveTransformer(
 
 
     override fun transformCatch(catch: FirCatch, data: Any?): CompositeTransformResult<FirCatch> {
+        dataFlowAnalyzer.enterCatchClause(catch)
         return withScopeCleanup(localScopes) {
             localScopes += FirLocalScope()
             catch.transformParameter(this, noExpectedType)
-            catch.transformBlock(this, null).compose()
-        }
+            catch.transformBlock(this, null)
+        }.also { dataFlowAnalyzer.exitCatchClause(it) }.compose()
     }
 
     override fun transformTryExpression(tryExpression: FirTryExpression, data: Any?): CompositeTransformResult<FirStatement> {
         if (tryExpression.calleeReference is FirResolvedCallableReference && tryExpression.resultType !is FirImplicitTypeRef) {
             return tryExpression.compose()
         }
+
+        dataFlowAnalyzer.enterTryExpression(tryExpression)
         tryExpression.transformTryBlock(this, null)
+        dataFlowAnalyzer.exitTryMainBlock(tryExpression)
         tryExpression.transformCatches(this, null)
 
         @Suppress("NAME_SHADOWING")
-        val tryExpression = syntheticCallGenerator.generateCalleeForTryExpression(tryExpression) ?: run {
-            tryExpression.resultType = FirErrorTypeRefImpl(null, "Inapplicable when expression")
-            return tryExpression.compose()
+        var result = syntheticCallGenerator.generateCalleeForTryExpression(tryExpression)?.let {
+            val expectedTypeRef = data as FirTypeRef?
+            callCompleter.completeCall(it, expectedTypeRef)
+        } ?: run {
+            tryExpression.resultType = FirErrorTypeRefImpl(null, "")
+            tryExpression
         }
-        val expectedTypeRef = data as FirTypeRef?
-        val result = callCompleter.completeCall(tryExpression, expectedTypeRef)
 
-        return result.transformFinallyBlock(this, noExpectedType).compose()
+        result = if (result.finallyBlock != null) {
+            result.also(dataFlowAnalyzer::enterFinallyBlock)
+                .transformFinallyBlock(this, noExpectedType)
+                .also(dataFlowAnalyzer::exitFinallyBlock)
+        } else {
+            result
+        }
+        dataFlowAnalyzer.exitTryExpression(result)
+        return result.compose()
     }
 
     override fun transformFunctionCall(functionCall: FirFunctionCall, data: Any?): CompositeTransformResult<FirStatement> {
+        dataFlowAnalyzer.enterFunctionCall(functionCall)
         if (functionCall.calleeReference is FirResolvedCallableReference && functionCall.resultType is FirImplicitTypeRef) {
             storeTypeFromCallee(functionCall)
         }
@@ -424,7 +458,7 @@ open class FirBodyResolveTransformer(
                 throw RuntimeException("While resolving call ${functionCall.render()}", e)
             }
 
-
+        dataFlowAnalyzer.exitFunctionCall(completeInference)
         return completeInference.compose()
 
     }
@@ -437,6 +471,7 @@ open class FirBodyResolveTransformer(
 
 
     override fun transformBlock(block: FirBlock, data: Any?): CompositeTransformResult<FirStatement> {
+        dataFlowAnalyzer.enterBlock(block)
         @Suppress("NAME_SHADOWING")
         val block = block.transformChildren(this, data) as FirBlock
         val statement = block.statements.lastOrNull()
@@ -451,18 +486,20 @@ open class FirBodyResolveTransformer(
         } else {
             (resultExpression.resultType as? FirResolvedTypeRef) ?: FirErrorTypeRefImpl(null, "No type for block")
         }
-
+        dataFlowAnalyzer.exitBlock(block)
         return block.compose()
     }
 
-    @Deprecated("should be removed after try/when completion")
-    private fun commonSuperType(types: List<FirTypeRef>): FirTypeRef? {
-        val commonSuperType = with(NewCommonSuperTypeCalculator) {
-            with(inferenceComponents.ctx) {
-                commonSuperType(types.map { it.coneTypeUnsafe() })
-            }
-        } as ConeKotlinType
-        return FirResolvedTypeRefImpl(null, commonSuperType, emptyList())
+    override fun <E : FirTargetElement> transformJump(jump: FirJump<E>, data: Any?): CompositeTransformResult<FirStatement> {
+        val result = super.transformJump(jump, data)
+        dataFlowAnalyzer.exitJump(jump)
+        return result
+    }
+
+    override fun transformThrowExpression(throwExpression: FirThrowExpression, data: Any?): CompositeTransformResult<FirStatement> {
+        return super.transformThrowExpression(throwExpression, data).also {
+            dataFlowAnalyzer.exitThrowExceptionNode(it.single as FirThrowExpression)
+        }
     }
 
     override fun transformWhenExpression(whenExpression: FirWhenExpression, data: Any?): CompositeTransformResult<FirStatement> {
@@ -470,6 +507,7 @@ open class FirBodyResolveTransformer(
             return whenExpression.compose()
         }
 
+        dataFlowAnalyzer.enterWhenExpression(whenExpression)
         return withScopeCleanup(localScopes) with@{
             if (whenExpression.subjectVariable != null) {
                 localScopes += FirLocalScope()
@@ -479,16 +517,25 @@ open class FirBodyResolveTransformer(
 
             @Suppress("NAME_SHADOWING")
             val whenExpression = syntheticCallGenerator.generateCalleeForWhenExpression(whenExpression) ?: run {
-                whenExpression.resultType = FirErrorTypeRefImpl(null, "Inapplicable when expression")
+                dataFlowAnalyzer.exitWhenExpression(whenExpression)
+                whenExpression.resultType = FirErrorTypeRefImpl(null, "")
                 return@with whenExpression.compose()
             }
 
             val expectedTypeRef = data as FirTypeRef?
             val result = callCompleter.completeCall(whenExpression, expectedTypeRef)
+            dataFlowAnalyzer.exitWhenExpression(result)
             result.compose()
         }
     }
 
+    override fun transformWhenBranch(whenBranch: FirWhenBranch, data: Any?): CompositeTransformResult<FirWhenBranch> {
+        return whenBranch.also { dataFlowAnalyzer.enterWhenBranchCondition(whenBranch) }
+            .transformCondition(this, data).also { dataFlowAnalyzer.exitWhenBranchCondition(it) }
+            .transformResult(this, data).also { dataFlowAnalyzer.exitWhenBranchResult(it) }
+            .compose()
+    }
+
     override fun transformWhenSubjectExpression(
         whenSubjectExpression: FirWhenSubjectExpression,
         data: Any?
@@ -529,7 +576,9 @@ open class FirBodyResolveTransformer(
         }
 
 
-        return super.transformConstExpression(constExpression, data)
+        return super.transformConstExpression(constExpression, data).also {
+            dataFlowAnalyzer.exitConstExpresion(it.single as FirConstExpression<*>)
+        }
     }
 
     override fun transformDeclaration(declaration: FirDeclaration, data: Any?): CompositeTransformResult<FirDeclaration> {
@@ -539,13 +588,21 @@ open class FirBodyResolveTransformer(
     }
 
     override fun transformAnnotationCall(annotationCall: FirAnnotationCall, data: Any?): CompositeTransformResult<FirStatement> {
-        return (annotationCall.transformChildren(this, data) as FirStatement).compose()
+        dataFlowAnalyzer.enterAnnotationCall(annotationCall)
+        return (annotationCall.transformChildren(this, data) as FirAnnotationCall).also {
+            dataFlowAnalyzer.exitAnnotationCall(it)
+        }.compose()
     }
 
     override fun <F : FirFunction<F>> transformFunction(function: FirFunction<F>, data: Any?): CompositeTransformResult<FirDeclaration> {
         return withScopeCleanup(localScopes) {
             localScopes += FirLocalScope()
-            super.transformFunction(function, data)
+            dataFlowAnalyzer.enterFunction(function)
+            super.transformFunction(function, data).also {
+                val result = it.single
+                // TODO: handle graph
+                val controlFlowGraph = dataFlowAnalyzer.exitFunction(result as FirFunction<*>)
+            }
         }
     }
 
@@ -559,7 +616,7 @@ open class FirBodyResolveTransformer(
         }
         val result = transformFunction(function, returnTypeRef).single as FirFunction<*>
         val body = result.body
-        return if (result is FirTypedDeclaration && result.returnTypeRef is FirImplicitTypeRef && body != null) {
+        return if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
             result.transformReturnTypeRef(this, body.resultType)
             result
         } else {
@@ -680,6 +737,7 @@ open class FirBodyResolveTransformer(
             localScopes.lastOrNull()?.storeDeclaration(variable)
         }
         variable.resolvePhase = transformerPhase
+        dataFlowAnalyzer.exitVariableDeclaration(variable)
         return variable.compose()
     }
 
@@ -687,6 +745,7 @@ open class FirBodyResolveTransformer(
         val returnTypeRef = property.returnTypeRef
         if (returnTypeRef !is FirImplicitTypeRef && implicitTypeOnly) return property.compose()
         if (property.resolvePhase == transformerPhase) return property.compose()
+        dataFlowAnalyzer.enterProperty(property)
         if (returnTypeRef is FirImplicitTypeRef) {
             property.transformReturnTypeRef(StoreType, FirComputingImplicitTypeRef)
         }
@@ -706,6 +765,8 @@ open class FirBodyResolveTransformer(
                     }
                 }
                 property.resolvePhase = transformerPhase
+                // TODO: handle graph
+                val controlFlowGraph = dataFlowAnalyzer.exitProperty(property)
                 property.compose()
             }
         }
@@ -755,18 +816,47 @@ open class FirBodyResolveTransformer(
         binaryLogicExpression: FirBinaryLogicExpression,
         data: Any?
     ): CompositeTransformResult<FirStatement> {
-        return super.transformBinaryLogicExpression(binaryLogicExpression, builtinTypes.booleanType).also {
-            (it.single as FirBinaryLogicExpression).resultType = builtinTypes.booleanType
-        }
+        val booleanType = builtinTypes.booleanType
+        return when (binaryLogicExpression.kind) {
+            FirBinaryLogicExpression.OperationKind.AND ->
+                binaryLogicExpression.also(dataFlowAnalyzer::enterBinaryAnd)
+                    .transformLeftOperand(this, booleanType).also(dataFlowAnalyzer::exitLeftBinaryAndArgument)
+                    .transformRightOperand(this, booleanType).also(dataFlowAnalyzer::exitBinaryAnd)
+
+            FirBinaryLogicExpression.OperationKind.OR ->
+                binaryLogicExpression.also(dataFlowAnalyzer::enterBinaryOr)
+                    .transformLeftOperand(this, booleanType).also(dataFlowAnalyzer::exitLeftBinaryOrArgument)
+                    .transformRightOperand(this, booleanType).also(dataFlowAnalyzer::exitBinaryOr)
+
+        }.transformRestChildren(this, booleanType).also {
+            it.resultType = booleanType
+        }.compose()
     }
 
     override fun transformOperatorCall(operatorCall: FirOperatorCall, data: Any?): CompositeTransformResult<FirStatement> {
-        if (operatorCall.operation in FirOperation.BOOLEANS) {
-            return (operatorCall.transformChildren(this, noExpectedType) as FirOperatorCall).also {
+        val result = if (operatorCall.operation in FirOperation.BOOLEANS) {
+            (operatorCall.transformChildren(this, noExpectedType) as FirOperatorCall).also {
                 it.resultType = builtinTypes.booleanType
-            }.compose()
-        }
-        return super.transformOperatorCall(operatorCall, data)
+            }
+        } else {
+            super.transformOperatorCall(operatorCall, data).single
+        } as FirOperatorCall
+        dataFlowAnalyzer.exitOperatorCall(result)
+        return result.compose()
+    }
+
+    override fun transformWhileLoop(whileLoop: FirWhileLoop, data: Any?): CompositeTransformResult<FirStatement> {
+        return whileLoop.also(dataFlowAnalyzer::enterWhileLoop)
+            .transformCondition(this, data).also(dataFlowAnalyzer::exitWhileLoopCondition)
+            .transformBlock(this, data).also(dataFlowAnalyzer::exitWhileLoop)
+            .transformRestChildren(this, data).compose()
+    }
+
+    override fun transformDoWhileLoop(doWhileLoop: FirDoWhileLoop, data: Any?): CompositeTransformResult<FirStatement> {
+        return doWhileLoop.also(dataFlowAnalyzer::enterDoWhileLoop)
+            .transformBlock(this, data).also(dataFlowAnalyzer::enterDoWhileLoopCondition)
+            .transformCondition(this, data).also(dataFlowAnalyzer::exitDoWhileLoop)
+            .transformRestChildren(this, data).compose()
     }
 
     // ----------------------- Util functions -----------------------

compiler/fir/resolve/testData/resolve/recursiveCallOnWhenWithSealedClass.txt

@@ -27,7 +27,7 @@ FILE: recursiveCallOnWhenWithSealedClass.kt
 
         }
 
-        public final fun unwrap(): <ERROR TYPE REF: Inapplicable when expression> {
+        public final fun unwrap(): <ERROR TYPE REF: > {
             ^unwrap when (this@R|/Maybe|) {
                 ($subj$ is R|Maybe.Nope|) ->  {
                     throw R|java/lang/Exception.Exception|(String())

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/impl/FirErrorLoop.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirErrorExpressionImpl
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
 class FirErrorLoop(
@@ -29,4 +30,16 @@ class FirErrorLoop(
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         super<FirLoop>.acceptChildren(visitor, data)
     }
+
+    override fun <D> transformCondition(transformer: FirTransformer<D>, data: D): FirLoop {
+        return this
+    }
+
+    override fun <D> transformBlock(transformer: FirTransformer<D>, data: D): FirLoop {
+        return this
+    }
+
+    override fun <D> transformRestChildren(transformer: FirTransformer<D>, data: D): FirLoop {
+        return this
+    }
 }

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

@@ -27,6 +27,7 @@ abstract class FirBinaryLogicExpression(psi: PsiElement?) : FirUnknownTypeExpres
 
     abstract fun <D> transformLeftOperand(transformer: FirTransformer<D>, data: D): FirBinaryLogicExpression
     abstract fun <D> transformRightOperand(transformer: FirTransformer<D>, data: D): FirBinaryLogicExpression
+    abstract fun <D> transformRestChildren(transformer: FirTransformer<D>, data: D): FirBinaryLogicExpression
 
     enum class OperationKind(val token: String) {
         AND("&&"), OR("||")

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.expressions
 
 import org.jetbrains.kotlin.fir.FirLabeledElement
 import org.jetbrains.kotlin.fir.VisitedSupertype
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
 interface FirLoop : @VisitedSupertype FirStatement, FirLabeledElement, FirAnnotationContainer {
@@ -22,4 +23,8 @@ interface FirLoop : @VisitedSupertype FirStatement, FirLabeledElement, FirAnnota
         block.accept(visitor, data)
         super<FirLabeledElement>.acceptChildren(visitor, data)
     }
+
+    fun <D> transformCondition(transformer: FirTransformer<D>, data: D): FirLoop
+    fun <D> transformBlock(transformer: FirTransformer<D>, data: D): FirLoop
+    fun <D> transformRestChildren(transformer: FirTransformer<D>, data: D): FirLoop
 }

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

@@ -6,10 +6,12 @@
 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.FirLabel
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirLoop
+import org.jetbrains.kotlin.fir.transformSingle
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
@@ -27,9 +29,23 @@ abstract class FirAbstractLoop(
     }
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
+        return transformCondition(transformer, data)
+            .transformBlock(transformer, data)
+            .transformRestChildren(transformer, data)
+    }
+
+    override fun <D> transformCondition(transformer: FirTransformer<D>, data: D): FirAbstractLoop {
         condition = condition.transformSingle(transformer, data)
+        return this
+    }
+
+    override fun <D> transformBlock(transformer: FirTransformer<D>, data: D): FirAbstractLoop {
         block = block.transformSingle(transformer, data)
+        return this
+    }
+
+    override fun <D> transformRestChildren(transformer: FirTransformer<D>, data: D): FirAbstractLoop {
         label = label?.transformSingle(transformer, data)
-        return super<FirAnnotatedStatement>.transformChildren(transformer, data)
+        return super<FirAnnotatedStatement>.transformChildren(transformer, data) as FirAbstractLoop
     }
 }

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

@@ -21,7 +21,7 @@ class FirBinaryLogicExpressionImpl(
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
         transformLeftOperand(transformer, data)
         transformRightOperand(transformer, data)
-        return super.transformChildren(transformer, data)
+        return transformRestChildren(transformer, data)
     }
 
     override fun <D> transformLeftOperand(transformer: FirTransformer<D>, data: D): FirBinaryLogicExpression {
@@ -33,4 +33,8 @@ class FirBinaryLogicExpressionImpl(
         rightOperand = rightOperand.transformSingle(transformer, data)
         return this
     }
+
+    override fun <D> transformRestChildren(transformer: FirTransformer<D>, data: D): FirBinaryLogicExpression {
+        return super.transformChildren(transformer, data) as FirBinaryLogicExpression
+    }
 }

idea/testData/fir/lazyResolve/simpleProperty/results.txt

@@ -1 +1 @@
-R|root/myProperty|
+R|root/myProperty|