K2 Scripting: add support for imported scripts

2023-06-19 10:13:03 +02:00
parent d24fc3b581
commit 96bde033e1
21 changed files with 306 additions and 38 deletions
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.scripting.compiler.plugin

 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.scripting.compiler.plugin.services.Fir2IrScriptConfiguratorExtensionImpl
 import org.jetbrains.kotlin.scripting.compiler.plugin.services.FirScriptConfiguratorExtensionImpl
 import org.jetbrains.kotlin.scripting.compiler.plugin.services.FirScriptDefinitionProviderService
 import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
@@ -27,5 +28,6 @@ class FirScriptingCompilerExtensionRegistrar(
        }

        +FirScriptConfiguratorExtensionImpl.getFactory(hostConfiguration)
+        +Fir2IrScriptConfiguratorExtensionImpl.getFactory(hostConfiguration)
    }
 }
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.pipeline.*
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
@@ -39,6 +40,9 @@ import org.jetbrains.kotlin.scripting.compiler.plugin.dependencies.ScriptsCompil
 import org.jetbrains.kotlin.scripting.compiler.plugin.services.scriptDefinitionProviderService
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
+import org.jetbrains.kotlin.scripting.resolve.resolvedImportScripts
+import org.jetbrains.kotlin.utils.topologicalSort
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
 import kotlin.script.experimental.jvm.JvmDependency
@@ -370,14 +374,49 @@ private fun doCompileWithK2(
        }
    ).single().session

