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

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