Implement simple "add to compilation" functionality for scripts

(incomplete import - scripts added to the compilation but not yet
imported into scope)
Base functionality needed for #KT-27960

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt

@@ -22,7 +22,6 @@ import com.intellij.codeInsight.InferredAnnotationsManager
 import com.intellij.codeInsight.runner.JavaMainMethodProvider
 import com.intellij.core.*
 import com.intellij.ide.highlighter.JavaFileType
-import com.intellij.ide.plugins.PluginManagerCore
 import com.intellij.lang.MetaLanguage
 import com.intellij.lang.java.JavaParserDefinition
 import com.intellij.openapi.Disposable
@@ -241,16 +240,11 @@ class KotlinCoreEnvironment private constructor(
         }
 
         sourceFiles += createKtFiles(project)
-        sourceFiles.sortBy { it.virtualFile.path }
 
         if (scriptDefinitionProvider != null) {
-            ScriptDependenciesProvider.getInstance(project).let { importsProvider ->
+            sourceFiles += collectRequiredSourcesFromDependencies(configuration, project, sourceFiles)
-                configuration.addJvmClasspathRoots(
-                    sourceFiles.mapNotNull(importsProvider::getScriptDependencies)
-                        .flatMap { it.classpath }
-                        .distinctBy { it.absolutePath })
-            }
         }
+        sourceFiles.sortBy { it.virtualFile.path }
 
         val jdkHome = configuration.get(JVMConfigurationKeys.JDK_HOME)
         val jrtFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.JRT_PROTOCOL)
@@ -427,61 +421,13 @@ class KotlinCoreEnvironment private constructor(
 
     fun getSourceFiles(): List<KtFile> = sourceFiles
 
-    private fun createKtFiles(project: Project): List<KtFile> {
+    private fun createKtFiles(project: Project): List<KtFile> =
-        val sourceRoots = getSourceRootsCheckingForDuplicates()
+        createSourceFilesFromSourceRoots(configuration, project, getSourceRootsCheckingForDuplicates())
 
-        val localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
+    internal fun report(severity: CompilerMessageSeverity, message: String) = report(configuration, severity, message)
-        val psiManager = PsiManager.getInstance(project)
-
-        val processedFiles = hashSetOf<VirtualFile>()
-        val result = mutableListOf<KtFile>()
-
-        val virtualFileCreator = PreprocessedFileCreator(project)
-
-        for ((sourceRootPath, isCommon) in sourceRoots) {
-            val vFile = localFileSystem.findFileByPath(sourceRootPath)
-            if (vFile == null) {
-                val message = "Source file or directory not found: $sourceRootPath"
-
-                val buildFilePath = configuration.get(JVMConfigurationKeys.MODULE_XML_FILE)
-                if (buildFilePath != null && Logger.isInitialized()) {
-                    LOG.warn("$message\n\nbuild file path: $buildFilePath\ncontent:\n${buildFilePath.readText()}")
-                }
-
-                report(ERROR, message)
-                continue
-            }
-
-            if (!vFile.isDirectory && vFile.fileType != KotlinFileType.INSTANCE) {
-                report(ERROR, "Source entry is not a Kotlin file: $sourceRootPath")
-                continue
-            }
-
-            for (file in File(sourceRootPath).walkTopDown()) {
-                if (!file.isFile) continue
-
-                val virtualFile = localFileSystem.findFileByPath(file.absolutePath)?.let(virtualFileCreator::create)
-                if (virtualFile != null && processedFiles.add(virtualFile)) {
-                    val psiFile = psiManager.findFile(virtualFile)
-                    if (psiFile is KtFile) {
-                        result.add(psiFile)
-                        if (isCommon) {
-                            psiFile.isCommonSource = true
-                        }
-                    }
-                }
-            }
-        }
-
-        return result
-    }
-
-    private fun report(severity: CompilerMessageSeverity, message: String) {
-        configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY).report(severity, message)
-    }
 
     companion object {
-        private val LOG = Logger.getInstance(KotlinCoreEnvironment::class.java)
+        internal val LOG = Logger.getInstance(KotlinCoreEnvironment::class.java)
 
         private val APPLICATION_LOCK = Object()
         private var ourApplicationEnvironment: JavaCoreApplicationEnvironment? = null
@@ -701,3 +647,85 @@ class KotlinCoreEnvironment private constructor(
         }
     }
 }
