From 83179f448208ff643830c815ca5bf7a12a639c81 Mon Sep 17 00:00:00 2001 From: Dmitriy Novozhilov Date: Fri, 18 Feb 2022 13:59:51 +0300 Subject: [PATCH] [FIR] Add extension for injecting new implicit receivers after resolution of function call --- .../fir/extensions/FirExtensionRegistrar.kt | 6 ++++ .../FirExpressionResolutionExtension.kt | 28 +++++++++++++++++++ .../FirExpressionsResolveTransformer.kt | 25 +++++++++++++++++ 3 files changed, 59 insertions(+) create mode 100644 compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExpressionResolutionExtension.kt diff --git a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt index 1337d122878..9dd0aadc22f 100644 --- a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt +++ b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt @@ -28,6 +28,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() { FirAdditionalCheckersExtension::class, FirSupertypeGenerationExtension::class, FirTypeAttributeExtension::class, + FirExpressionResolutionExtension::class, ) } @@ -61,6 +62,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() { operator fun ((FirSession) -> FirTypeAttributeExtension).unaryPlus() { registerExtension(FirTypeAttributeExtension::class, FirTypeAttributeExtension.Factory { this.invoke(it) }) } + + @JvmName("plusExpressionResolutionExtension") + operator fun ((FirSession) -> FirExpressionResolutionExtension).unaryPlus() { + registerExtension(FirExpressionResolutionExtension::class, FirExpressionResolutionExtension.Factory { this.invoke(it) }) + } } @OptIn(PluginServicesInitialization::class) diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExpressionResolutionExtension.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExpressionResolutionExtension.kt new file mode 100644 index 00000000000..745c8e405d9 --- /dev/null +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExpressionResolutionExtension.kt @@ -0,0 +1,28 @@ +/* + * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package org.jetbrains.kotlin.fir.extensions + +import org.jetbrains.kotlin.fir.FirSession +import org.jetbrains.kotlin.fir.expressions.FirFunctionCall +import org.jetbrains.kotlin.fir.types.ConeKotlinType +import kotlin.reflect.KClass + +abstract class FirExpressionResolutionExtension(session: FirSession) : FirExtension(session) { + companion object { + val NAME = FirExtensionPointName("ExpressionResolutionExtension") + } + + final override val name: FirExtensionPointName + get() = NAME + + final override val extensionType: KClass = FirExpressionResolutionExtension::class + + abstract fun addNewImplicitReceivers(functionCall: FirFunctionCall): List + + fun interface Factory : FirExtension.Factory +} + +val FirExtensionService.expressionResolutionExtensions: List by FirExtensionService.registeredExtensions() diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt index a3a62b301f5..eceebef36a7 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt @@ -16,6 +16,8 @@ import org.jetbrains.kotlin.fir.expressions.* import org.jetbrains.kotlin.fir.expressions.builder.* import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList import org.jetbrains.kotlin.fir.expressions.impl.toAnnotationArgumentMapping +import org.jetbrains.kotlin.fir.extensions.expressionResolutionExtensions +import org.jetbrains.kotlin.fir.extensions.extensionService import org.jetbrains.kotlin.fir.references.* import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference import org.jetbrains.kotlin.fir.references.builder.buildExplicitSuperReference @@ -32,6 +34,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer import org.jetbrains.kotlin.fir.resolve.transformers.StoreReceiver import org.jetbrains.kotlin.fir.scopes.impl.isWrappedIntegerOperator import org.jetbrains.kotlin.fir.scopes.impl.isWrappedIntegerOperatorForUnsignedType +import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol import org.jetbrains.kotlin.fir.types.* import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef @@ -56,6 +59,8 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform var enableArrayOfCallTransformation = false var containingSafeCallExpression: FirSafeCallExpression? = null + private val expressionResolutionExtensions = session.extensionService.expressionResolutionExtensions.takeIf { it.isNotEmpty() } + init { @Suppress("LeakingThis") components.callResolver.initTransformer(this) @@ -378,6 +383,9 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform } val result = completeInference.transformToIntegerOperatorCallOrApproximateItIfNeeded(data) dataFlowAnalyzer.exitFunctionCall(result, callCompleted) + + addReceiversFromExtensions(result) + if (callCompleted) { if (enableArrayOfCallTransformation) { return arrayOfCallTransformer.transformFunctionCall(result, null) @@ -386,6 +394,23 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform return result } + @OptIn(PrivateForInline::class) + private fun addReceiversFromExtensions(functionCall: FirFunctionCall) { + val extensions = expressionResolutionExtensions ?: return + val boundSymbol = context.containerIfAny?.symbol as? FirCallableSymbol<*> ?: return + for (extension in extensions) { + for (receiverType in extension.addNewImplicitReceivers(functionCall)) { + val receiverValue = ImplicitExtensionReceiverValue( + boundSymbol, + receiverType, + session, + scopeSession + ) + context.addReceiver(name = null, receiverValue) + } + } + } + private fun FirFunctionCall.transformToIntegerOperatorCallOrApproximateItIfNeeded(resolutionMode: ResolutionMode): FirFunctionCall { if (!explicitReceiver.isIntegerLiteralOrOperatorCall()) return this val resolvedSymbol = when (val reference = calleeReference) {