From dd0e3cd1358d4f7fd30c116e9f25e1c4612bcf51 Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <Nicolay.Mitropolsky@jetbrains.com>
Date: Sat, 7 Jul 2018 11:13:48 +0300
Subject: [PATCH] Uast: `ULambdaExpression` returns proper
 `functionalInterfaceType` for `SamAdapterExtensionFunctionDescriptor`
 (KT-25297)

---
 .../uast/kotlin/internal/kotlinInternalUastUtils.kt | 13 +++++++++++--
 plugins/uast-kotlin/tests/KotlinUastApiTest.kt      |  4 ++++
 plugins/uast-kotlin/tests/KotlinUastApiTest.kt.172  |  4 ++++
 plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173  |  4 ++++
 4 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinInternalUastUtils.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinInternalUastUtils.kt
index 67e9a84aa9b..fedfe12e5d4 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinInternalUastUtils.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinInternalUastUtils.kt
@@ -46,10 +46,12 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.ArgumentMatch
 import org.jetbrains.kotlin.resolve.constants.IntegerValueTypeConstructor
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedCallableMemberDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
+import org.jetbrains.kotlin.synthetic.SamAdapterExtensionFunctionDescriptor
 import org.jetbrains.kotlin.type.MapPsiToAsmDesc
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.typeUtil.isInterface
@@ -302,8 +304,15 @@ internal fun KotlinULambdaExpression.getFunctionalInterfaceType(): PsiType? {
     if (parent is KtLambdaArgument) run {
         val callExpression = parent.parent as? KtCallExpression ?: return@run
         val resolvedCall = callExpression.getResolvedCall(callExpression.analyze()) ?: return@run
-        val samConstructorDescriptor = resolvedCall.candidateDescriptor as? SamConstructorDescriptor ?: return@run
-        return samConstructorDescriptor.returnType?.getFunctionalInterfaceType(this, psi)
+        val candidateDescriptor = resolvedCall.candidateDescriptor
+        when (candidateDescriptor) {
+            is SamConstructorDescriptor -> return candidateDescriptor.returnType?.getFunctionalInterfaceType(this, psi)
+            is SamAdapterExtensionFunctionDescriptor -> {
+                val index = (resolvedCall.getArgumentMapping(parent) as? ArgumentMatch)?.valueParameter?.index ?: return@run
+                val parameterDescriptor = candidateDescriptor.baseDescriptorForSynthetic.valueParameters.getOrNull(index) ?: return@run
+                return parameterDescriptor.type.getFunctionalInterfaceType(this, psi)
+            }
+        }
     }
     return psi.getExpectedType()?.getFunctionalInterfaceType(this, psi)
 }
diff --git a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt
index 639a2dcb869..285d9b9ea9e 100644
--- a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt
+++ b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt
@@ -371,6 +371,10 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
     fun testUtilsStreamLambda() {
         doTest("Lambdas") { _, file ->
             val lambda = file.findElementByTextFromPsi<ULambdaExpression>("{ it.isEmpty() }")
+            assertEquals(
+                "java.util.function.Predicate<? super java.lang.String>",
+                lambda.functionalInterfaceType?.canonicalText
+            )
             assertEquals(
                 "kotlin.jvm.functions.Function1<? super java.lang.String,? extends java.lang.Boolean>",
                 lambda.getExpressionType()?.canonicalText
diff --git a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.172 b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.172
index cfae9202d30..2003d81ba29 100644
--- a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.172
+++ b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.172
@@ -220,6 +220,10 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
     fun testUtilsStreamLambda() {
         doTest("Lambdas") { _, file ->
             val lambda = file.findElementByTextFromPsi<ULambdaExpression>("{ it.isEmpty() }")
+            assertEquals(
+                "java.util.function.Predicate<? super java.lang.String>",
+                lambda.functionalInterfaceType?.canonicalText
+            )
             assertEquals(
                 "kotlin.jvm.functions.Function1<? super java.lang.String,? extends java.lang.Boolean>",
                 lambda.getExpressionType()?.canonicalText
diff --git a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173 b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173
index a752b425c5f..19857fea93a 100644
--- a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173
+++ b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173
@@ -335,6 +335,10 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
     fun testUtilsStreamLambda() {
         doTest("Lambdas") { _, file ->
             val lambda = file.findElementByTextFromPsi<ULambdaExpression>("{ it.isEmpty() }")
+            assertEquals(
+                "java.util.function.Predicate<? super java.lang.String>",
+                lambda.functionalInterfaceType?.canonicalText
+            )
             assertEquals(
                 "kotlin.jvm.functions.Function1<? super java.lang.String,? extends java.lang.Boolean>",
                 lambda.getExpressionType()?.canonicalText