FIR: Add callable references nodes to CFG

Otherwise, they are not being considered as return expressions from lambdas

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

@@ -1112,6 +1112,12 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         graphBuilder.exitElvis().mergeIncomingFlow()
     }
 
+    // Callable reference
+
+    fun exitCallableReference(callableReferenceAccess: FirCallableReferenceAccess) {
+        graphBuilder.exitCallableReference(callableReferenceAccess).mergeIncomingFlow()
+    }
+
     // ------------------------------------------------------ Utils ------------------------------------------------------
 
     private var CFGNode<*>.flow: FLOW

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/CFGNode.kt

@@ -564,6 +564,16 @@ class FunctionCallNode(
     }
 }
 
+class CallableReferenceNode(
+    owner: ControlFlowGraph,
+    override val fir: FirCallableReferenceAccess,
+    level: Int, id: Int
+) : CFGNode<FirCallableReferenceAccess>(owner, level, id) {
+    override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
+        return visitor.visitCallableReferenceNode(this, data)
+    }
+}
+
 class DelegatedConstructorCallNode(
     owner: ControlFlowGraph,
     override val fir: FirDelegatedConstructorCall,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/CFGNodeRenderer.kt

@@ -116,6 +116,8 @@ fun CFGNode<*>.render(): String =
                 is ElvisRhsEnterNode -> "Enter rhs of ?:"
                 is ElvisExitNode -> "Exit ?:"
 
+                is CallableReferenceNode -> "Callable reference: ${fir.render(CfgRenderMode)}"
+
                 is AbstractBinaryExitNode -> throw IllegalStateException()
             },
         )
@@ -141,4 +143,4 @@ private fun FirLoop.type(): String = when (this) {
     is FirWhileLoop -> "while"
     is FirDoWhileLoop -> "do-while"
     else -> throw IllegalArgumentException()
-}
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -1039,6 +1039,12 @@ class ControlFlowGraphBuilder {
         return node
     }
 
+    // ----------------------------------- Callable references -----------------------------------
+
+    fun exitCallableReference(callableReferenceAccess: FirCallableReferenceAccess): CallableReferenceNode {
+        return createCallableReferenceNode(callableReferenceAccess).also { addNewSimpleNode(it) }
+    }
+
     // ----------------------------------- Block -----------------------------------
 
     fun enterInitBlock(initBlock: FirAnonymousInitializer): Pair<InitBlockEnterNode, CFGNode<*>?> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphNodeBuilder.kt

@@ -119,6 +119,9 @@ fun ControlFlowGraphBuilder.createLoopBlockExitNode(fir: FirLoop): LoopBlockExit
 fun ControlFlowGraphBuilder.createFunctionCallNode(fir: FirFunctionCall): FunctionCallNode =
     FunctionCallNode(currentGraph, fir, levelCounter, createId())
 
+fun ControlFlowGraphBuilder.createCallableReferenceNode(fir: FirCallableReferenceAccess): CallableReferenceNode =
+    CallableReferenceNode(currentGraph, fir, levelCounter, createId())
+
 fun ControlFlowGraphBuilder.createDelegatedConstructorCallNode(fir: FirDelegatedConstructorCall): DelegatedConstructorCallNode =
     DelegatedConstructorCallNode(currentGraph, fir, levelCounter, createId())
 
@@ -231,4 +234,4 @@ fun ControlFlowGraphBuilder.createExitDefaultArgumentsNode(fir: FirValueParamete
     ExitDefaultArgumentsNode(currentGraph, fir, levelCounter, createId())
 
 fun ControlFlowGraphBuilder.createComparisonExpressionNode(fir: FirComparisonExpression): ComparisonExpressionNode =
-    ComparisonExpressionNode(currentGraph, fir, levelCounter, createId())
+    ComparisonExpressionNode(currentGraph, fir, levelCounter, createId())

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphVisitor.kt

@@ -277,6 +277,10 @@ abstract class ControlFlowGraphVisitor<out R, in D> {
         return visitNode(node, data)
     }
 
+    open fun visitCallableReferenceNode(node: CallableReferenceNode, data: D): R {
+        return visitNode(node, data)
+    }
+
     open fun visitDelegatedConstructorCallNode(node: DelegatedConstructorCallNode, data: D): R {
         return visitNode(node, data)
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -660,9 +660,11 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                     callableReferenceAccess, data.expectedType, resolutionContext,
                 ) ?: callableReferenceAccess
 
+            dataFlowAnalyzer.exitCallableReference(resolvedReference)
             return resolvedReference.compose()
         }
 
+        dataFlowAnalyzer.exitCallableReference(callableReferenceAccessWithTransformedLHS)
         return callableReferenceAccessWithTransformedLHS.compose()
     }