FIR IDE: consider if & while conditions expected type as boolean

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirTypeProvider.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.psi.PsiElement
-import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
@@ -21,12 +20,10 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
-import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtQualifiedExpression
-import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
+import kotlin.reflect.KClass
 
 internal class KtFirTypeProvider(
     override val analysisSession: KtFirAnalysisSession,
@@ -43,6 +40,7 @@ internal class KtFirTypeProvider(
 
     override fun getExpectedType(expression: PsiElement): KtType? =
         getExpectedTypeByReturnExpression(expression)
+            ?: getExpressionTypeByIfOrBooleanCondition(expression)
 
     private fun getExpectedTypeByReturnExpression(expression: PsiElement): KtType? {
         val returnParent = expression.getReturnExpressionWithThisType() ?: return null
@@ -50,15 +48,20 @@ internal class KtFirTypeProvider(
         return targetSymbol.type
     }
 
-    private fun PsiElement.getReturnExpressionWithThisType(): KtReturnExpression? {
+    private fun PsiElement.getReturnExpressionWithThisType(): KtReturnExpression? =
-        val parent = parent
+        unwrapQualified { returnExpr, target -> returnExpr.returnedExpression == target }
-        return when {
+
-            parent is KtReturnExpression && parent.returnedExpression == this -> parent
+    private fun getExpressionTypeByIfOrBooleanCondition(expression: PsiElement): KtType? = when {
-            parent is KtQualifiedExpression && parent.selectorExpression == this -> parent.getReturnExpressionWithThisType()
+        expression.isWhileLoopCondition() || expression.isIfCondition() -> builtinTypes.BOOLEAN
-            else -> null
+        else -> null
-        }
     }
 
+    private fun PsiElement.isWhileLoopCondition() =
+        unwrapQualified<KtWhileExpressionBase> { whileExpr, cond -> whileExpr.condition == cond } != null
+
+    private fun PsiElement.isIfCondition() =
+        unwrapQualified<KtIfExpression> { ifExpr, cond -> ifExpr.condition == cond } != null
+
     override fun isEqualTo(first: KtType, second: KtType): Boolean = withValidityAssertion {
         second.assertIsValid()
         check(first is KtFirType)
@@ -94,4 +97,25 @@ internal class KtFirTypeProvider(
         isStubTypeEqualsToAnything = true,
         analysisSession.firResolveState.rootModuleSession //TODO use correct session here
     )
-}
+}
+
+private inline fun <reified R : Any> PsiElement.unwrapQualified(check: (R, PsiElement) -> Boolean): R? {
+    val parent = nonContainerParent
+    return when {
+        parent is R && check(parent, this) -> parent
+        parent is KtQualifiedExpression && parent.selectorExpression == this -> {
+            val grandParent = parent.nonContainerParent
+            when {
+                grandParent is R && check(grandParent, parent) -> grandParent
+                else -> null
+            }
+        }
+        else -> null
+    }
+}
+
+private val PsiElement.nonContainerParent: PsiElement?
+    get() = when (val parent = parent) {
+        is KtContainerNode -> parent.parent
+        else -> parent
+    }

idea/idea-frontend-fir/testData/components/expectedExpressionType/ifCondition.kt

@@ -0,0 +1,5 @@
+fun x() {
+    if(x<caret>y){
+}
+
+// EXPECTED_TYPE: kotlin/Boolean

idea/idea-frontend-fir/testData/components/expectedExpressionType/ifConditionQualified.kt

@@ -0,0 +1,5 @@
+fun x() {
+    if(x.fd<caret>fd){
+}
+
+// EXPECTED_TYPE: kotlin/Boolean

idea/idea-frontend-fir/testData/components/expectedExpressionType/whileCondition.kt

@@ -0,0 +1,5 @@
+fun x() {
+    while(x<caret>y){
+}
+
+// EXPECTED_TYPE: kotlin/Boolean

idea/idea-frontend-fir/testData/components/expectedExpressionType/whileConditionQualified.kt

@@ -0,0 +1,5 @@
+fun x() {
+    while(x.fd<caret>fd){
+}
+
+// EXPECTED_TYPE: kotlin/Boolean

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/components/ExpectedExpressionTypeTestGenerated.java

@@ -28,6 +28,16 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/expectedExpressionType"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @TestMetadata("ifCondition.kt")
+    public void testIfCondition() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/ifCondition.kt");
+    }
+
+    @TestMetadata("ifConditionQualified.kt")
+    public void testIfConditionQualified() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/ifConditionQualified.kt");
+    }
+
     @TestMetadata("returnFromFunction.kt")
     public void testReturnFromFunction() throws Exception {
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromFunction.kt");
@@ -47,4 +57,14 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
     public void testReturnFromLambda() throws Exception {
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromLambda.kt");
     }
+
+    @TestMetadata("whileCondition.kt")
+    public void testWhileCondition() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/whileCondition.kt");
+    }
+
+    @TestMetadata("whileConditionQualified.kt")
+    public void testWhileConditionQualified() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/whileConditionQualified.kt");
+    }
 }