Pass provided script configuration to refining code

when script compilation configuration refinement happening during parsing, the updated configuration passed to the script compiler/REPL compile function is now used. #KT-44580 fixed
2021-02-08 21:47:20 +01:00
parent 6dd331d7e8
commit ef4fa3381d
8 changed files with 82 additions and 18 deletions
@@ -14,6 +14,8 @@ import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.toScriptSource
 import kotlin.script.experimental.jvm.BasicJvmReplEvaluator
 import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
+import kotlin.script.experimental.jvm.updateClasspath
+import kotlin.script.experimental.jvm.util.classpathFromClass

 class ReplTest : TestCase() {

@@ -191,6 +193,44 @@ class ReplTest : TestCase() {
        )
    }

+    @Test
+    fun testAddNewAnnotationHandler() {
+        val replCompiler = KJvmReplCompilerBase.create(defaultJvmScriptingHostConfiguration)
+        val replEvaluator = BasicJvmReplEvaluator()
+        val compilationConfiguration = ScriptCompilationConfiguration().with {
+            updateClasspath(classpathFromClass<NewAnn>())
+        }
+        val evaluationConfiguration = ScriptEvaluationConfiguration()
+
+        val res0 = runBlocking {
+            replCompiler.compile("1".toScriptSource("Line_0.kts"), compilationConfiguration).onSuccess {
+                replEvaluator.eval(it, evaluationConfiguration)
+            }
+        }
+        assertTrue("Expecting 1 got $res0", res0 is ResultWithDiagnostics.Success && (res0.value.get().result as ResultValue.Value).value == 1)
+
+        var handlerInvoked = false
+
+        val compilationConfiguration2 = compilationConfiguration.with {
+            refineConfiguration {
+//                defaultImports(NewAnn::class) // TODO: fix support for default imports
+                onAnnotations<NewAnn> {
+                    handlerInvoked = true
+                    it.compilationConfiguration.asSuccess()
+                }
+            }
+        }
+
+        val res1 = runBlocking {
+            replCompiler.compile("@file:kotlin.script.experimental.jvmhost.test.NewAnn()\n2".toScriptSource("Line_1.kts"), compilationConfiguration2).onSuccess {
+                replEvaluator.eval(it, evaluationConfiguration)
+            }
+        }
+        assertTrue("Expecting 2 got $res1", res1 is ResultWithDiagnostics.Success && (res1.value.get().result as ResultValue.Value).value == 2)
+
+        assertTrue("Refinement handler on annotation is not invoked", handlerInvoked)
+    }
+
    companion object {
        private fun evaluateInRepl(
            snippets: Sequence<String>,
@@ -300,3 +340,6 @@ class ReplTest : TestCase() {
        )
    }
 }
+
+@Target(AnnotationTarget.FILE)
+annotation class NewAnn
@@ -13,6 +13,7 @@ import com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
+import kotlin.script.experimental.api.ScriptCompilationConfiguration
 import kotlin.script.experimental.api.valueOrNull
 import kotlin.script.experimental.dependencies.ScriptDependencies

