FIR IDE: add expected type calculation for function calls

does not work in case of incomplete call expression :(

idea/idea-frontend-fir/build.gradle.kts

@@ -32,6 +32,7 @@ dependencies {
     testCompile(projectTests(":idea:idea-test-framework"))
     testCompile(project(":kotlin-test:kotlin-test-junit"))
     testCompile(commonDep("junit:junit"))
+    testCompile(projectTests(":idea:idea-frontend-independent"))
 
     compile(intellijPluginDep("java"))
 }

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

@@ -6,12 +6,19 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.psi.PsiElement
+import com.jetbrains.rd.util.firstOrNull
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.expressions.FirCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.argumentMapping
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.types.ConeTypeCheckerContext
 import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.assertIsValid
 import org.jetbrains.kotlin.idea.frontend.api.components.KtBuiltinTypes
@@ -23,6 +30,7 @@ import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 import kotlin.reflect.KClass
 
 internal class KtFirTypeProvider(
@@ -41,6 +49,23 @@ internal class KtFirTypeProvider(
     override fun getExpectedType(expression: PsiElement): KtType? =
         getExpectedTypeByReturnExpression(expression)
             ?: getExpressionTypeByIfOrBooleanCondition(expression)
+            ?: getExpectedTypeOfFunctionParameter(expression)
+
+    private fun getExpectedTypeOfFunctionParameter(expression: PsiElement): KtType? {
+        val (ktCallExpression, ktArgument) = expression.getFunctionCallAsWithThisAsParameter() ?: return null
+        val firCall = ktCallExpression.getOrBuildFirSafe<FirFunctionCall>(firResolveState) ?: return null
+        val arguments = firCall.argumentMapping ?: return null
+        val firParameterForExpression = arguments.entries.firstOrNull { (arg, _) -> arg.psi == ktArgument }?.value ?: return null
+        return firParameterForExpression.returnTypeRef.coneType.asKtType()
+    }
+
+    private fun PsiElement.getFunctionCallAsWithThisAsParameter(): KtCallWithArgument? {
+        val valueArgument = unwrapQualified<KtValueArgument> { valueArg, expr -> valueArg.getArgumentExpression() == expr } ?: return null
+        val argumentsList = valueArgument.parent as? KtValueArgumentList ?: return null
+        val callExpression = argumentsList.parent as? KtCallExpression ?: return null
+        val argumentExpression = valueArgument.getArgumentExpression() ?: return null
+        return KtCallWithArgument(callExpression, argumentExpression)
+    }
 
     private fun getExpectedTypeByReturnExpression(expression: PsiElement): KtType? {
         val returnParent = expression.getReturnExpressionWithThisType() ?: return null
@@ -99,6 +124,8 @@ internal class KtFirTypeProvider(
     )
 }
 
+private data class KtCallWithArgument(val call: KtCallExpression, val argument: KtExpression)
+
 private inline fun <reified R : Any> PsiElement.unwrapQualified(check: (R, PsiElement) -> Boolean): R? {
     val parent = nonContainerParent
     return when {

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

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

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

@@ -0,0 +1,9 @@
+// FIX_ME: should work on non fully resolved calls
+
+fun x() {
+    toCall(1, z = a<caret>v)
+}
+
+fun toCall(x: Int, y: String, z: Boolean): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/Boolean

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

@@ -0,0 +1,9 @@
+// FIX_ME: should work on non fully resolved calls
+
+fun x() {
+    toCall(a<caret>b, 12)
+}
+
+fun <T> toCall(x: T, y: T): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/Int

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

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

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/components/AbstractExpectedExpressionTypeTest.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
+import org.jetbrains.kotlin.test.uitls.IgnoreTests
 import java.io.File
 
 abstract class AbstractExpectedExpressionTypeTest : KotlinLightCodeInsightFixtureTestCase() {
@@ -37,7 +38,9 @@ abstract class AbstractExpectedExpressionTypeTest : KotlinLightCodeInsightFixtur
             }
         }
 
-        KotlinTestUtils.assertEqualsToFile(File(path), testDataFile.getTextWithActualType(actualExpectedTypeText))
+        IgnoreTests.runTestWithFixMeSupport(testDataFile.toPath()) {
+            KotlinTestUtils.assertEqualsToFile(File(path), testDataFile.getTextWithActualType(actualExpectedTypeText))
+        }
     }
 
     private fun File.getTextWithActualType(actualType: String?) : String {

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

@@ -28,6 +28,26 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/expectedExpressionType"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @TestMetadata("functionLambdaParam.kt")
+    public void testFunctionLambdaParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionLambdaParam.kt");
+    }
+
+    @TestMetadata("functionNamedlParam.kt")
+    public void testFunctionNamedlParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionNamedlParam.kt");
+    }
+
+    @TestMetadata("functionParamWithTypeParam.kt")
+    public void testFunctionParamWithTypeParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionParamWithTypeParam.kt");
+    }
+
+    @TestMetadata("functionPositionalParam.kt")
+    public void testFunctionPositionalParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionPositionalParam.kt");
+    }
+
     @TestMetadata("ifCondition.kt")
     public void testIfCondition() throws Exception {
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/ifCondition.kt");

idea/idea-frontend-independent/tests/org/jetbrains/kotlin/test/uitls/IgnoreTests.kt

@@ -27,6 +27,15 @@ object IgnoreTests {
         )
     }
 
+    fun runTestWithFixMeSupport(testFile: Path, test: () -> Unit) {
+        runTestIfEnabledByDirective(
+            testFile,
+            EnableOrDisableTestDirective.Disable(DIRECTIVES.FIX_ME),
+            additionalFilesExtensions = emptyList(),
+            test = test
+        )
+    }
+
     fun runTestIfNotDisabledByFileDirective(
         testFile: Path,
         disableTestDirective: String,
@@ -128,5 +137,6 @@ object IgnoreTests {
     object DIRECTIVES {
         const val FIR_COMPARISON = "// FIR_COMPARISON"
         const val IGNORE_FIR = "// IGNORE_FIR"
+        const val FIX_ME = "// FIX_ME: "
     }
 }