Create from Usage: Infer expected type for lambda body

#KT-18186 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/typeUtils.kt

@@ -17,11 +17,9 @@
 package org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder
 
 import com.intellij.refactoring.psi.SearchUtils
+import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.cfg.pseudocode.*
-import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
+import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.descriptors.ModuleDescriptor
-import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
-import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.idea.project.builtIns
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
@@ -41,6 +39,7 @@ import org.jetbrains.kotlin.resolve.scopes.utils.findClassifier
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
+import org.jetbrains.kotlin.types.typeUtil.supertypes
 import java.util.*
 
 internal operator fun KotlinType.contains(inner: KotlinType): Boolean {
@@ -117,8 +116,11 @@ fun KtExpression.guessTypes(
         context: BindingContext,
         module: ModuleDescriptor,
         pseudocode: Pseudocode? = null,
-        coerceUnusedToUnit: Boolean = true
+        coerceUnusedToUnit: Boolean = true,
+        allowErrorTypes: Boolean = false
 ): Array<KotlinType> {
+    fun isAcceptable(type: KotlinType) = allowErrorTypes || !ErrorUtils.containsErrorType(type)
+
     if (coerceUnusedToUnit
         && this !is KtDeclaration
         && isUsedAsStatement(context)
@@ -126,7 +128,7 @@ fun KtExpression.guessTypes(
 
     // if we know the actual type of the expression
     val theType1 = context.getType(this)
-    if (theType1 != null) {
+    if (theType1 != null && isAcceptable(theType1)) {
         val dataFlowInfo = context.getDataFlowInfoAfter(this)
         val possibleTypes = dataFlowInfo.getCollectedTypes(DataFlowValueFactory.createDataFlowValue(this, theType1, context, module))
         return if (possibleTypes.isNotEmpty()) possibleTypes.toTypedArray() else arrayOf(theType1)
@@ -134,7 +136,7 @@ fun KtExpression.guessTypes(
 
     // expression has an expected type
     val theType2 = context[BindingContext.EXPECTED_EXPRESSION_TYPE, this]
-    if (theType2 != null) return arrayOf(theType2)
+    if (theType2 != null && isAcceptable(theType2)) return arrayOf(theType2)
 
     val parent = parent
     return when {
@@ -198,6 +200,22 @@ fun KtExpression.guessTypes(
         parent is KtStringTemplateEntryWithExpression && parent.expression == this -> {
             arrayOf(module.builtIns.stringType)
         }
+        parent is KtBlockExpression && parent.statements.lastOrNull() == this && parent.parent is KtFunctionLiteral -> {
+            parent.guessTypes(context, module, pseudocode, coerceUnusedToUnit)
+        }
+        parent is KtFunction -> {
+            val functionDescriptor = context[BindingContext.DECLARATION_TO_DESCRIPTOR, parent] as? FunctionDescriptor ?: return arrayOf()
+            val returnType = functionDescriptor.returnType
+            if (returnType != null && isAcceptable(returnType)) return arrayOf(returnType)
+            val functionalExpression: KtExpression? = when {
+                parent is KtFunctionLiteral -> parent.parent as? KtLambdaExpression
+                parent is KtNamedFunction && parent.name == null -> parent
+                else -> null
+            }
+            if (functionalExpression == null) return arrayOf()
+            val lambdaTypes = functionalExpression.guessTypes(context, module, pseudocode?.parent, coerceUnusedToUnit)
+            lambdaTypes.mapNotNull { it.getFunctionType()?.arguments?.lastOrNull()?.type }.toTypedArray()
+        }
         else -> {
             pseudocode?.getElementValue(this)?.let {
                 getExpectedTypePredicate(it, context, module.builtIns).getRepresentativeTypes().toTypedArray()
@@ -206,6 +224,8 @@ fun KtExpression.guessTypes(
     }
 }
 
+private fun KotlinType.getFunctionType() = if (isFunctionType) this else supertypes().firstOrNull { it.isFunctionType }
+
 private fun KtNamedDeclaration.guessType(context: BindingContext): Array<KotlinType> {
     val expectedTypes = SearchUtils.findAllReferences(this, useScope)!!.mapNotNullTo(HashSet<KotlinType>()) { ref ->
         if (ref is KtSimpleNameReference) {

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createCallable/CreateCallableFromCallActionFactory.kt

@@ -19,12 +19,13 @@ package org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiClass
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
-import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.caches.resolve.analyzeAndGetResult
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder.*
@@ -55,7 +56,7 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
 ) : CreateCallableMemberFromUsageFactory<E>(extensionsEnabled) {
     protected abstract fun doCreateCallableInfo(
             expression: E,
-            context: BindingContext,
+            analysisResult: AnalysisResult,
             name: String,
             receiverType: TypeInfo,
             possibleContainers: List<KtElement>
@@ -95,9 +96,9 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
 
         if (calleeExpr.getReferencedNameElementType() != KtTokens.IDENTIFIER) return null
 
-        val context = calleeExpr.analyze()
+        val analysisResult = calleeExpr.analyzeAndGetResult()
-        val receiver = element.getCall(context)?.explicitReceiver
+        val receiver = element.getCall(analysisResult.bindingContext)?.explicitReceiver
-        val receiverType = getReceiverTypeInfo(context, project, receiver) ?: return null
+        val receiverType = getReceiverTypeInfo(analysisResult.bindingContext, project, receiver) ?: return null
 
         val possibleContainers =
                 if (receiverType is TypeInfo.Empty) {
@@ -108,7 +109,7 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
                 }
                 else Collections.emptyList()
 
-        return doCreateCallableInfo(element, context, calleeExpr.getReferencedName(), receiverType, possibleContainers)
+        return doCreateCallableInfo(element, analysisResult, calleeExpr.getReferencedName(), receiverType, possibleContainers)
     }
 
     private fun getReceiverTypeInfo(context: BindingContext, project: Project, receiver: Receiver?): TypeInfo? {
@@ -179,7 +180,7 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
 
         override fun doCreateCallableInfo(
                 expression: KtSimpleNameExpression,
-                context: BindingContext,
+                analysisResult: AnalysisResult,
                 name: String,
                 receiverType: TypeInfo,
                 possibleContainers: List<KtElement>
@@ -196,14 +197,14 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
         object Default : Property() {
             override fun doCreateCallableInfo(
                     expression: KtSimpleNameExpression,
-                    context: BindingContext,
+                    analysisResult: AnalysisResult,
                     name: String,
                     receiverType: TypeInfo,
                     possibleContainers: List<KtElement>
             ): CallableInfo? {
                 return super.doCreateCallableInfo(
                         expression,
-                        context,
+                        analysisResult,
                         name,
                         receiverType,
                         possibleContainers.filterNot { it is KtClassBody && (it.parent as KtClassOrObject).isInterfaceClass() }
@@ -214,11 +215,11 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
         object Abstract : Property() {
             override fun doCreateCallableInfo(
                     expression: KtSimpleNameExpression,
-                    context: BindingContext,
+                    analysisResult: AnalysisResult,
                     name: String,
                     receiverType: TypeInfo,
                     possibleContainers: List<KtElement>
-            ) = super.doCreateCallableInfo(expression, context, name, receiverType, possibleContainers)?.let {
+            ) = super.doCreateCallableInfo(expression, analysisResult, name, receiverType, possibleContainers)?.let {
                 getAbstractCallableInfo(it, expression)
             }
         }
@@ -226,12 +227,12 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
         object ByImplicitExtensionReceiver : Property() {
             override fun doCreateCallableInfo(
                     expression: KtSimpleNameExpression,
-                    context: BindingContext,
+                    analysisResult: AnalysisResult,
                     name: String,
                     receiverType: TypeInfo,
                     possibleContainers: List<KtElement>
-            ) = super.doCreateCallableInfo(expression, context, name, receiverType, possibleContainers)?.let {
+            ) = super.doCreateCallableInfo(expression, analysisResult, name, receiverType, possibleContainers)?.let {
-                ByImplicitExtensionReceiver.getCallableWithReceiverInsideExtension(it, expression, context, receiverType)
+                ByImplicitExtensionReceiver.getCallableWithReceiverInsideExtension(it, expression, analysisResult.bindingContext, receiverType)
             }
         }
     }
@@ -243,14 +244,20 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
 
         override fun doCreateCallableInfo(
                 expression: KtCallExpression,
-                context: BindingContext,
+                analysisResult: AnalysisResult,
                 name: String,
                 receiverType: TypeInfo,
                 possibleContainers: List<KtElement>
         ): CallableInfo? {
             val parameters = expression.getParameterInfos()
             val typeParameters = expression.getTypeInfoForTypeArguments()
-            val returnType = TypeInfo(expression.getQualifiedExpressionForSelectorOrThis(), Variance.OUT_VARIANCE)
+            val fullCallExpression = expression.getQualifiedExpressionForSelectorOrThis()
+            val expectedType = fullCallExpression.guessTypes(analysisResult.bindingContext, analysisResult.moduleDescriptor).singleOrNull()
+            val returnType = if (expectedType != null) {
+                TypeInfo(expectedType, Variance.OUT_VARIANCE)
+            } else {
+                TypeInfo(fullCallExpression, Variance.OUT_VARIANCE)
+            }
             return FunctionInfo(name, receiverType, returnType, possibleContainers, parameters, typeParameters)
         }
 
@@ -259,11 +266,11 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
         object Abstract : Function() {
             override fun doCreateCallableInfo(
                     expression: KtCallExpression,
-                    context: BindingContext,
+                    analysisResult: AnalysisResult,
                     name: String,
                     receiverType: TypeInfo,
                     possibleContainers: List<KtElement>
-            ) = super.doCreateCallableInfo(expression, context, name, receiverType, possibleContainers)?.let {
+            ) = super.doCreateCallableInfo(expression, analysisResult, name, receiverType, possibleContainers)?.let {
                 getAbstractCallableInfo(it, expression)
             }
         }
@@ -271,12 +278,12 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
         object ByImplicitExtensionReceiver : Function() {
             override fun doCreateCallableInfo(
                     expression: KtCallExpression,
-                    context: BindingContext,
+                    analysisResult: AnalysisResult,
                     name: String,
                     receiverType: TypeInfo,
                     possibleContainers: List<KtElement>
-            ) = super.doCreateCallableInfo(expression, context, name, receiverType, possibleContainers)?.let {
+            ) = super.doCreateCallableInfo(expression, analysisResult, name, receiverType, possibleContainers)?.let {
-                getCallableWithReceiverInsideExtension(it, expression, context, receiverType)
+                getCallableWithReceiverInsideExtension(it, expression, analysisResult.bindingContext, receiverType)
             }
         }
     }
@@ -288,7 +295,7 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
 
         override fun doCreateCallableInfo(
                 expression: KtCallExpression,
-                context: BindingContext,
+                analysisResult: AnalysisResult,
                 name: String,
                 receiverType: TypeInfo,
                 possibleContainers: List<KtElement>
@@ -297,14 +304,14 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
 
             val classDescriptor = expression
                     .calleeExpression
-                    ?.getReferenceTargets(context)
+                    ?.getReferenceTargets(analysisResult.bindingContext)
                     ?.mapNotNull { (it as? ConstructorDescriptor)?.containingDeclaration }
                     ?.distinct()
                     ?.singleOrNull() as? ClassDescriptor
             val klass = classDescriptor?.source?.getPsi()
             if ((klass !is KtClass && klass !is PsiClass) || !klass.canRefactor()) return null
 
-            val expectedType = context[BindingContext.EXPECTED_EXPRESSION_TYPE, expression.getQualifiedExpressionForSelectorOrThis()]
+            val expectedType = analysisResult.bindingContext[BindingContext.EXPECTED_EXPRESSION_TYPE, expression.getQualifiedExpressionForSelectorOrThis()]
                                ?: classDescriptor.builtIns.nullableAnyType
             if (!classDescriptor.defaultType.isSubtypeOf(expectedType)) return null
 

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createCallable/CreateFunctionFromCallableReferenceActionFactory.kt

@@ -45,6 +45,7 @@ object CreateFunctionFromCallableReferenceActionFactory : CreateCallableMemberFr
         val context = resolutionFacade.analyze(element, BodyResolveMode.PARTIAL)
         return element
                 .guessTypes(context, resolutionFacade.moduleDescriptor)
+                .ifEmpty { element.guessTypes(context, resolutionFacade.moduleDescriptor, allowErrorTypes = true) } // approximate with Any
                 .filter(KotlinType::isFunctionType)
                 .mapNotNull {
                     val expectedReceiverType = it.getReceiverTypeFromFunctionType()

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createVariable/CreateParameterFromUsageFix.kt

@@ -97,8 +97,8 @@ open class CreateParameterFromUsageFix<E : KtElement>(
             if (!receiverClass.canRefactor()) return null
             val constructorDescriptor = receiverClassDescriptor.unsubstitutedPrimaryConstructor ?: return null
 
-            val paramType = info.returnTypeInfo.getPossibleTypes(builder).firstOrNull() ?: return null
+            val paramType = info.returnTypeInfo.getPossibleTypes(builder).firstOrNull()
-            if (paramType.hasTypeParametersToAdd(constructorDescriptor, builder.currentFileContext)) return null
+            if (paramType != null && paramType.hasTypeParametersToAdd(constructorDescriptor, builder.currentFileContext)) return null
 
             val paramInfo = KotlinParameterInfo(
                     callableDescriptor = constructorDescriptor,

idea/src/org/jetbrains/kotlin/idea/refactoring/kotlinRefactoringUtil.kt

@@ -194,7 +194,7 @@ fun PsiElement.getExtractionContainers(strict: Boolean = true, includeAll: Boole
     fun getEnclosingDeclaration(element: PsiElement, strict: Boolean): PsiElement? {
         return (if (strict) element.parents else element.parentsWithSelf)
                 .filter {
-                    (it is KtDeclarationWithBody && it !is KtFunctionLiteral)
+                    (it is KtDeclarationWithBody && it !is KtFunctionLiteral && !(it is KtNamedFunction && it.name == null))
                     || it is KtAnonymousInitializer
                     || it is KtClassBody
                     || it is KtFile

idea/testData/quickfix/createFromUsage/createFunction/call/inAnonymousFunBlockWithExpectedType.kt

@@ -0,0 +1,9 @@
+// "Create function 'bar'" "true"
+
+class A<T>(val t: T)
+
+fun <T, U> A<T>.convert(f: (T) -> U) = A(f(t))
+
+fun foo(l: A<String>): A<Int> {
+    return l.convert(fun(it: String): Int { return <caret>bar(it) })
+}

idea/testData/quickfix/createFromUsage/createFunction/call/inAnonymousFunBlockWithExpectedType.kt.after

@@ -0,0 +1,13 @@
+// "Create function 'bar'" "true"
+
+class A<T>(val t: T)
+
+fun <T, U> A<T>.convert(f: (T) -> U) = A(f(t))
+
+fun foo(l: A<String>): A<Int> {
+    return l.convert(fun(it: String): Int { return bar(it) })
+}
+
+fun bar(it: String): Int {
+    TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+}

idea/testData/quickfix/createFromUsage/createFunction/call/inAnonymousFunExprWithExpectedType.kt

@@ -0,0 +1,9 @@
+// "Create function 'bar'" "true"
+
+class A<T>(val t: T)
+
+fun <T, U> A<T>.convert(f: (T) -> U) = A(f(t))
+
+fun foo(l: A<String>): A<Int> {
+    return l.convert(fun(it: String) = <caret>bar(it))
+}

idea/testData/quickfix/createFromUsage/createFunction/call/inAnonymousFunExprWithExpectedType.kt.after

@@ -0,0 +1,13 @@
+// "Create function 'bar'" "true"
+
+class A<T>(val t: T)
+
+fun <T, U> A<T>.convert(f: (T) -> U) = A(f(t))
+
+fun foo(l: A<String>): A<Int> {
+    return l.convert(fun(it: String) = bar(it))
+}
+
+fun bar(it: String): Int {
+    TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+}

idea/testData/quickfix/createFromUsage/createFunction/call/inLambdaWithExpectedType.kt

@@ -0,0 +1,9 @@
+// "Create function 'bar'" "true"
+
+class A<T>(val t: T)
+
+fun <T, U> A<T>.convert(f: (T) -> U) = A(f(t))
+
+fun foo(l: A<String>): A<Int> {
+    return l.convert { <caret>bar(it) }
+}

idea/testData/quickfix/createFromUsage/createFunction/call/inLambdaWithExpectedType.kt.after

@@ -0,0 +1,13 @@
+// "Create function 'bar'" "true"
+
+class A<T>(val t: T)
+
+fun <T, U> A<T>.convert(f: (T) -> U) = A(f(t))
+
+fun foo(l: A<String>): A<Int> {
+    return l.convert { bar(it) }
+}
+
+fun bar(it: String): Int {
+    TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -2577,12 +2577,30 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("inAnonymousFunBlockWithExpectedType.kt")
+                public void testInAnonymousFunBlockWithExpectedType() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/call/inAnonymousFunBlockWithExpectedType.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("inAnonymousFunExprWithExpectedType.kt")
+                public void testInAnonymousFunExprWithExpectedType() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/call/inAnonymousFunExprWithExpectedType.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("inLambda.kt")
                 public void testInLambda() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/call/inLambda.kt");
                     doTest(fileName);
                 }
 
+                @TestMetadata("inLambdaWithExpectedType.kt")
+                public void testInLambdaWithExpectedType() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/call/inLambdaWithExpectedType.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("inconsistentTypes.kt")
                 public void testInconsistentTypes() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/call/inconsistentTypes.kt");