From 4cf2bf6ca59d05558e973467e6601645b67afaa3 Mon Sep 17 00:00:00 2001
From: Vyacheslav Gerasimov <vyacheslav.gerasimov@jetbrains.com>
Date: Fri, 14 Jul 2017 16:12:52 +0300
Subject: [PATCH] UAST: Implement getFunctionalInterfaceType for
 KotlinULambdaExpression

---
 .../expressions/KotlinULambdaExpression.kt     |  6 +++++-
 .../uast/kotlin/expressions/LocalFunction.kt   |  3 +++
 .../kotlin/internal/kotlinInternalUastUtils.kt | 18 ++++++++++++++++++
 plugins/uast-kotlin/testData/SAM.kt            | 11 +++++++++++
 4 files changed, 37 insertions(+), 1 deletion(-)
 create mode 100644 plugins/uast-kotlin/testData/SAM.kt

diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULambdaExpression.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULambdaExpression.kt
index b20df3eb2b0..6ec088a76bb 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULambdaExpression.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULambdaExpression.kt
@@ -16,6 +16,7 @@
 
 package org.jetbrains.uast.kotlin
 
+import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.uast.UBlockExpression
 import org.jetbrains.uast.UElement
@@ -27,6 +28,9 @@ class KotlinULambdaExpression(
         override val psi: KtLambdaExpression,
         override val uastParent: UElement?
 ) : KotlinAbstractUExpression(), ULambdaExpression, KotlinUElementWithType {
+    val functionalInterfaceType: PsiType?
+        get() = getFunctionalInterfaceType()
+
     override val body by lz { KotlinConverter.convertOrEmpty(psi.bodyExpression, this) }
     
     override val valueParameters by lz {
@@ -43,6 +47,6 @@ class KotlinULambdaExpression(
         val expressions = (body as? UBlockExpression)?.expressions
                                   ?.joinToString("\n") { it.asRenderString().withMargin } ?: body.asRenderString()
 
-        return "{ " + renderedValueParameters + "\n" + expressions + "\n}"
+        return "{ $renderedValueParameters\n$expressions\n}"
     }
 }
\ No newline at end of file
diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt
index 2fedb9d330b..f82533c6b93 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.uast.kotlin.expressions
 
+import com.intellij.psi.PsiType
 import com.intellij.psi.PsiVariable
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.uast.*
@@ -29,6 +30,8 @@ private class KotlinLocalFunctionULambdaExpression(
         override val psi: KtFunction,
         override val uastParent: UElement?
 ): KotlinAbstractUExpression(), ULambdaExpression {
+    val functionalInterfaceType: PsiType?
+        get() = null
 
     override val body by lz { KotlinConverter.convertOrEmpty(psi.bodyExpression, this) }
 
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 76a3abd39e8..04c0221e9d6 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
@@ -30,11 +30,13 @@ import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.PsiTypesUtil
 import org.jetbrains.kotlin.asJava.toLightClass
 import org.jetbrains.kotlin.asJava.toLightElements
+import org.jetbrains.kotlin.builtins.isBuiltinFunctionalTypeOrSubtype
 import org.jetbrains.kotlin.codegen.signature.BothSignatureWriter
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -43,6 +45,7 @@ import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.isError
+import org.jetbrains.kotlin.types.typeUtil.isInterface
 import org.jetbrains.uast.*
 import java.lang.ref.WeakReference
 import java.text.StringCharacterIterator
@@ -180,3 +183,18 @@ internal inline fun <reified T : UDeclaration, reified P : PsiElement> unwrap(el
     assert(unwrapped !is UElement)
     return unwrapped as P
 }
+
+internal fun KtExpression.getExpectedType(): KotlinType? = analyze()[BindingContext.EXPECTED_EXPRESSION_TYPE, this]
+
+internal fun KtTypeReference.getType(): KotlinType? = analyze()[BindingContext.TYPE, this]
+
+internal fun KotlinType.getFunctionalInterfaceType(source: UElement, element: KtElement): PsiType? =
+        takeIf { it.isInterface() && !it.isBuiltinFunctionalTypeOrSubtype }?.toPsiType(source, element, false)
+
+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)
+    }
+}
diff --git a/plugins/uast-kotlin/testData/SAM.kt b/plugins/uast-kotlin/testData/SAM.kt
new file mode 100644
index 00000000000..b794f2c9286
--- /dev/null
+++ b/plugins/uast-kotlin/testData/SAM.kt
@@ -0,0 +1,11 @@
+
+val notSam = { /* Not SAM */ }
+var foo: java.lang.Runnable = {/* Variable */}
+fun bar(): java.lang.Runnable {
+    foo = {/* Assignment */}
+    val a = {/* Type Cast */} as java.lang.Runnable
+    runRunnable {/* Argument */}
+    return {/* Return */}
+}
+
+fun runRunnable(r: java.lang.Runnable) = r()
\ No newline at end of file