FIR IDE: Additional support for

KtFirExpressionTypeProvider.getExpectedType():

- Infix function parameter
- Variable assignment
- Property declaration
- Function expression body

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.frontend.api.components
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
@@ -25,6 +24,10 @@ interface KtExpressionTypeProviderMixIn : KtAnalysisSessionMixIn {
     fun KtDeclaration.getReturnKtType(): KtType =
         analysisSession.expressionTypeProvider.getReturnTypeForKtDeclaration(this)
 
+    /**
+     * Returns the expected [KtType] of this [PsiElement] if it is an expression. The returned value should not be a
+     * [org.jetbrains.kotlin.idea.frontend.api.types.KtErrorType].
+     */
     fun PsiElement.getExpectedType(): KtType? =
         analysisSession.expressionTypeProvider.getExpectedType(this)
 }

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

@@ -7,7 +7,10 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
-import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.expressions.argumentMapping
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirNamedReference
@@ -23,9 +26,11 @@ import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.components.KtExpressionTypeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtTypedSymbol
+import org.jetbrains.kotlin.idea.frontend.api.types.KtErrorType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper
+import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 
 internal class KtFirExpressionTypeProvider(
@@ -57,10 +62,16 @@ internal class KtFirExpressionTypeProvider(
             ?: ConeClassErrorType(ConeUnresolvedNameError(name))
     }
 
-    override fun getExpectedType(expression: PsiElement): KtType? =
+    override fun getExpectedType(expression: PsiElement): KtType? {
-        getExpectedTypeByReturnExpression(expression)
+        val expectedType = getExpectedTypeByReturnExpression(expression)
             ?: getExpressionTypeByIfOrBooleanCondition(expression)
             ?: getExpectedTypeOfFunctionParameter(expression)
+            ?: getExpectedTypeOfInfixFunctionParameter(expression)
+            ?: getExpectedTypeByVariableAssignment(expression)
+            ?: getExpectedTypeByPropertyDeclaration(expression)
+            ?: getExpectedTypeByFunctionExpressionBody(expression)
+        return expectedType.takeIf { it !is KtErrorType }
+    }
 
     private fun getExpectedTypeOfFunctionParameter(expression: PsiElement): KtType? {
         val (ktCallExpression, ktArgument) = expression.getFunctionCallAsWithThisAsParameter() ?: return null
@@ -78,6 +89,17 @@ internal class KtFirExpressionTypeProvider(
         return KtCallWithArgument(callExpression, argumentExpression)
     }
 
+    private fun getExpectedTypeOfInfixFunctionParameter(expression: PsiElement): KtType? {
+        val infixCallExpression =
+            expression.unwrapQualified<KtBinaryExpression> { binaryExpr, expr -> binaryExpr.right == expr } ?: return null
+        val firCall = infixCallExpression.getOrBuildFirSafe<FirFunctionCall>(firResolveState) ?: return null
+
+        // There is only one parameter for infix functions; get its type
+        val arguments = firCall.argumentMapping ?: return null
+        val firParameterForExpression = arguments.values.singleOrNull() ?: return null
+        return firParameterForExpression.returnTypeRef.coneType.asKtType()
+    }
+
     private fun getExpectedTypeByReturnExpression(expression: PsiElement): KtType? {
         val returnParent = expression.getReturnExpressionWithThisType() ?: return null
         val targetSymbol = with(analysisSession) { returnParent.getReturnTargetSymbol() } ?: return null
@@ -92,6 +114,30 @@ internal class KtFirExpressionTypeProvider(
         else -> null
     }
 
+    private fun getExpectedTypeByVariableAssignment(expression: PsiElement): KtType? {
+        // Given: `x = expression`
+        // Expected type of `expression` is type of `x`
+        val assignmentExpression =
+            expression.unwrapQualified<KtBinaryExpression> { binaryExpr, expr -> binaryExpr.right == expr && binaryExpr.operationToken == KtTokens.EQ }
+                ?: return null
+        val variableExpression = assignmentExpression.left as? KtNameReferenceExpression ?: return null
+        return getKtExpressionType(variableExpression)
+    }
+
+    private fun getExpectedTypeByPropertyDeclaration(expression: PsiElement): KtType? {
+        // Given: `val x: T = expression`
+        // Expected type of `expression` is `T`
+        val property = expression.unwrapQualified<KtProperty> { property, expr -> property.initializer == expr } ?: return null
+        return getReturnTypeForKtDeclaration(property)
+    }
+
+    private fun getExpectedTypeByFunctionExpressionBody(expression: PsiElement): KtType? {
+        // Given: `fun f(): T = expression`
+        // Expected type of `expression` is `T`
+        val function = expression.unwrapQualified<KtFunction> { function, expr -> function.bodyExpression == expr } ?: return null
+        return getReturnTypeForKtDeclaration(function)
+    }
+
     private fun PsiElement.isWhileLoopCondition() =
         unwrapQualified<KtWhileExpressionBase> { whileExpr, cond -> whileExpr.condition == cond } != null
 

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

@@ -0,0 +1,3 @@
+fun foo(): String = a<caret>v
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,3 @@
+fun foo(): String = x.a<caret>v
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,5 @@
+val av = "foo"
+
+fun foo() = a<caret>v
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,3 @@
+fun foo() = a<caret>v
+
+// EXPECTED_TYPE: null

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

@@ -0,0 +1,7 @@
+fun x() {
+    toCall(1, x.a<caret>v, true)
+}
+
+fun toCall(x: Int, y: String, z: Boolean): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,7 @@
+fun x() {
+    1.toCall(x.a<caret>v)
+}
+
+infix fun Int.toCall(y: String): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,7 @@
+fun x() {
+    1 toCall a<caret>v
+}
+
+infix fun Int.toCall(y: String): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,7 @@
+fun x() {
+    1 toCall x.a<caret>v
+}
+
+infix fun Int.toCall(y: String): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,5 @@
+class C {
+    var y: String = a<caret>v
+}
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,5 @@
+class C {
+    var y: String = x.a<caret>v
+}
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,6 @@
+class C {
+    val av = "foo"
+    var y = a<caret>v
+}
+
+// EXPECTED_TYPE: kotlin/String

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

@@ -0,0 +1,5 @@
+class C {
+    var y = x.a<caret>v
+}
+
+// EXPECTED_TYPE: null

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

@@ -0,0 +1,6 @@
+fun foo() {
+    var y = 1
+    y = a<caret>v
+}
+
+// EXPECTED_TYPE: kotlin/Int

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

@@ -0,0 +1,6 @@
+fun foo() {
+    var y = 1
+    y = x.a<caret>v
+}
+
+// EXPECTED_TYPE: kotlin/Int

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

@@ -29,6 +29,26 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/expectedExpressionType"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @TestMetadata("functionExpressionBody.kt")
+    public void testFunctionExpressionBody() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBody.kt");
+    }
+
+    @TestMetadata("functionExpressionBodyQualified.kt")
+    public void testFunctionExpressionBodyQualified() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBodyQualified.kt");
+    }
+
+    @TestMetadata("functionExpressionBodyWithTypeFromRHS.kt")
+    public void testFunctionExpressionBodyWithTypeFromRHS() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBodyWithTypeFromRHS.kt");
+    }
+
+    @TestMetadata("functionExpressionBodyWithoutExplicitType.kt")
+    public void testFunctionExpressionBodyWithoutExplicitType() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBodyWithoutExplicitType.kt");
+    }
+
     @TestMetadata("functionLambdaParam.kt")
     public void testFunctionLambdaParam() throws Exception {
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionLambdaParam.kt");
@@ -49,6 +69,11 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionPositionalParam.kt");
     }
 
