From a20dea1b0a52d865fa13232c212c82c172bf5848 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Thu, 7 Sep 2023 15:04:01 +0200
Subject: [PATCH] KT-61728 [LL API] Ignore CFG nodes with uninitialized flow

This is done until the KT-61794 is fixed
---
 .../level/api/fir/util/ContextCollector.kt    | 10 ++++--
 .../testdata/contextCollector/KT-61728.kt     |  5 +++
 .../testdata/contextCollector/KT-61728.txt    | 35 +++++++++++++++++++
 .../ContextCollectorSourceTestGenerated.java  |  6 ++++
 4 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 analysis/low-level-api-fir/testdata/contextCollector/KT-61728.kt
 create mode 100644 analysis/low-level-api-fir/testdata/contextCollector/KT-61728.txt

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
index 8fd23efaea1..1f6edad67cb 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.fir.resolve.dfa.PropertyStability
 import org.jetbrains.kotlin.fir.resolve.dfa.RealVariable
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ClassExitNode
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.MergePostponedLambdaExitsNode
 import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph
 import org.jetbrains.kotlin.fir.resolve.dfa.smartCastedType
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
@@ -283,8 +284,13 @@ private class ContextCollectorVisitor(
         return null
     }
 
-    private fun isAcceptedControlFlowNode(node: CFGNode<*>): Boolean {
-        return node !is ClassExitNode
+    private fun isAcceptedControlFlowNode(node: CFGNode<*>): Boolean = when {
+        node is ClassExitNode -> false
+
+        // TODO Remove as soon as KT-61728 is fixed
+        node is MergePostponedLambdaExitsNode && !node.flowInitialized -> false
+
+        else -> true
     }
 
     override fun visitScript(script: FirScript) {
diff --git a/analysis/low-level-api-fir/testdata/contextCollector/KT-61728.kt b/analysis/low-level-api-fir/testdata/contextCollector/KT-61728.kt
new file mode 100644
index 00000000000..87333365a63
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/contextCollector/KT-61728.kt
@@ -0,0 +1,5 @@
+fun foo(param: String?): String? {
+    val tmp = param?.let {
+        <expr>it</expr>
+    } ?: return null
+}
diff --git a/analysis/low-level-api-fir/testdata/contextCollector/KT-61728.txt b/analysis/low-level-api-fir/testdata/contextCollector/KT-61728.txt
new file mode 100644
index 00000000000..3db2c14994a
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/contextCollector/KT-61728.txt
@@ -0,0 +1,35 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol param: R|kotlin/String?|
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol it: R|kotlin/String|
+    Element 10
+        Scope: FirLocalScope
+
+FILE: [ResolvedTo(IMPORTS)] KT-61728.kt
+    public final [ResolvedTo(BODY_RESOLVE)] fun foo([ResolvedTo(BODY_RESOLVE)] param: R|kotlin/String?|): R|kotlin/String?| {
+        [ResolvedTo(BODY_RESOLVE)] lval tmp: R|kotlin/String| = R|<local>/param|?.{ $subj$.R|kotlin/let|<R|kotlin/String|, R|kotlin/String|>(<L> = [ResolvedTo(BODY_RESOLVE)] let@fun <anonymous>([ResolvedTo(BODY_RESOLVE)] it: R|kotlin/String|): R|kotlin/String| <inline=Inline, kind=EXACTLY_ONCE>  {
+            ^ R|<local>/it|
+        }
+        ) } ?: ^foo Null(null)
+    }
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java
index 614ba332626..a0e4e6be8be 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorSourceTestGenerated.java
@@ -66,6 +66,12 @@ public class ContextCollectorSourceTestGenerated extends AbstractContextCollecto
         runTest("analysis/low-level-api-fir/testdata/contextCollector/innerClasses.kt");
     }
 
+    @Test
+    @TestMetadata("KT-61728.kt")
+    public void testKT_61728() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/KT-61728.kt");
+    }
+
     @Test
     @TestMetadata("lambdaArguments.kt")
     public void testLambdaArguments() throws Exception {