From 2c7687195ca0175a528aa88902766b466fb8e8f8 Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <Nicolay.Mitropolsky@jetbrains.com>
Date: Sat, 7 Jul 2018 00:36:03 +0300
Subject: [PATCH] Uast: properly handling explicit SAM (KT-25297)

---
 .../uast/kotlin/internal/kotlinInternalUastUtils.kt   | 11 ++++++++---
 plugins/uast-kotlin/testData/SAM.kt                   |  2 ++
 plugins/uast-kotlin/tests/KotlinUastApiTest.kt        |  5 +++++
 plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173    |  5 +++++
 4 files changed, 20 insertions(+), 3 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 fb33a27cec5..67e9a84aa9b 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
@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.codegen.signature.BothSignatureWriter
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.project.languageVersionSettings
 import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaPackageFragment
+import org.jetbrains.kotlin.load.java.sam.SamConstructorDescriptor
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinaryPackageSourceElement
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.metadata.ProtoBuf
@@ -297,10 +298,14 @@ internal fun KotlinType.getFunctionalInterfaceType(source: UElement, element: Kt
 
 internal fun KotlinULambdaExpression.getFunctionalInterfaceType(): PsiType? {
     val parent = psi.parent
-    return when (parent) {
-        is KtBinaryExpressionWithTypeRHS -> parent.right?.getType()?.getFunctionalInterfaceType(this, psi)
-        else -> psi.getExpectedType()?.getFunctionalInterfaceType(this, psi)
+    if (parent is KtBinaryExpressionWithTypeRHS) return parent.right?.getType()?.getFunctionalInterfaceType(this, psi)
+    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)
     }
+    return psi.getExpectedType()?.getFunctionalInterfaceType(this, psi)
 }
 
 internal fun unwrapFakeFileForLightClass(file: PsiFile): PsiFile = (file as? FakeFileForLightClass)?.ktFile ?: file
diff --git a/plugins/uast-kotlin/testData/SAM.kt b/plugins/uast-kotlin/testData/SAM.kt
index b794f2c9286..e390962292a 100644
--- a/plugins/uast-kotlin/testData/SAM.kt
+++ b/plugins/uast-kotlin/testData/SAM.kt
@@ -8,4 +8,6 @@ fun bar(): java.lang.Runnable {
     return {/* Return */}
 }
 
+val baz = java.lang.Runnable { /* SAM */ }
+
 fun runRunnable(r: java.lang.Runnable) = r()
\ No newline at end of file
diff --git a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt
index d6196ce62f7..639a2dcb869 100644
--- a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt
+++ b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt
@@ -98,6 +98,11 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
 
             assertEquals("java.lang.Runnable",
                          file.findElementByText<ULambdaExpression>("{/* Return */}").functionalInterfaceType?.canonicalText)
+
+            assertEquals(
+                "java.lang.Runnable",
+                file.findElementByText<ULambdaExpression>("{ /* SAM */ }").functionalInterfaceType?.canonicalText
+            )
         }
     }
 
diff --git a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173 b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173
index 8688555176d..a752b425c5f 100644
--- a/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173
+++ b/plugins/uast-kotlin/tests/KotlinUastApiTest.kt.173
@@ -99,6 +99,11 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
 
             assertEquals("java.lang.Runnable",
                          file.findElementByText<ULambdaExpression>("{/* Return */}").functionalInterfaceType?.canonicalText)
+
+            assertEquals(
+                "java.lang.Runnable",
+                file.findElementByText<ULambdaExpression>("{ /* SAM */ }").functionalInterfaceType?.canonicalText
+            )
         }
     }