+
+private fun report(configuration: CompilerConfiguration, severity: CompilerMessageSeverity, message: String) {
+    configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY).report(severity, message)
+}
+
+private fun createSourceFilesFromSourceRoots(
+    configuration: CompilerConfiguration,
+    project: Project,
+    sourceRoots: List<KotlinSourceRoot>
+): MutableList<KtFile> {
+    val localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
+    val psiManager = PsiManager.getInstance(project)
+
+    val processedFiles = hashSetOf<VirtualFile>()
+    val result = mutableListOf<KtFile>()
+
+    val virtualFileCreator = PreprocessedFileCreator(project)
+
+    for ((sourceRootPath, isCommon) in sourceRoots) {
+        val vFile = localFileSystem.findFileByPath(sourceRootPath)
+        if (vFile == null) {
+            val message = "Source file or directory not found: $sourceRootPath"
+
+            val buildFilePath = configuration.get(JVMConfigurationKeys.MODULE_XML_FILE)
+            if (buildFilePath != null && Logger.isInitialized()) {
+                KotlinCoreEnvironment.LOG.warn("$message\n\nbuild file path: $buildFilePath\ncontent:\n${buildFilePath.readText()}")
+            }
+
+            report(configuration, ERROR, message)
+            continue
+        }
+
+        if (!vFile.isDirectory && vFile.fileType != KotlinFileType.INSTANCE) {
+            report(configuration, ERROR, "Source entry is not a Kotlin file: $sourceRootPath")
+            continue
+        }
+
+        for (file in File(sourceRootPath).walkTopDown()) {
+            if (!file.isFile) continue
+
+            val virtualFile = localFileSystem.findFileByPath(file.absolutePath)?.let(virtualFileCreator::create)
+            if (virtualFile != null && processedFiles.add(virtualFile)) {
+                val psiFile = psiManager.findFile(virtualFile)
+                if (psiFile is KtFile) {
+                    result.add(psiFile)
+                    if (isCommon) {
+                        psiFile.isCommonSource = true
+                    }
+                }
+            }
+        }
+    }
+
+    return result
+}
+
+fun collectRequiredSourcesFromDependencies(
+    configuration: CompilerConfiguration,
+    project: Project,
+    initialSources: Iterable<KtFile>
+): List<KtFile> {
+    val collectedSources = ArrayList<KtFile>()
+    val importsProvider = ScriptDependenciesProvider.getInstance(project)
+    var remainingSources = initialSources.sortedBy { it.virtualFile.path }
+    val knownSourcePaths = HashSet<String>()
+    while (true) {
+        val dependencies = remainingSources.mapNotNull(importsProvider::getScriptDependencies)
+        configuration.addJvmClasspathRoots(dependencies.flatMap { it.classpath }.distinctBy { it.absolutePath })
+        val addedSourceRoots = dependencies.flatMap { it.scripts.map { KotlinSourceRoot(it.path, false) } }
+        val addedSources = createSourceFilesFromSourceRoots(configuration, project, addedSourceRoots)
+        val newSources = if (addedSources.isEmpty()) emptyList() else {
+            remainingSources.forEach { knownSourcePaths.add(it.virtualFile.path) }
+            addedSources.filterNot { knownSourcePaths.contains(it.virtualFile.path) }
+        }
+        if (newSources.isEmpty()) break
+        else {
+            collectedSources += newSources
+            remainingSources = newSources
+        }
+    }
+    return collectedSources
+}

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt

@@ -267,7 +267,7 @@ object KotlinToJVMBytecodeCompiler {
         }
     }
 