+    @TestMetadata("functionPositionalParamQualified.kt")
+    public void testFunctionPositionalParamQualified() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionPositionalParamQualified.kt");
+    }
+
     @TestMetadata("ifCondition.kt")
     public void testIfCondition() throws Exception {
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/ifCondition.kt");
@@ -59,6 +84,41 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/ifConditionQualified.kt");
     }
 
+    @TestMetadata("infixFunctionAsRegularCallParam.kt")
+    public void testInfixFunctionAsRegularCallParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/infixFunctionAsRegularCallParam.kt");
+    }
+
+    @TestMetadata("infixFunctionParam.kt")
+    public void testInfixFunctionParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/infixFunctionParam.kt");
+    }
+
+    @TestMetadata("infixFunctionParamQualified.kt")
+    public void testInfixFunctionParamQualified() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/infixFunctionParamQualified.kt");
+    }
+
+    @TestMetadata("propertyDeclaration.kt")
+    public void testPropertyDeclaration() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclaration.kt");
+    }
+
+    @TestMetadata("propertyDeclarationQualified.kt")
+    public void testPropertyDeclarationQualified() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationQualified.kt");
+    }
+
+    @TestMetadata("propertyDeclarationWithTypeFromRHS.kt")
+    public void testPropertyDeclarationWithTypeFromRHS() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithTypeFromRHS.kt");
+    }
+
+    @TestMetadata("propertyDeclarationWithoutExplicitType.kt")
+    public void testPropertyDeclarationWithoutExplicitType() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithoutExplicitType.kt");
+    }
+
     @TestMetadata("returnFromFunction.kt")
     public void testReturnFromFunction() throws Exception {
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromFunction.kt");
@@ -79,6 +139,16 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromLambda.kt");
     }
 
+    @TestMetadata("variableAssignment.kt")
+    public void testVariableAssignment() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/variableAssignment.kt");
+    }
+
+    @TestMetadata("variableAssignmentQualified.kt")
+    public void testVariableAssignmentQualified() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/variableAssignmentQualified.kt");
+    }
+
     @TestMetadata("whileCondition.kt")
     public void testWhileCondition() throws Exception {
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/whileCondition.kt");