-    session.scriptDefinitionProviderService?.run {
+    val scriptDefinitionProviderService = session.scriptDefinitionProviderService
+
+    scriptDefinitionProviderService?.run {
        definitionProvider = ScriptDefinitionProvider.getInstance(context.environment.project)
        configurationProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
    }

-    val rawFir = session.buildFirFromKtFiles(sourceFiles)
+    val rawFir = session.buildFirFromKtFiles(sourceFiles) //.reversed()

-    val (scopeSession, fir) = session.runResolution(rawFir)
+    val orderedRawFir =
+        if (scriptDefinitionProviderService == null) rawFir
+        else {
+            val rawFirDeps = rawFir.associateWith { firFile ->
+                ((firFile.sourceFile as? KtPsiSourceFile)?.psiFile as? KtFile)?.let { ktFile ->
+                    val scriptCompilationConfiguration = scriptDefinitionProviderService.configurationProvider?.getScriptConfiguration(ktFile)?.configuration
+                    scriptCompilationConfiguration?.get(ScriptCompilationConfiguration.resolvedImportScripts)?.mapNotNull { depSource ->
+                        (depSource as? VirtualFileScriptSource)?.virtualFile?.let { depVFile ->
+                            rawFir.find { ((it.sourceFile as? KtPsiSourceFile)?.psiFile as? KtFile)?.virtualFile == depVFile }
+                        }
+                    }
+                }.orEmpty()
+            }
+
+            class CycleDetected(val node: FirFile) : Throwable()
+
+            try {
+                topologicalSort(
+                    rawFir, reportCycle = { throw CycleDetected(it) }
+                ) {
+                    rawFirDeps[this] ?: emptyList()
+                }.reversed()
+            } catch (e: CycleDetected) {
+                return ResultWithDiagnostics.Failure(
+                    ScriptDiagnostic(
+                        ScriptDiagnostic.unspecifiedError,
+                        "Unable to handle recursive script dependencies, cycle detected on file ${e.node.name}",
+                        sourcePath = e.node.sourceFile?.path
+                    )
+                )
+            }
+        }
+
+    val (scopeSession, fir) = session.runResolution(orderedRawFir)
    // checkers
    session.runCheckers(scopeSession, fir, diagnosticsReporter)

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2010-2023 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.scripting.compiler.plugin.services
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.backend.Fir2IrScriptConfiguratorExtension
+import org.jetbrains.kotlin.fir.backend.Fir2IrScriptConfiguratorExtension.Factory
+import org.jetbrains.kotlin.fir.declarations.FirScript
+import org.jetbrains.kotlin.fir.resolve.providers.firProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
+import org.jetbrains.kotlin.ir.declarations.IrScript
+import org.jetbrains.kotlin.ir.symbols.IrScriptSymbol
+import org.jetbrains.kotlin.scripting.resolve.resolvedImportScripts
+import kotlin.script.experimental.api.ScriptCompilationConfiguration
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+
+class Fir2IrScriptConfiguratorExtensionImpl(
+    session: FirSession,
+    @Suppress("UNUSED_PARAMETER") hostConfiguration: ScriptingHostConfiguration
+) : Fir2IrScriptConfiguratorExtension(session) {
+    override fun IrScript.configure(script: FirScript, getIrScriptByFirSymbol: (FirScriptSymbol) -> IrScriptSymbol?) {
+        // processing only refined scripts here
+        val scriptFile = session.firProvider.getFirScriptContainerFile(script.symbol) ?: return
+        val scriptSourceFile = scriptFile.sourceFile?.toSourceCode() ?: return
+        val compilationConfiguration = session.getScriptCompilationConfiguration(scriptSourceFile, getDefault = { null }) ?: return
+
+        // assuming that if the script is compiled, the import files should be all resolved already
+        val importedScripts = compilationConfiguration[ScriptCompilationConfiguration.resolvedImportScripts]?.takeIf { it.isNotEmpty() } ?: return
+        val importedScriptSymbols = importedScripts.mapNotNull {
+            session.firProvider.getFirScriptByFilePath(it.locationId!!)  // TODO: all !! should be converted to diagnostics
+        }
+
+        this.importedScripts = importedScriptSymbols.map { getIrScriptByFirSymbol(it)!! }.takeIf { it.isNotEmpty() }
+    }
+
+    companion object {
+        fun getFactory(hostConfiguration: ScriptingHostConfiguration): Factory {
+            return Factory { session -> Fir2IrScriptConfiguratorExtensionImpl(session, hostConfiguration) }
+        }
+    }
+}
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2023 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.scripting.compiler.plugin.services
+
+import org.jetbrains.kotlin.*
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
+import kotlin.script.experimental.api.ScriptCompilationConfiguration
+import kotlin.script.experimental.api.SourceCode
+import kotlin.script.experimental.api.valueOrNull
+import kotlin.script.experimental.host.FileScriptSource
+import kotlin.script.experimental.host.StringScriptSource
+
+fun FirSession.getScriptCompilationConfiguration(
+    sourceCode: SourceCode?,
+    getDefault: FirScriptDefinitionProviderService.() -> ScriptCompilationConfiguration? = { definitionProvider?.getDefaultDefinition()?.compilationConfiguration }
+) =
+    scriptDefinitionProviderService?.let { providerService ->
+        sourceCode?.let { script ->
+            val ktFile = (script as? KtFileScriptSource)?.ktFile ?: error("only PSI scripts are supported at the moment")
+            providerService.configurationProvider?.getScriptConfigurationResult(ktFile)?.valueOrNull()?.configuration
+                ?: providerService.getDefault()
+        } ?: providerService.getDefault()
+    }
@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.script.loadScriptingPlugin
 import org.jetbrains.kotlin.scripting.compiler.plugin.TestDisposable
-import org.jetbrains.kotlin.scripting.compiler.plugin.expectTestToFailOnK2
 import org.jetbrains.kotlin.scripting.compiler.plugin.impl.ScriptJvmCompilerFromEnvironment
 import org.jetbrains.kotlin.scripting.compiler.plugin.updateWithBaseCompilerArguments
 import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
@@ -24,7 +23,6 @@ import org.jetbrains.kotlin.test.TestJdkKind
 import org.junit.Assert
 import java.io.File
 import java.nio.file.Files
-import kotlin.io.path.*
 import kotlin.script.experimental.annotations.KotlinScript
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -38,7 +36,7 @@ private const val testDataPath = "plugins/scripting/scripting-compiler/testData/
 class CompileTimeFibonacciTest : TestCase() {
    private val testRootDisposable: Disposable = TestDisposable()

-    fun testFibonacciWithSupportedNumbersImplementsTheCorrectConstants() = expectTestToFailOnK2 {
+    fun testFibonacciWithSupportedNumbersImplementsTheCorrectConstants() {
        val outputLines = runScript("supported.fib.kts")
            .valueOr { failure ->
                val message = failure.reports.joinToString("\n") { it.message }