-    internal fun compileAndExecuteScript(environment: KotlinCoreEnvironment, scriptArgs: List<String>): ExitCode {
+    fun compileAndExecuteScript(environment: KotlinCoreEnvironment, scriptArgs: List<String>): ExitCode {
         val scriptClass = compileScript(environment) ?: return ExitCode.COMPILATION_ERROR
 
         try {

compiler/testData/script/hello.kts

@@ -0,0 +1,2 @@
+
+print("Hello from basic script!")

compiler/testData/script/hello.req1.kts

@@ -0,0 +1,4 @@
+
+@file:Import("hello_req.kt")
+
+print("Hello from ${RequiredFromHello().name}!")

compiler/testData/script/hello_args.kts

@@ -0,0 +1,2 @@
+
+print("Hello, ${args[0]}!")

compiler/testData/script/hello_req.kt

@@ -0,0 +1,5 @@
+
+class RequiredFromHello {
+
+    val name = "required"
+}

compiler/tests/org/jetbrains/kotlin/scripts/ScriptCliCompilationTest.kt

@@ -0,0 +1,148 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.scripts
+
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY
+import org.jetbrains.kotlin.cli.common.ExitCode
+import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoot
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinToJVMBytecodeCompiler
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.daemon.TestMessageCollector
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.scripting.compiler.plugin.KotlinScriptDefinitionAdapterFromNewAPI
+import org.jetbrains.kotlin.test.ConfigurationKind
+import org.jetbrains.kotlin.test.KotlinTestUtils
+import org.jetbrains.kotlin.test.TestJdkKind
+import org.jetbrains.kotlin.test.testFramework.KtUsefulTestCase
+import org.jetbrains.kotlin.utils.PathUtil
+import org.junit.Assert
+import org.junit.Test
+import java.io.File
+import kotlin.reflect.KClass
+import kotlin.script.experimental.annotations.KotlinScript
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.FileScriptSource
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+import kotlin.script.experimental.host.configurationDependencies
+import kotlin.script.experimental.host.createCompilationConfigurationFromTemplate
+import kotlin.script.experimental.jvm.JvmDependency
+import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
+import kotlin.script.experimental.jvm.dependenciesFromCurrentContext
+import kotlin.script.experimental.jvm.jvm
+
+private const val testDataPath = "compiler/testData/script/"
+
+class ScriptCliCompilationTest : KtUsefulTestCase() {
+
+    @Test
+    fun testPrerequisites() {
+        Assert.assertTrue(thisClasspath.isNotEmpty())
+    }
+
+    @Test
+    fun testSimpleScript() {
+        val out = checkRun("hello.kts")
+        Assert.assertEquals("Hello from basic script!", out)
+    }
+
+    @Test
+    fun testSimpleScriptWithArgs() {
+        val out = checkRun("hello_args.kts", listOf("kotlin"))
+        Assert.assertEquals("Hello, kotlin!", out)
+    }
+
+    @Test
+    fun testScriptWithRequire() {
+        val out = checkRun("hello.req1.kts", scriptDef = TestScriptWithRequire::class)
+        Assert.assertEquals("Hello from required!", out)
+    }
+
+
+    private val thisClasspath = listOf(PathUtil.getResourcePathForClass(ScriptCliCompilationTest::class.java))
+
+    private fun runCompiler(
+        script: File,
+        args: List<String> = emptyList(),
+        scriptDef: KClass<*>? = null,
+        classpath: List<File> = emptyList()
+    ): Pair<ExitCode, MessageCollector> {
+
+        val collector = TestMessageCollector()
+
+        val configuration = KotlinTestUtils.newConfiguration(ConfigurationKind.NO_KOTLIN_REFLECT, TestJdkKind.FULL_JDK).apply {
+            put(MESSAGE_COLLECTOR_KEY, collector)
+            addKotlinSourceRoot(script.path)
+            if (scriptDef != null) {
+                val hostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {
+                    configurationDependencies(JvmDependency(classpath))
+                }
+                val scriptDefinition = KotlinScriptDefinitionAdapterFromNewAPI(
+                    createCompilationConfigurationFromTemplate(
+                        KotlinType(scriptDef),
+                        hostConfiguration, KotlinScriptDefinition::class
+                    ),
+                    hostConfiguration
+                )
+                add(JVMConfigurationKeys.SCRIPT_DEFINITIONS, scriptDefinition)
+            }
+        }
+
+        val environment = KotlinCoreEnvironment.createForTests(testRootDisposable, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES)
+
+        return KotlinToJVMBytecodeCompiler.compileAndExecuteScript(environment, args) to collector
+    }
+
+    private fun checkRun(
+        scriptFileName: String,
+        args: List<String> = emptyList(),
+        scriptDef: KClass<*>? = null,
+        classpath: List<File> = emptyList()
+    ): String =
+        captureOut {
+            val res = runCompiler(File(testDataPath, scriptFileName), args, scriptDef, classpath)
+            val resMessage = lazy {
+                "Compilation results:\n" + res.second.toString()
+            }
+            Assert.assertEquals(resMessage.value, ExitCode.OK, res.first)
+            Assert.assertFalse(resMessage.value, res.second.hasErrors())
+        }
+}
+
+@KotlinScript(
+    fileExtension = "req1.kts",
+    compilationConfiguration = TestScriptWithRequireConfiguration::class
+)
+abstract class TestScriptWithRequire
+
+object TestScriptWithRequireConfiguration : ScriptCompilationConfiguration(
+    {
+        defaultImports(Import::class)
+        jvm {
+            dependenciesFromCurrentContext(wholeClasspath = true)
+        }
+        refineConfiguration {
+            onAnnotations(Import::class) { context: ScriptConfigurationRefinementContext ->
+                val sources = context.collectedData?.get(ScriptCollectedData.foundAnnotations)
+                    ?.flatMap {
+                        (it as? Import)?.sources?.map { sourceName -> FileScriptSource(File(testDataPath, sourceName)) } ?: emptyList()
+                    }
+                    ?.takeIf { it.isNotEmpty() }
+                    ?: return@onAnnotations context.compilationConfiguration.asSuccess()
+                ScriptCompilationConfiguration(context.compilationConfiguration) {
+                    importScripts.append(sources)
+                }.asSuccess()
+            }
+        }
+    }
+)
+
+@Target(AnnotationTarget.FILE)
+@Repeatable
+@Retention(AnnotationRetention.SOURCE)
+annotation class Import(vararg val sources: String)

libraries/scripting/common/src/kotlin/script/experimental/api/scriptCompilation.kt

@@ -77,6 +77,11 @@ val ScriptCompilationConfigurationKeys.providedProperties by PropertiesCollectio
  */
 val ScriptCompilationConfigurationKeys.defaultImports by PropertiesCollection.key<List<String>>()
 
+/**
+ * The list of script sources that should be compiled along with the script and imported into it
+ */
+val ScriptCompilationConfigurationKeys.importScripts by PropertiesCollection.key<List<SourceCode>>()
+
 /**
  * The list of script dependencies - platform specific
  */

libraries/scripting/jvm-host/build.gradle.kts

@@ -15,6 +15,7 @@ dependencies {
     compileOnly(project(":compiler:cli"))
     compileOnly(intellijCoreDep())
     runtime(projectRuntimeJar(":kotlin-compiler"))
+    testCompile(projectTests(":compiler:tests-common"))
     testCompile(commonDep("junit"))
 }
 
@@ -27,3 +28,6 @@ standardPublicJars()
 
 publish()
 
+projectTest {
+    workingDir = rootDir
+}

libraries/scripting/jvm-host/src/kotlin/script/experimental/jvmhost/impl/KJvmCompilerImpl.kt

@@ -16,15 +16,13 @@ import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
-import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.*
-import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
-import org.jetbrains.kotlin.cli.jvm.compiler.NoScopeRecordCliBindingTrace
-import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.cli.jvm.config.JvmClasspathRoot
 import org.jetbrains.kotlin.cli.jvm.config.JvmModulePathRoot
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
 import org.jetbrains.kotlin.cli.jvm.modules.CoreJrtFileSystem
 import org.jetbrains.kotlin.codegen.ClassBuilderFactories
+import org.jetbrains.kotlin.codegen.CompilationErrorHandler
 import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.*
@@ -133,7 +131,9 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
             val psiFile: KtFile = psiFileFactory.trySetupPsiForFile(virtualFile, KotlinLanguage.INSTANCE, true, false) as KtFile?
                 ?: return failure("Unable to make PSI file from script".asErrorDiagnostics())
 
-            val sourceFiles = listOf(psiFile)
+            val sourceFiles = arrayListOf(psiFile).also {
+                it.addAll(collectRequiredSourcesFromDependencies(kotlinCompilerConfiguration, environment.project, it))
+            }
 
             analyzerWithCompilerReport.analyzeAndReport(sourceFiles) {
                 val project = environment.project
@@ -158,13 +158,7 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
                 sourceFiles,
                 kotlinCompilerConfiguration
             ).build()
-            generationState.beforeCompile()
+            KotlinCodegenFacade.compileCorrectFiles(generationState, CompilationErrorHandler.THROW_EXCEPTION)
-            KotlinCodegenFacade.generatePackage(
-                generationState,
-                psiFile.script!!.containingKtFile.packageFqName,
-                setOf(psiFile.script!!.containingKtFile),
-                org.jetbrains.kotlin.codegen.CompilationErrorHandler.THROW_EXCEPTION
-            )
 
             val res = KJvmCompiledScript<Any>(updatedConfiguration, generationState, scriptFileName.capitalize())
 

libraries/scripting/jvm-host/test/kotlin/script/experimental/jvmhost/test/ScriptingHostTest.kt

@@ -5,6 +5,7 @@
 
 package kotlin.script.experimental.jvmhost.test
 
+import junit.framework.TestCase
 import kotlinx.coroutines.runBlocking
 import org.junit.Assert
 import org.junit.Test
@@ -21,7 +22,11 @@ import kotlin.script.experimental.jvmhost.impl.CompiledScriptClassLoader
 import kotlin.script.experimental.jvmhost.impl.KJvmCompiledScript
 import kotlin.script.templates.standard.SimpleScriptTemplate
 
-class ScriptingHostTest {
+class ScriptingHostTest : TestCase() {
+
+    companion object {
+        const val TEST_DATA_DIR = "libraries/scripting/jvm-host/testData"
+    }
 
     @Test
     fun testSimpleUsage() {
@@ -32,6 +37,19 @@ class ScriptingHostTest {
         Assert.assertEquals(greeting, output)
     }
 
+    @Test
+    fun testSimpleRequire() {
+        val greeting = "Hello from required!"
+        val script = "val subj = RequiredClass().value\nprintln(\"Hello from \$subj!\")"
+        val compilationConfiguration = createJvmCompilationConfigurationFromTemplate<SimpleScriptTemplate> {
+            importScripts(File(TEST_DATA_DIR, "importTest/requiredSrc.kt").toScriptSource())
+        }
+        val output = captureOut {
+            BasicJvmScriptingHost().eval(script.toScriptSource(), compilationConfiguration, null).throwOnFailure()
+        }
+        Assert.assertEquals(greeting, output)
+    }
+
     @Test
     fun testMemoryCache() {
         val script = "val x = 1\nprintln(\"x = \$x\")"

libraries/scripting/jvm-host/testData/importTest/requiredSrc.kt

@@ -0,0 +1,5 @@
+
+class RequiredClass {
+
+    val value = "required"
+}

libraries/scripting/jvm/src/kotlin/script/experimental/jvm/impl/BridgeDependenciesResolver.kt

@@ -44,11 +44,13 @@ class BridgeDependenciesResolver(
 
             val defaultImports = scriptCompilationConfiguration[ScriptCompilationConfiguration.defaultImports]?.toList() ?: emptyList()
 
-            fun ScriptCompilationConfiguration.toDependencies(classpath: List<File>): ScriptDependencies = ScriptDependencies(
+            fun ScriptCompilationConfiguration.toDependencies(classpath: List<File>): ScriptDependencies =
-                classpath = classpath,
+                ScriptDependencies(
-                sources = this[ScriptCompilationConfiguration.ide.dependenciesSources].toClassPathOrEmpty(),
+                    classpath = classpath,
-                imports = defaultImports
+                    sources = this[ScriptCompilationConfiguration.ide.dependenciesSources].toClassPathOrEmpty(),
-            )
+                    imports = defaultImports,
+                    scripts = this[ScriptCompilationConfiguration.importScripts].toFilesOrEmpty()
+                )
 
             val refineResults = scriptCompilationConfiguration.refineWith(
                 scriptCompilationConfiguration[ScriptCompilationConfiguration.refineConfigurationOnAnnotations]?.handler,
@@ -76,11 +78,8 @@ class BridgeDependenciesResolver(
             }
 
             return DependenciesResolver.ResolveResult.Success(
-                ScriptDependencies(
+                // TODO: consider returning only increment from the initial config
-                    classpath = newClasspath, // TODO: maybe it should return only increment from the initial config
+                refinedConfiguration.toDependencies(newClasspath),
-                    sources = refinedConfiguration[ScriptCompilationConfiguration.ide.dependenciesSources].toClassPathOrEmpty(),
-                    imports = defaultImports
-                ),
                 diagnostics
             )
         } catch (e: Throwable) {