From 7df289746c78f50f4ab29d10153d318cc0c6323d Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Thu, 10 Dec 2020 15:05:47 -0800
Subject: [PATCH] FIR: fix invoke lookup for SAM resolution

---
 .../kotlin/fir/resolve/calls/Arguments.kt     | 32 +++++++++++++++++--
 .../tests/j+k/sam/recursiveSamsAndInvoke.kt   |  2 +-
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt
index adcb10f28d2..511470096ca 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt
@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.addSubtypeConstraintIfCompatible
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
+import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.model.CaptureStatus
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
@@ -409,8 +410,35 @@ fun FirExpression.isFunctional(
             val invokeSymbol =
                 coneType.findContributedInvokeSymbol(
                     session, scopeSession, classLikeExpectedFunctionType, shouldCalculateReturnTypesOfFakeOverrides = false
-                )
-            return invokeSymbol != null
+                ) ?: return false
+            // Make sure the contributed `invoke` is indeed a wanted functional type by checking if types are compatible.
+            val expectedReturnType = classLikeExpectedFunctionType.returnType(session)!!.lowerBoundIfFlexible()
+            val returnTypeCompatible =
+                expectedReturnType is ConeTypeParameterType ||
+                        AbstractTypeChecker.isSubtypeOf(
+                            session.inferenceComponents.ctx,
+                            invokeSymbol.fir.returnTypeRef.coneType,
+                            expectedReturnType,
+                            isFromNullabilityConstraint = false
+                        )
+            if (!returnTypeCompatible) {
+                return false
+            }
+            if (invokeSymbol.fir.valueParameters.size != classLikeExpectedFunctionType.typeArguments.size - 1) {
+                return false
+            }
+            val parameterPairs =
+                invokeSymbol.fir.valueParameters.zip(classLikeExpectedFunctionType.valueParameterTypesIncludingReceiver(session))
+            return parameterPairs.all { (invokeParameter, expectedParameter) ->
+                val expectedParameterType = expectedParameter!!.lowerBoundIfFlexible()
+                expectedParameterType is ConeTypeParameterType ||
+                        AbstractTypeChecker.isSubtypeOf(
+                            session.inferenceComponents.ctx,
+                            invokeParameter.returnTypeRef.coneType,
+                            expectedParameterType,
+                            isFromNullabilityConstraint = false
+                        )
+            }
         }
     }
 }
diff --git a/compiler/testData/diagnostics/tests/j+k/sam/recursiveSamsAndInvoke.kt b/compiler/testData/diagnostics/tests/j+k/sam/recursiveSamsAndInvoke.kt
index 37fcf10fcc9..f4d975f42ef 100644
--- a/compiler/testData/diagnostics/tests/j+k/sam/recursiveSamsAndInvoke.kt
+++ b/compiler/testData/diagnostics/tests/j+k/sam/recursiveSamsAndInvoke.kt
@@ -16,7 +16,7 @@ public interface MyListener<F extends MyFuture<?>>  {
 typealias Handler = (cause: Throwable?) -> Unit
 
 fun <T> MyFuture<T>.setup() {
-    <!INAPPLICABLE_CANDIDATE!>addListener<!>(ListenerImpl<T, MyFuture<T>>())
+    addListener(ListenerImpl<T, MyFuture<T>>())
     addListener { }
 }