Uast: ULambdaExpression returns proper functionalInterfaceType for SamAdapterExtensionFunctionDescriptor (KT-25297)

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
+        val candidateDescriptor = resolvedCall.candidateDescriptor
-        return samConstructorDescriptor.returnType?.getFunctionalInterfaceType(this, psi)
+        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)
 }

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

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

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