From fd54706e1a0ccd659afa248b1c43d830b43c2b90 Mon Sep 17 00:00:00 2001
From: Nikolay Lunyak <nikolay.lunyak@jetbrains.com>
Date: Mon, 15 May 2023 11:54:39 +0300
Subject: [PATCH] [FIR] KT-58674: Set expected type for loop conditions

^KT-58674 Fixed
---
 .../resolveWithStdlib/diagnostics/forLoopChecker.fir.txt    | 4 ++--
 .../resolve/FirControlFlowStatementsResolveTransformer.kt   | 6 ++++--
 .../diagnostics/tests/whileConditionExpectedType.fir.kt     | 4 ++--
 .../diagnostics/tests/whileConditionExpectedType.kt         | 4 ++--
 4 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/compiler/fir/analysis-tests/testData/resolveWithStdlib/diagnostics/forLoopChecker.fir.txt b/compiler/fir/analysis-tests/testData/resolveWithStdlib/diagnostics/forLoopChecker.fir.txt
index 390a054d42f..35573b4787e 100644
--- a/compiler/fir/analysis-tests/testData/resolveWithStdlib/diagnostics/forLoopChecker.fir.txt
+++ b/compiler/fir/analysis-tests/testData/resolveWithStdlib/diagnostics/forLoopChecker.fir.txt
@@ -217,7 +217,7 @@ FILE: forLoopChecker.kt
 
          {
             lval <iterator>: R|ImproperIterator3| = R|<local>/notRange5|.R|/NotRange5.iterator|()
-            while(R|<local>/<iterator>|.R|/ImproperIterator3.hasNext|()) {
+            while(R|<local>/<iterator>|.R|/ImproperIterator3.hasNext<CS errors: /ImproperIterator3.hasNext>#|()) {
                 lval i: R|kotlin/Int| = R|<local>/<iterator>|.R|/ImproperIterator3.next|()
                  {
                 }
@@ -239,7 +239,7 @@ FILE: forLoopChecker.kt
 
          {
             lval <iterator>: R|ImproperIterator3| = R|<local>/notRange7|.R|/NotRange7.iterator|()
-            while(R|<local>/<iterator>|.R|/ImproperIterator3.hasNext|()) {
+            while(R|<local>/<iterator>|.R|/ImproperIterator3.hasNext<CS errors: /ImproperIterator3.hasNext>#|()) {
                 lval i: R|kotlin/Int| = R|<local>/<iterator>|.R|/ImproperIterator3.next|()
                  {
                 }
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
index cbb1e739575..e2f508b3f3c 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
@@ -30,7 +30,8 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirAbstractBodyRes
     override fun transformWhileLoop(whileLoop: FirWhileLoop, data: ResolutionMode): FirStatement {
         val context = ResolutionMode.ContextIndependent
         return whileLoop.also(dataFlowAnalyzer::enterWhileLoop)
-            .transformCondition(transformer, context).also(dataFlowAnalyzer::exitWhileLoopCondition)
+            .transformCondition(transformer, withExpectedType(session.builtinTypes.booleanType))
+            .also(dataFlowAnalyzer::exitWhileLoopCondition)
             .transformBlock(transformer, context).also(dataFlowAnalyzer::exitWhileLoop)
             .transformOtherChildren(transformer, context)
     }
@@ -43,7 +44,8 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirAbstractBodyRes
                 .also {
                     transformer.expressionsTransformer.transformBlockInCurrentScope(it.block, context)
                 }
-                .also(dataFlowAnalyzer::enterDoWhileLoopCondition).transformCondition(transformer, context)
+                .also(dataFlowAnalyzer::enterDoWhileLoopCondition)
+                .transformCondition(transformer, withExpectedType(session.builtinTypes.booleanType))
                 .also(dataFlowAnalyzer::exitDoWhileLoop)
                 .transformOtherChildren(transformer, context)
         }
diff --git a/compiler/testData/diagnostics/tests/whileConditionExpectedType.fir.kt b/compiler/testData/diagnostics/tests/whileConditionExpectedType.fir.kt
index a9732dde4a2..d7a8daddd1d 100644
--- a/compiler/testData/diagnostics/tests/whileConditionExpectedType.fir.kt
+++ b/compiler/testData/diagnostics/tests/whileConditionExpectedType.fir.kt
@@ -2,13 +2,13 @@
 // INFERENCE_HELPERS
 
 fun test() {
-    while (<!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>materialize<!>()) { // K1: OK, K2: NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER
+    while (materialize()) {
 
     }
 
     do {
 
-    } while (<!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>materialize<!>()) // K1: OK, K2: NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER
+    } while (materialize())
 
     if (materialize()) {
 
diff --git a/compiler/testData/diagnostics/tests/whileConditionExpectedType.kt b/compiler/testData/diagnostics/tests/whileConditionExpectedType.kt
index 14244337fdd..62bcfc7e1d3 100644
--- a/compiler/testData/diagnostics/tests/whileConditionExpectedType.kt
+++ b/compiler/testData/diagnostics/tests/whileConditionExpectedType.kt
@@ -2,13 +2,13 @@
 // INFERENCE_HELPERS
 
 fun test() {
-    while (materialize()) { // K1: OK, K2: NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER
+    while (materialize()) {
 
     }
 
     do {
 
-    } while (materialize()) // K1: OK, K2: NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER
+    } while (materialize())
 
     if (materialize()) {