From 90a4f6002e1ab48f6d3eda6e223aebc625d97fe5 Mon Sep 17 00:00:00 2001 From: Ilya Chernikov Date: Tue, 6 Dec 2022 15:53:31 +0100 Subject: [PATCH] K2 Scripting: support sam-with-receiver in scripting plugin --- .../experimental/jvmhost/test/FeaturesTest.kt | 4 +- .../plugin/impl/compilationContext.kt | 59 +++++++++++++++++-- .../compiler/plugin/pluginRegisrar.kt | 2 + .../FirScriptConfigurationExtensionImpl.kt | 11 +++- 4 files changed, 68 insertions(+), 8 deletions(-) diff --git a/libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/FeaturesTest.kt b/libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/FeaturesTest.kt index b06add30868..988fe265ef2 100644 --- a/libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/FeaturesTest.kt +++ b/libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/FeaturesTest.kt @@ -37,8 +37,8 @@ class FeaturesTest : TestCase() { when (res) { is ResultWithDiagnostics.Success -> fail("Expecting \"Unresolved reference\" error, got successful compilation") is ResultWithDiagnostics.Failure -> - if (res.reports.none { it.message.contains("Unresolved reference") }) { - fail("Expecting \"Unresolved reference\" error, got:\n ${res.reports.joinToString("\n ")}") + if (res.reports.none { it.message.contains("Unresolved reference") || it.message.contains("'this' is not defined in this context") }) { + fail("Expecting \"Unresolved reference\" or \"'this' is not defined in this context\" error, got:\n ${res.reports.joinToString("\n ")}") } } } diff --git a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/compilationContext.kt b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/compilationContext.kt index 9c330b9a13a..eb18d9f1f71 100644 --- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/compilationContext.kt +++ b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/compilationContext.kt @@ -19,9 +19,9 @@ import org.jetbrains.kotlin.cli.common.setupCommonArguments import org.jetbrains.kotlin.cli.jvm.* import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment -import org.jetbrains.kotlin.cli.jvm.config.configureJdkClasspathRoots import org.jetbrains.kotlin.cli.jvm.config.JvmClasspathRoot import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots +import org.jetbrains.kotlin.cli.jvm.config.configureJdkClasspathRoots import org.jetbrains.kotlin.cli.jvm.config.jvmClasspathRoots import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar import org.jetbrains.kotlin.config.CommonConfigurationKeys @@ -35,6 +35,18 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor import org.jetbrains.kotlin.descriptors.ModuleDescriptor import org.jetbrains.kotlin.extensions.AnnotationBasedExtension import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor +import org.jetbrains.kotlin.fir.FirSession +import org.jetbrains.kotlin.fir.builder.scriptConfigurators +import org.jetbrains.kotlin.fir.containingClassLookupTag +import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction +import org.jetbrains.kotlin.fir.declarations.utils.isSuspend +import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar +import org.jetbrains.kotlin.fir.extensions.extensionService +import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension +import org.jetbrains.kotlin.fir.resolve.createFunctionalType +import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol +import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType +import org.jetbrains.kotlin.fir.types.coneType import org.jetbrains.kotlin.platform.TargetPlatform import org.jetbrains.kotlin.platform.jvm.isJvm import org.jetbrains.kotlin.psi.KtFile @@ -43,10 +55,12 @@ import org.jetbrains.kotlin.resolve.sam.SamWithReceiverResolver import org.jetbrains.kotlin.scripting.compiler.plugin.ScriptingCompilerConfigurationComponentRegistrar import org.jetbrains.kotlin.scripting.compiler.plugin.dependencies.ScriptsCompilationDependencies import org.jetbrains.kotlin.scripting.compiler.plugin.dependencies.collectScriptsCompilationDependencies +import org.jetbrains.kotlin.scripting.compiler.plugin.services.FirScriptConfiguratorExtensionImpl import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider import org.jetbrains.kotlin.scripting.definitions.annotationsForSamWithReceivers +import org.jetbrains.kotlin.utils.addToStdlib.runIf import kotlin.script.experimental.api.ScriptCompilationConfiguration import kotlin.script.experimental.api.compilerOptions import kotlin.script.experimental.api.dependencies @@ -122,13 +136,48 @@ internal class ScriptingSamWithReceiverComponentContributor(val annotations: Lis } } +internal class FirScriptingSamWithReceiverExtensionRegistrar() : FirExtensionRegistrar() { + override fun ExtensionRegistrarContext.configurePlugin() { + +::FirScriptSamWithReceiverConventionTransformer + } + + class FirScriptSamWithReceiverConventionTransformer( + session: FirSession + ) : FirSamConversionTransformerExtension(session) { + + val knownAnnotations: Set by lazy { + session.extensionService.scriptConfigurators.flatMapTo(mutableSetOf()) { + (it as? FirScriptConfiguratorExtensionImpl)?.knownAnnotationsForSamWithReceiver ?: emptySet() + } + } + + override fun getCustomFunctionalTypeForSamConversion(function: FirSimpleFunction): ConeLookupTagBasedType? { + val containingClassSymbol = function.containingClassLookupTag()?.toFirRegularClassSymbol(session) ?: return null + return runIf(containingClassSymbol.resolvedAnnotationClassIds.any { it.asSingleFqName().asString() in knownAnnotations }) { + val parameterTypes = function.valueParameters.map { it.returnTypeRef.coneType } + if (parameterTypes.isEmpty()) return null + createFunctionalType( + parameters = parameterTypes.subList(1, parameterTypes.size), + receiverType = parameterTypes[0], + rawReturnType = function.returnTypeRef.coneType, + isSuspend = function.isSuspend + ) + } + } + } + +} + internal fun SharedScriptCompilationContext.applyConfigure(): SharedScriptCompilationContext = apply { val samWithReceiverAnnotations = baseScriptCompilationConfiguration[ScriptCompilationConfiguration.annotationsForSamWithReceivers] if (samWithReceiverAnnotations?.isEmpty() == false) { - StorageComponentContainerContributor.registerExtension( - environment.project, - ScriptingSamWithReceiverComponentContributor(samWithReceiverAnnotations.map { it.typeName }) - ) + val annotations = samWithReceiverAnnotations.map { it.typeName } + if (!environment.configuration.getBoolean(CommonConfigurationKeys.USE_FIR)) { + StorageComponentContainerContributor.registerExtension( + environment.project, + ScriptingSamWithReceiverComponentContributor(annotations) + ) + } } } diff --git a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/pluginRegisrar.kt b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/pluginRegisrar.kt index 0738ade22d6..80340b4a921 100644 --- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/pluginRegisrar.kt +++ b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/pluginRegisrar.kt @@ -30,6 +30,7 @@ import org.jetbrains.kotlin.scripting.compiler.plugin.definitions.CliScriptRepor import org.jetbrains.kotlin.scripting.compiler.plugin.extensions.JvmStandardReplFactoryExtension import org.jetbrains.kotlin.scripting.compiler.plugin.extensions.ScriptingCollectAdditionalSourcesExtension import org.jetbrains.kotlin.scripting.compiler.plugin.extensions.ScriptingProcessSourcesBeforeCompilingExtension +import org.jetbrains.kotlin.scripting.compiler.plugin.impl.FirScriptingSamWithReceiverExtensionRegistrar import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider import org.jetbrains.kotlin.scripting.extensions.ScriptExtraImportsProviderExtension @@ -87,6 +88,7 @@ class ScriptingK2CompilerPluginRegistrar : CompilerPluginRegistrar() { // TODO: add jdk path and other params if needed } FirExtensionRegistrarAdapter.registerExtension(FirScriptingCompilerExtensionRegistrar(hostConfiguration, compilerConfiguration)) + FirExtensionRegistrarAdapter.registerExtension(FirScriptingSamWithReceiverExtensionRegistrar()) } } diff --git a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/services/FirScriptConfigurationExtensionImpl.kt b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/services/FirScriptConfigurationExtensionImpl.kt index f7d078ea41e..fcd53e7b4eb 100644 --- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/services/FirScriptConfigurationExtensionImpl.kt +++ b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/services/FirScriptConfigurationExtensionImpl.kt @@ -27,6 +27,7 @@ import org.jetbrains.kotlin.name.ClassId import org.jetbrains.kotlin.name.FqName import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.psi.KtFile +import org.jetbrains.kotlin.scripting.definitions.annotationsForSamWithReceivers import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource import kotlin.script.experimental.api.* @@ -116,6 +117,10 @@ class FirScriptConfiguratorExtensionImpl( } ) } + + compilationConfiguration[ScriptCompilationConfiguration.annotationsForSamWithReceivers]?.forEach { + _knownAnnotationsForSamWithReceiver.add(it.typeName) + } } } @@ -131,12 +136,16 @@ class FirScriptConfiguratorExtensionImpl( } } + private val _knownAnnotationsForSamWithReceiver = hashSetOf() + + internal val knownAnnotationsForSamWithReceiver: Set + get() = _knownAnnotationsForSamWithReceiver + companion object { fun getFactory(hostConfiguration: ScriptingHostConfiguration): Factory { return Factory { session -> FirScriptConfiguratorExtensionImpl(session, hostConfiguration) } } } - } fun KtSourceFile.toSourceCode(): SourceCode? = when (this) {