Implement import test, fix import support in compiler and evaluator

compiler/backend/src/org/jetbrains/kotlin/codegen/context/ScriptContext.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.codegen.context
 
-import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.FieldInfo
 import org.jetbrains.kotlin.codegen.OwnerKind
 import org.jetbrains.kotlin.codegen.StackValue
@@ -35,7 +34,6 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassNotAny
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.lazy.descriptors.script.ScriptProvidedPropertiesDescriptor
 import org.jetbrains.org.objectweb.asm.Type
-import kotlin.reflect.KClass
 
 class ScriptContext(
     val typeMapper: KotlinTypeMapper,
@@ -83,9 +81,8 @@ class ScriptContext(
         scriptDescriptor.unsubstitutedPrimaryConstructor.valueParameters[ctorImplicitReceiversParametersStart + index].name.identifier
 
     fun getImplicitReceiverType(index: Int): Type? {
-        val receivers = script.kotlinScriptDefinition.implicitReceivers
+        val receiverParam = scriptDescriptor.unsubstitutedPrimaryConstructor.valueParameters[ctorImplicitReceiversParametersStart + index]
-        val kClass = receivers.getOrNull(index)?.classifier as? KClass<*>
+        return state.typeMapper.mapType(receiverParam.type)
-        return kClass?.java?.classId?.let(AsmUtil::asmTypeByClassId)
     }
 
     fun getProvidedPropertyName(index: Int): String =
@@ -113,10 +110,8 @@ class ScriptContext(
 
     fun getScriptFieldName(scriptDescriptor: ScriptDescriptor): String {
         val index = earlierScripts.indexOf(scriptDescriptor)
-        if (index < 0) {
+        return if (index >= 0) "script$" + (index + 1)
-            throw IllegalStateException("Unregistered script: $scriptDescriptor")
+        else "\$\$importedScript${scriptDescriptor.name.identifier}"
-        }
-        return "script$" + (index + 1)
     }
 
     override fun toString(): String {

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

@@ -711,7 +711,7 @@ data class ScriptsCompilationDependencies(
     val sourceDependencies: List<SourceDependencies>
 ) {
     data class SourceDependencies(
-        val script: KtFile,
+        val scriptFile: KtFile,
         val sourceDependencies: List<KtFile>
     )
 }
@@ -721,7 +721,7 @@ fun collectScriptsCompilationDependencies(
     configuration: CompilerConfiguration,
     project: Project,
     initialSources: Iterable<KtFile>
-): ScriptsCompilationDependencies{
+): ScriptsCompilationDependencies {
     val collectedClassPath = ArrayList<File>()
     val collectedSources = ArrayList<KtFile>()
     val collectedSourceDependencies = ArrayList<ScriptsCompilationDependencies.SourceDependencies>()

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

@@ -234,6 +234,9 @@ interface CompiledScript<out ScriptBase : Any> {
      */
     suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>>
 
-    val importedScripts: List<CompiledScript<*>>?
+    /**
-        get() = null
+     * The scripts compiled along with this one in one module, imported or otherwise included into compilation
+     */
+    val otherScripts: List<CompiledScript<*>>
+        get() = emptyList()
 }

libraries/scripting/common/src/kotlin/script/experimental/host/scriptHostUtil.kt

@@ -42,9 +42,9 @@ fun getMergedScriptText(script: SourceCode, configuration: ScriptCompilationConf
 /**
  * The implementation of the SourceCode for a script located in a file
  */
-open class FileScriptSource(val file: File) : ExternalSourceCode {
+open class FileScriptSource(val file: File, private val preloadedText: String? = null) : ExternalSourceCode {
     override val externalLocation: URL get() = file.toURI().toURL()
-    override val text: String by lazy { file.readText() }
+    override val text: String by lazy { preloadedText ?: file.readText() }
 }
 
 /**

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

@@ -14,37 +14,49 @@ import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.jvm.JvmDependency
 import kotlin.script.experimental.jvmhost.JvmScriptEvaluationConfiguration
+import kotlin.script.experimental.jvmhost.actualClassLoader
 import kotlin.script.experimental.jvmhost.baseClassLoader
 
+class KJvmCompiledModule(
+    generationState: GenerationState
+) : Serializable {
+    val compilerOutputFiles: Map<String, ByteArray> =
+        generationState.factory.asList()
+            .associateTo(sortedMapOf<String, ByteArray>()) { it.relativePath to it.asByteArray() }
+
+    companion object {
+        @JvmStatic
+        private val serialVersionUID = 0L
+    }
+}
+
 class KJvmCompiledScript<out ScriptBase : Any>(
     compilationConfiguration: ScriptCompilationConfiguration,
-    generationState: GenerationState,
     private var scriptClassFQName: String,
-    override val importedScripts: List<CompiledScript<*>>? = null
+    override val otherScripts: List<CompiledScript<*>> = emptyList(),
+    private var compiledModule: KJvmCompiledModule? = null
 ) : CompiledScript<ScriptBase>, Serializable {
 
     private var _compilationConfiguration: ScriptCompilationConfiguration? = compilationConfiguration
-    private var compilerOutputFiles: Map<String, ByteArray> = run {
-        val res = sortedMapOf<String, ByteArray>()
-        for (it in generationState.factory.asList()) {
-            res[it.relativePath] = it.asByteArray()
-        }
-        res
-    }
 
     override val compilationConfiguration: ScriptCompilationConfiguration
         get() = _compilationConfiguration!!
 
     override suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>> = try {
-        val baseClassLoader = scriptEvaluationConfiguration?.get(JvmScriptEvaluationConfiguration.baseClassLoader)
+        val classLoader = scriptEvaluationConfiguration?.get(JvmScriptEvaluationConfiguration.actualClassLoader)
-            ?: Thread.currentThread().contextClassLoader
+            ?: run {
-        val dependencies = compilationConfiguration[ScriptCompilationConfiguration.dependencies]
+                if (compiledModule == null)
-            ?.flatMap { (it as? JvmDependency)?.classpath?.map { it.toURI().toURL() } ?: emptyList() }
+                    return ResultWithDiagnostics.Failure("Unable to load class $scriptClassFQName: no compiled module is provided".asErrorDiagnostics())
-        // TODO: previous dependencies and classloaders should be taken into account here
+                val baseClassLoader = scriptEvaluationConfiguration?.get(JvmScriptEvaluationConfiguration.baseClassLoader)
-        val classLoaderWithDeps =
+                    ?: Thread.currentThread().contextClassLoader
-            if (dependencies == null) baseClassLoader
+                val dependencies = compilationConfiguration[ScriptCompilationConfiguration.dependencies]
-            else URLClassLoader(dependencies.toTypedArray(), baseClassLoader)
+                    ?.flatMap { (it as? JvmDependency)?.classpath?.map { it.toURI().toURL() } ?: emptyList() }
-        val classLoader = CompiledScriptClassLoader(classLoaderWithDeps, compilerOutputFiles)
+                // TODO: previous dependencies and classloaders should be taken into account here
+                val classLoaderWithDeps =
+                    if (dependencies == null) baseClassLoader
+                    else URLClassLoader(dependencies.toTypedArray(), baseClassLoader)
+                CompiledScriptClassLoader(classLoaderWithDeps, compiledModule!!.compilerOutputFiles)
+            }
 
         val clazz = classLoader.loadClass(scriptClassFQName).kotlin
         clazz.asSuccess()
@@ -65,13 +77,13 @@ class KJvmCompiledScript<out ScriptBase : Any>(
     }
 
     private fun writeObject(outputStream: ObjectOutputStream) {
-        outputStream.writeObject(compilerOutputFiles)
+        outputStream.writeObject(compiledModule)
         outputStream.writeObject(scriptClassFQName)
     }
 
     private fun readObject(inputStream: ObjectInputStream) {
         _compilationConfiguration = null
-        compilerOutputFiles = inputStream.readObject() as Map<String, ByteArray>
+        compiledModule = inputStream.readObject() as KJvmCompiledModule
         scriptClassFQName = inputStream.readObject() as String
     }
 

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

@@ -38,6 +38,7 @@ import org.jetbrains.kotlin.script.KotlinScriptDefinition
 import org.jetbrains.kotlin.script.util.KotlinJars
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import java.io.File
+import java.util.*
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
 import kotlin.reflect.full.starProjectedType
@@ -172,20 +173,38 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
                 sourceFiles,
                 kotlinCompilerConfiguration
             ).build()
+
             KotlinCodegenFacade.compileCorrectFiles(generationState, CompilationErrorHandler.THROW_EXCEPTION)
 
-            fun makeCompiledScript(script: KtScript): KJvmCompiledScript<*> {
+            val scriptDependenciesStack = ArrayDeque<KtScript>()
 
-                val importedScripts = sourceDependencies.find { it.script == script }?.sourceDependencies?.mapNotNull { sourceFile ->
+            fun makeOtherScripts(script: KtScript): List<KJvmCompiledScript<*>> {
-                    sourceFile.declarations.firstIsInstanceOrNull<KtScript>()?.let { makeCompiledScript(it) }
-                } ?: emptyList()
 
-                return KJvmCompiledScript<Any>(updatedConfiguration, generationState, script.fqName.asString(), importedScripts)
+                // TODO: ensure that it is caught earlier (as well) since it would be more economical
+                if (scriptDependenciesStack.contains(script))
+                    throw IllegalArgumentException("Unable to handle recursive script dependencies")
+                scriptDependenciesStack.push(script)
+
+                val containingKtFile = script.containingKtFile
+                val otherScripts: List<KJvmCompiledScript<*>> =
+                    sourceDependencies.find { it.scriptFile == containingKtFile }?.sourceDependencies?.mapNotNull { sourceFile ->
+                        sourceFile.declarations.firstIsInstanceOrNull<KtScript>()?.let {
+                            KJvmCompiledScript<Any>(updatedConfiguration, it.fqName.asString(), makeOtherScripts(it))
+                        }
+                    } ?: emptyList()
+
+                scriptDependenciesStack.pop()
+                return otherScripts
             }
 
-            val res = makeCompiledScript(ktScript)
+            val compiledScript = KJvmCompiledScript<Any>(
+                updatedConfiguration,
+                ktScript.fqName.asString(),
+                makeOtherScripts(ktScript),
+                KJvmCompiledModule(generationState)
+            )
 
-            return ResultWithDiagnostics.Success(res, messageCollector.diagnostics)
+            return ResultWithDiagnostics.Success(compiledScript, messageCollector.diagnostics)
         } catch (ex: Throwable) {
             return failure(ex.asDiagnostics())
         }

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

@@ -17,6 +17,8 @@ open class JvmScriptEvaluationConfiguration : PropertiesCollection.Builder() {
 
 val JvmScriptEvaluationConfiguration.baseClassLoader by PropertiesCollection.key<ClassLoader?>(Thread.currentThread().contextClassLoader)
 
+val JvmScriptEvaluationConfiguration.actualClassLoader by PropertiesCollection.key<ClassLoader?>()
+
 val ScriptEvaluationConfiguration.jvm get() = JvmScriptEvaluationConfiguration()
 
 open class BasicJvmScriptEvaluator : ScriptEvaluator {
@@ -43,12 +45,30 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
                     }
                     val importedScriptsReports = ArrayList<ScriptDiagnostic>()
                     var importedScriptsLoadingFailed = false
-                    compiledScript.importedScripts?.forEach {
+
-                        val importedScriptEvalRes = invoke(it, scriptEvaluationConfiguration)
+                    // for other scripts we need evaluation configuration with actualClassloader set,
+                    // so they are loaded in the same classloader as the "main" script
+                    val updatedEvalConfiguration = when {
+                        scriptEvaluationConfiguration == null -> ScriptEvaluationConfiguration {
+                            // TODO: find out why dsl syntax doesn't work here
+                            set(JvmScriptEvaluationConfiguration.actualClassLoader, scriptClass.java.classLoader)
+                        }
+                        scriptEvaluationConfiguration.getNoDefault(JvmScriptEvaluationConfiguration.actualClassLoader) == null -> ScriptEvaluationConfiguration(scriptEvaluationConfiguration) {
+                            // TODO: find out why dsl syntax doesn't work here
+                            set(JvmScriptEvaluationConfiguration.actualClassLoader, scriptClass.java.classLoader)
+                        }
+                        else -> scriptEvaluationConfiguration
+                    }
+
+                    compiledScript.otherScripts.forEach {
+                        // TODO: in the future other scripts could be used for other purposes, so args here should be added only for actually imported scripts
+                        // (it means that we should keep mapping somewhere (or reuse one with source dependencies) between imported scrips and e.g. fqnames)
+                        val importedScriptEvalRes = invoke(it, updatedEvalConfiguration)
                         importedScriptsReports.addAll(importedScriptEvalRes.reports)
                         when (importedScriptEvalRes) {
                             is ResultWithDiagnostics.Success -> {
-                                args.add(importedScriptEvalRes.value)
+                                // TODO: checks and diagnostics
+                                args.add((importedScriptEvalRes.value.returnValue as ResultValue.Value).value)
                             }
                             else -> {
                                 importedScriptsLoadingFailed = true
@@ -60,14 +80,14 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
                         ResultWithDiagnostics.Failure(importedScriptsReports)
                     } else {
 
-                        scriptEvaluationConfiguration?.get(ScriptEvaluationConfiguration.constructorArgs)?.let {
+                        updatedEvalConfiguration[ScriptEvaluationConfiguration.constructorArgs]?.let {
                             args.addAll(it)
                         }
                         val ctor = scriptClass.java.constructors.single()
                         val instance = ctor.newInstance(*args.toArray())
 
                         // TODO: fix result value
-                        ResultWithDiagnostics.Success(EvaluationResult(ResultValue.Value("", instance, ""), scriptEvaluationConfiguration))
+                        ResultWithDiagnostics.Success(EvaluationResult(ResultValue.Value("", instance, ""), updatedEvalConfiguration))
                     }
                 }
             }

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

@@ -15,6 +15,7 @@ import java.security.MessageDigest
 import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.BasicScriptingHost
+import kotlin.script.experimental.host.FileScriptSource
 import kotlin.script.experimental.host.toScriptSource
 import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
 import kotlin.script.experimental.jvmhost.*
@@ -50,6 +51,31 @@ class ScriptingHostTest : TestCase() {
         Assert.assertEquals(greeting, output)
     }
 
+    @Test
+    fun testSimpleImport() {
+        val greeting = "Hello from helloWithVal script!\nHello from imported helloWithVal script!"
+        val script = "println(\"Hello from imported \$helloScriptName script!\")"
+        val compilationConfiguration = createJvmCompilationConfigurationFromTemplate<SimpleScriptTemplate> {
+            refineConfiguration {
+                beforeCompiling { ctx ->
+                    val importedScript = File(TEST_DATA_DIR, "importTest/helloWithVal.kts")
+                    if ((ctx.script as? FileScriptSource)?.file?.canonicalFile == importedScript.canonicalFile) {
+                        ctx.compilationConfiguration
+                    } else {
+                        ScriptCompilationConfiguration(ctx.compilationConfiguration) {
+                            importScripts(importedScript.toScriptSource())
+                        }
+                    }.asSuccess()
+                }
+            }
+        }
+        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\")"
@@ -159,7 +185,11 @@ class ScriptingHostTest : TestCase() {
 
 fun ResultWithDiagnostics<*>.throwOnFailure(): ResultWithDiagnostics<*> = apply {
     if (this is ResultWithDiagnostics.Failure) {
-        throw Exception("Compilation/evaluation failed:\n  ${reports.joinToString("\n  ") { it.exception?.toString() ?: it.message }}")
+        val firstExceptionFromReports = reports.find { it.exception != null }?.exception
+        throw Exception(
+            "Compilation/evaluation failed:\n  ${reports.joinToString("\n  ") { it.exception?.toString() ?: it.message }}",
+            firstExceptionFromReports
+        )
     }
 }
 

libraries/scripting/jvm-host/testData/importTest/helloWithVal.kts

@@ -0,0 +1,4 @@
+
+val helloScriptName = "helloWithVal"
+
+println("Hello from $helloScriptName script!")

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

@@ -14,6 +14,7 @@ import kotlin.script.experimental.dependencies.AsyncDependenciesResolver
 import kotlin.script.experimental.dependencies.DependenciesResolver
 import kotlin.script.experimental.dependencies.ScriptDependencies
 import kotlin.script.experimental.dependencies.ScriptReport
+import kotlin.script.experimental.host.FileScriptSource
 import kotlin.script.experimental.host.toScriptSource
 import kotlin.script.experimental.jvm.JvmDependency
 import kotlin.script.experimental.jvm.compat.mapToLegacyScriptReportPosition
@@ -94,8 +95,8 @@ internal fun List<ScriptDiagnostic>.mapScriptReportsToDiagnostics() =
     map { ScriptReport(it.message, mapToLegacyScriptReportSeverity(it.severity), mapToLegacyScriptReportPosition(it.location)) }
 
 internal fun ScriptContents.toScriptSource(): SourceCode = when {
+    file != null -> FileScriptSource(file!!, text?.toString())
     text != null -> text!!.toString().toScriptSource()
-    file != null -> file!!.toScriptSource()
     else -> throw IllegalArgumentException("Unable to convert script contents $this into script source")
 }