@@ -35,6 +36,11 @@ open class ScriptDependenciesProvider constructor(

    open fun getScriptConfigurationResult(file: KtFile): ScriptCompilationConfigurationResult? = null

+    // TODO: consider fixing implementations and removing default implementation
+    open fun getScriptConfigurationResult(
+        file: KtFile, providedConfiguration: ScriptCompilationConfiguration?
+    ): ScriptCompilationConfigurationResult? = getScriptConfigurationResult(file)
+
    open fun getScriptConfiguration(file: KtFile): ScriptCompilationConfigurationWrapper? = getScriptConfigurationResult(file)?.valueOrNull()

    companion object {
@@ -103,7 +103,7 @@ class ScriptLightVirtualFile(name: String, private val _path: String?, text: Str

    override fun getPath(): String = _path ?: if (parent != null) parent.path + "/" + name else name

-    override fun getCanonicalPath(): String? = path
+    override fun getCanonicalPath() = path
 }

 abstract class ScriptCompilationConfigurationWrapper(val script: SourceCode) {
@@ -187,8 +187,8 @@ abstract class ScriptCompilationConfigurationWrapper(val script: SourceCode) {
        override val configuration: ScriptCompilationConfiguration?
            get() {
                val legacy = legacyDependencies ?: return null
-                return definition?.compilationConfiguration?.let {
-                    ScriptCompilationConfiguration(it) {
+                return definition?.compilationConfiguration?.let { config ->
+                    ScriptCompilationConfiguration(config) {
                        updateClasspath(legacy.classpath)
                        defaultImports.append(legacy.imports)
                        importScripts.append(legacy.scripts.map { FileScriptSource(it) })
@@ -221,13 +221,14 @@ typealias ScriptCompilationConfigurationResult = ResultWithDiagnostics<ScriptCom
 fun refineScriptCompilationConfiguration(
    script: SourceCode,
    definition: ScriptDefinition,
-    project: Project
+    project: Project,
+    providedConfiguration: ScriptCompilationConfiguration? = null // if null - take from definition
 ): ScriptCompilationConfigurationResult {
    // TODO: add location information on refinement errors
    val ktFileSource = script.toKtFileSource(definition, project)
    val legacyDefinition = definition.asLegacyOrNull<KotlinScriptDefinition>()
    if (legacyDefinition == null) {
-        val compilationConfiguration = definition.compilationConfiguration
+        val compilationConfiguration = providedConfiguration ?: definition.compilationConfiguration
        val collectedData =
            runReadAction {
                getScriptCollectedData(ktFileSource.ktFile, compilationConfiguration, project, definition.contextClassLoader)
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.scripting.compiler.plugin.definitions

 import com.intellij.openapi.components.ServiceManager
-import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
@@ -15,28 +14,36 @@ import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
 import org.jetbrains.kotlin.scripting.resolve.ScriptReportSink
 import org.jetbrains.kotlin.scripting.resolve.refineScriptCompilationConfiguration
-import java.io.File
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.read
 import kotlin.concurrent.write
-import kotlin.script.experimental.api.ResultWithDiagnostics
+import kotlin.script.experimental.api.ScriptCompilationConfiguration

 class CliScriptDependenciesProvider(project: Project) : ScriptDependenciesProvider(project) {
    private val cacheLock = ReentrantReadWriteLock()
    private val cache = hashMapOf<String, ScriptCompilationConfigurationResult?>()

    override fun getScriptConfigurationResult(file: KtFile): ScriptCompilationConfigurationResult? = cacheLock.read {
-        calculateRefinedConfiguration(file)
+        calculateRefinedConfiguration(file, null)
    }

-    private fun calculateRefinedConfiguration(file: KtFile): ScriptCompilationConfigurationResult? {
+    override fun getScriptConfigurationResult(
+        file: KtFile,
+        providedConfiguration: ScriptCompilationConfiguration?
+    ): ScriptCompilationConfigurationResult? = cacheLock.read {
+        calculateRefinedConfiguration(file, providedConfiguration)
+    }
+
+    private fun calculateRefinedConfiguration(
+        file: KtFile, providedConfiguration: ScriptCompilationConfiguration?
+    ): ScriptCompilationConfigurationResult? {
        val path = file.virtualFilePath
        val cached = cache[path]
        return if (cached != null) cached
        else {
            val scriptDef = file.findScriptDefinition()
            if (scriptDef != null) {
-                val result = refineScriptCompilationConfiguration(KtFileScriptSource(file), scriptDef, project)
+                val result = refineScriptCompilationConfiguration(KtFileScriptSource(file), scriptDef, project, providedConfiguration)

                ServiceManager.getService(project, ScriptReportSink::class.java)?.attachReports(file.virtualFile, result.reports)

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import java.io.File
 import kotlin.script.experimental.api.ResultWithDiagnostics
+import kotlin.script.experimental.api.ScriptCompilationConfiguration
 import kotlin.script.experimental.api.asSuccess
 import kotlin.script.experimental.host.FileBasedScriptSource

@@ -32,7 +33,8 @@ data class ScriptsCompilationDependencies(
 fun collectScriptsCompilationDependencies(
    configuration: CompilerConfiguration,
    project: Project,
-    initialSources: Iterable<KtFile>
+    initialSources: Iterable<KtFile>,
+    providedConfiguration: ScriptCompilationConfiguration? = null
 ): ScriptsCompilationDependencies {
    val collectedClassPath = ArrayList<File>()
    val collectedSources = ArrayList<KtFile>()
@@ -44,7 +46,7 @@ fun collectScriptsCompilationDependencies(
        while (true) {
            val newRemainingSources = ArrayList<KtFile>()
            for (source in remainingSources) {
-                when (val refinedConfiguration = importsProvider.getScriptConfigurationResult(source)) {
+                when (val refinedConfiguration = importsProvider.getScriptConfigurationResult(source, providedConfiguration)) {
                    null -> {}
                    is ResultWithDiagnostics.Failure -> {
                        collectedSourceDependencies.add(ScriptsCompilationDependencies.SourceDependencies(source, refinedConfiguration))
@@ -19,10 +19,10 @@ import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
-import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerStageHistory
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerState
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplCodeAnalyzerBase
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicits
 import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicitsExceptInnermost
@@ -90,6 +90,7 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
                val (sourceFiles, sourceDependencies) = collectRefinedSourcesAndUpdateEnvironment(
                    context,
                    snippetKtFile,
+                    initialConfiguration,
                    messageCollector
                )

@@ -50,7 +50,7 @@ class ScriptJvmCompilerIsolated(val hostConfiguration: ScriptingHostConfiguratio
                    initialConfiguration, hostConfiguration, messageCollector, disposable
                )

-                compileImpl(script, context, messageCollector)
+                compileImpl(script, context, initialConfiguration, messageCollector)
            }
        }
 }
@@ -74,7 +74,7 @@ class ScriptJvmCompilerFromEnvironment(val environment: KotlinCoreEnvironment) :
                try {
                    environment.configuration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, messageCollector)

-                    compileImpl(script, context, messageCollector)
+                    compileImpl(script, context, initialConfiguration, messageCollector)
                } finally {
                    if (parentMessageCollector != null)
                        environment.configuration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, parentMessageCollector)
@@ -105,6 +105,7 @@ private fun withScriptCompilationCache(
 private fun compileImpl(
    script: SourceCode,
    context: SharedScriptCompilationContext,
+    initialConfiguration: ScriptCompilationConfiguration,
    messageCollector: ScriptDiagnosticsMessageCollector
 ): ResultWithDiagnostics<CompiledScript> {
    val mainKtFile =
@@ -121,6 +122,7 @@ private fun compileImpl(
    val (sourceFiles, sourceDependencies) = collectRefinedSourcesAndUpdateEnvironment(
        context,
        mainKtFile,
+        initialConfiguration,
        messageCollector
    )

@@ -32,11 +32,11 @@ 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.resolve.sam.SamWithReceiverResolver
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.jvm.isJvm
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtModifierListOwner
+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
@@ -273,6 +273,7 @@ private fun createInitialCompilerConfiguration(
 internal fun collectRefinedSourcesAndUpdateEnvironment(
    context: SharedScriptCompilationContext,
    mainKtFile: KtFile,
+    initialConfiguration: ScriptCompilationConfiguration,
    messageCollector: ScriptDiagnosticsMessageCollector
 ): Pair<List<KtFile>, List<ScriptsCompilationDependencies.SourceDependencies>> {
    val sourceFiles = arrayListOf(mainKtFile)
@@ -280,7 +281,8 @@ internal fun collectRefinedSourcesAndUpdateEnvironment(
        collectScriptsCompilationDependencies(
            context.environment.configuration,
            context.environment.project,
-            sourceFiles
+            sourceFiles,
+            initialConfiguration
        )

    context.environment.updateClasspath(classpath.map(::JvmClasspathRoot))