From dc453e6d8b0e6b12a1d3f5cf5be7b155dcbbce90 Mon Sep 17 00:00:00 2001 From: Dmitry Petrov Date: Thu, 8 Jun 2017 10:54:22 +0300 Subject: [PATCH] [NI] Update lambda argument expression type when return type is known --- .../tower/KotlinToResolvedCallTransformer.kt | 42 +++++++++++++------ .../kotlin/builtins/functionTypes.kt | 19 +++++---- 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt index bc97d47b49d..580461659ca 100644 --- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt +++ b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt @@ -16,6 +16,8 @@ package org.jetbrains.kotlin.resolve.calls.tower +import com.intellij.psi.PsiElement +import org.jetbrains.kotlin.builtins.replaceReturnType import org.jetbrains.kotlin.config.LanguageVersionSettings import org.jetbrains.kotlin.descriptors.* import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl @@ -40,9 +42,7 @@ import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue -import org.jetbrains.kotlin.types.ErrorUtils -import org.jetbrains.kotlin.types.KotlinType -import org.jetbrains.kotlin.types.TypeUtils +import org.jetbrains.kotlin.types.* import org.jetbrains.kotlin.types.expressions.DataFlowAnalyzer import java.util.* @@ -162,16 +162,7 @@ class KotlinToResolvedCallTransformer( for (lambdaArgument in lambdaArguments) { val returnType = lambdaArgument.finalReturnType - val psiCallArgument = lambdaArgument.argument.psiCallArgument - val ktFunction = when (psiCallArgument) { - is LambdaKotlinCallArgumentImpl -> psiCallArgument.ktLambdaExpression.functionLiteral - is FunctionExpressionImpl -> psiCallArgument.ktFunction - else -> throw AssertionError("Unexpected psiCallArgument for resolved lambda argument: $psiCallArgument") - } - - val functionDescriptor = trace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? FunctionDescriptorImpl ?: - throw AssertionError("No function descriptor for resolved lambda argument") - functionDescriptor.setReturnType(returnType) + updateTraceForLambdaReturnType(lambdaArgument, trace, returnType) for (lambdaResult in lambdaArgument.resultArguments) { val resultValueArgument = lambdaResult.psiCallArgument @@ -186,6 +177,31 @@ class KotlinToResolvedCallTransformer( } } + private fun updateTraceForLambdaReturnType(lambdaArgument: ResolvedLambdaArgument, trace: BindingTrace, returnType: UnwrappedType) { + val psiCallArgument = lambdaArgument.argument.psiCallArgument + + val ktArgumentExpression: KtExpression + val ktFunction: KtElement + when (psiCallArgument) { + is LambdaKotlinCallArgumentImpl -> { + ktArgumentExpression = psiCallArgument.ktLambdaExpression + ktFunction = ktArgumentExpression.functionLiteral + } + is FunctionExpressionImpl -> { + ktArgumentExpression = psiCallArgument.ktFunction + ktFunction = ktArgumentExpression + } + else -> throw AssertionError("Unexpected psiCallArgument for resolved lambda argument: $psiCallArgument") + } + + val functionDescriptor = trace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? FunctionDescriptorImpl ?: + throw AssertionError("No function descriptor for resolved lambda argument") + functionDescriptor.setReturnType(returnType) + + val existingLambdaType = trace.getType(ktArgumentExpression) ?: throw AssertionError("No type for resolved lambda argument") + trace.recordType(ktArgumentExpression, existingLambdaType.replaceReturnType(returnType)) + } + // todo very beginning code private fun runArgumentsChecks( context: BasicCallResolutionContext, diff --git a/core/descriptors/src/org/jetbrains/kotlin/builtins/functionTypes.kt b/core/descriptors/src/org/jetbrains/kotlin/builtins/functionTypes.kt index 242c1e5adfd..1944870f45b 100644 --- a/core/descriptors/src/org/jetbrains/kotlin/builtins/functionTypes.kt +++ b/core/descriptors/src/org/jetbrains/kotlin/builtins/functionTypes.kt @@ -29,10 +29,7 @@ import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.resolve.constants.ConstantValueFactory import org.jetbrains.kotlin.resolve.constants.StringValue import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe -import org.jetbrains.kotlin.types.KotlinType -import org.jetbrains.kotlin.types.KotlinTypeFactory -import org.jetbrains.kotlin.types.SimpleType -import org.jetbrains.kotlin.types.TypeProjection +import org.jetbrains.kotlin.types.* import org.jetbrains.kotlin.types.typeUtil.asTypeProjection import org.jetbrains.kotlin.types.typeUtil.replaceAnnotations import org.jetbrains.kotlin.utils.DFS @@ -122,21 +119,27 @@ private fun FqNameUnsafe.getFunctionalClassKind(): FunctionClassDescriptor.Kind? fun KotlinType.getReceiverTypeFromFunctionType(): KotlinType? { - assert(isBuiltinFunctionalType) { "Not a function type: ${this}" } + assert(isBuiltinFunctionalType) { "Not a function type: $this" } return if (isTypeAnnotatedWithExtensionFunctionType) arguments.first().type else null } fun KotlinType.getReturnTypeFromFunctionType(): KotlinType { - assert(isBuiltinFunctionalType) { "Not a function type: ${this}" } + assert(isBuiltinFunctionalType) { "Not a function type: $this" } return arguments.last().type } +fun KotlinType.replaceReturnType(newReturnType: KotlinType): KotlinType { + assert(isBuiltinFunctionalType) { "Not a function type: $this"} + val argumentsWithNewReturnType = arguments.toMutableList().apply { set(size - 1, TypeProjectionImpl(newReturnType)) } + return replace(newArguments = argumentsWithNewReturnType) +} + fun KotlinType.getValueParameterTypesFromFunctionType(): List { - assert(isBuiltinFunctionalType) { "Not a function type: ${this}" } + assert(isBuiltinFunctionalType) { "Not a function type: $this" } val arguments = arguments val first = if (isBuiltinExtensionFunctionalType) 1 else 0 val last = arguments.size - 1 - assert(first <= last) { "Not an exact function type: ${this}" } + assert(first <= last) { "Not an exact function type: $this" } return arguments.subList(first, last) }