Fix UAST tests for new-inference: adapt SAM-converted descriptors

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUSimpleReferenceExpression.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.ImportedFromObjectCallableDescriptor
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedCallableMemberDescriptor
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
@@ -226,7 +227,7 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
         get() = (resolved as? PsiNamedElement)?.name
 
     private val resolved by lazy {
-        when (val resultingDescriptor = sourcePsi.getResolvedCall(sourcePsi.analyze())?.resultingDescriptor) {
+        when (val resultingDescriptor = sourcePsi.getResolvedCall(sourcePsi.analyze())?.descriptorForResolveViaConstructor()) {
             is ConstructorDescriptor -> {
                 resultingDescriptor.constructedClass.toSource()?.getMaybeLightElement()
                     ?: (resultingDescriptor as? DeserializedCallableMemberDescriptor)?.let { resolveContainingDeserializedClass(sourcePsi, it) }
@@ -240,6 +241,11 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
     override fun resolve(): PsiElement? = resolved
 
     override fun asLogString(): String = log<USimpleNameReferenceExpression>("identifier = $identifier, resolvesTo = $resolvedName")
+
+    // In new inference, SAM constructor is substituted with a function descriptor, so we use candidate descriptor to preserve behavior
+    private fun ResolvedCall<*>.descriptorForResolveViaConstructor(): CallableDescriptor? {
+        return if (this is NewResolvedCallImpl) candidateDescriptor else resultingDescriptor
+    }
 }
 
 class KotlinStringUSimpleReferenceExpression(

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUSimpleReferenceExpression.kt.183

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.ImportedFromObjectCallableDescriptor
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedCallableMemberDescriptor
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
@@ -226,7 +227,7 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
         get() = (resolved as? PsiNamedElement)?.name
 
     private val resolved by lazy {
-        when (val resultingDescriptor = sourcePsi.getResolvedCall(sourcePsi.analyze())?.resultingDescriptor) {
+        when (val resultingDescriptor = sourcePsi.getResolvedCall(sourcePsi.analyze())?.descriptorForResolveViaConstructor()) {
             is ConstructorDescriptor -> {
                 resultingDescriptor.constructedClass.toSource()?.getMaybeLightElement()
                     ?: (resultingDescriptor as? DeserializedCallableMemberDescriptor)?.let { resolveContainingDeserializedClass(sourcePsi, it) }
@@ -240,6 +241,11 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
     override fun resolve(): PsiElement? = resolved
 
     override fun asLogString(): String = log<USimpleNameReferenceExpression>("identifier = $identifier, resolvesTo = $resolvedName")
+
+    // In new inference, SAM constructor is substituted with a function descriptor, so we use candidate descriptor to preserve behavior
+    private fun ResolvedCall<*>.descriptorForResolveViaConstructor(): CallableDescriptor? {
+        return if (this is NewResolvedCallImpl) candidateDescriptor else resultingDescriptor
+    }
 }
 
 class KotlinStringUSimpleReferenceExpression(

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinInternalUastUtils.kt

@@ -51,6 +51,8 @@ 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.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.constants.IntegerLiteralTypeConstructor
 import org.jetbrains.kotlin.resolve.constants.IntegerValueTypeConstructor
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
@@ -386,20 +388,43 @@ internal fun KotlinULambdaExpression.getFunctionalInterfaceType(): PsiType? {
     if (parent is KtValueArgument) run {
         val callExpression = parent.parents.take(2).firstIsInstanceOrNull<KtCallExpression>() ?: return@run
         val resolvedCall = callExpression.getResolvedCall(callExpression.analyze()) ?: return@run
+
+        // NewResolvedCallImpl can be used as a marker meaning that this code is working under *new* inference
+        if (resolvedCall is NewResolvedCallImpl) {
+            val samConvertedArgument = resolvedCall.getExpectedTypeForSamConvertedArgument(parent)
+
+            // Same as if in old inference we would get SamDescriptor
+            if (samConvertedArgument != null) {
+                val type = getTypeByArgument(resolvedCall, resolvedCall.candidateDescriptor, parent) ?: return@run
+                return type.getFunctionalInterfaceType(this, sourcePsi)
+            }
+        }
+
         val candidateDescriptor = resolvedCall.candidateDescriptor as? SyntheticMemberDescriptor<*> ?: return@run
         when (candidateDescriptor) {
             is SamConstructorDescriptor -> return candidateDescriptor.returnType?.getFunctionalInterfaceType(this, sourcePsi)
             is SamAdapterDescriptor<*>, is SamAdapterExtensionFunctionDescriptor -> {
-                val index = (resolvedCall.getArgumentMapping(parent) as? ArgumentMatch)?.valueParameter?.index ?: return@run
                 val functionDescriptor = candidateDescriptor.baseDescriptorForSynthetic as? FunctionDescriptor ?: return@run
-                val parameterDescriptor = functionDescriptor.valueParameters.getOrNull(index) ?: return@run
+
-                return parameterDescriptor.type.getFunctionalInterfaceType(this, sourcePsi)
+                val type = getTypeByArgument(resolvedCall, functionDescriptor, parent) ?: return@run
+                return type.getFunctionalInterfaceType(this, sourcePsi)
             }
         }
     }
     return sourcePsi.getExpectedType()?.getFunctionalInterfaceType(this, sourcePsi)
 }
 
+private fun getTypeByArgument(
+    resolvedCall: ResolvedCall<*>,
+    descriptor: CallableDescriptor,
+    argument: ValueArgument
+): KotlinType? {
+    val index = (resolvedCall.getArgumentMapping(argument) as? ArgumentMatch)?.valueParameter?.index ?: return null
+    val parameterDescriptor = descriptor.valueParameters.getOrNull(index) ?: return null
+
+    return parameterDescriptor.type
+}
+
 internal fun unwrapFakeFileForLightClass(file: PsiFile): PsiFile = (file as? FakeFileForLightClass)?.ktFile ?: file
 
 // mb merge with org.jetbrains.kotlin.idea.references.ReferenceAccess ?