Implement support for new script customizations in command line compiler

refactor and fix internals in many places to make it work
add appropriate tests for main-kts case

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/JvmCliScriptEvaluationExtension.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.scripting.compiler.plugin
 
 import com.intellij.core.JavaCoreProjectEnvironment
+import kotlinx.coroutines.runBlocking
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.ExitCode
 import org.jetbrains.kotlin.cli.common.ExitCode.COMPILATION_ERROR
@@ -16,12 +17,17 @@ import org.jetbrains.kotlin.cli.common.extensions.ScriptEvaluationExtension
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity.ERROR
 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.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.scripting.compiler.plugin.impl.ScriptJvmCompilerFromEnvironment
 import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider
 import java.io.File
+import kotlin.script.experimental.api.constructorArgs
+import kotlin.script.experimental.api.valueOr
+import kotlin.script.experimental.api.with
+import kotlin.script.experimental.host.toScriptSource
+import kotlin.script.experimental.jvm.BasicJvmScriptEvaluator
 
 class JvmCliScriptEvaluationExtension : ScriptEvaluationExtension {
     override fun isAccepted(arguments: CommonCompilerArguments): Boolean =
@@ -42,7 +48,7 @@ class JvmCliScriptEvaluationExtension : ScriptEvaluationExtension {
 
         configuration.addKotlinSourceRoot(sourcePath)
         configuration.put(JVMConfigurationKeys.RETAIN_OUTPUT_IN_MEMORY, true)
-        val coreEnvironment =
+        val environment =
             KotlinCoreEnvironment.createForProduction(projectEnvironment, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES)
 
         val scriptFile = File(sourcePath)
@@ -54,8 +60,32 @@ class JvmCliScriptEvaluationExtension : ScriptEvaluationExtension {
             return COMPILATION_ERROR
         }
 
+        val script = scriptFile.toScriptSource()
+
+        val definition = scriptDefinitionProvider.findDefinition(scriptFile) ?: scriptDefinitionProvider.getDefaultDefinition()
+
         val scriptArgs = arguments.freeArgs.subList(1, arguments.freeArgs.size)
 
-        return KotlinToJVMBytecodeCompiler.compileAndExecuteScript(coreEnvironment, scriptArgs)
+        val evaluationConfiguration = definition.evaluationConfiguration.with {
+            constructorArgs(scriptArgs.toTypedArray())
+        }
+        val scriptCompilationConfiguration = definition.compilationConfiguration
+
+        val scriptCompiler = ScriptJvmCompilerFromEnvironment(environment)
+
+        return runBlocking {
+            val compiledScript =
+                scriptCompiler.compile(script, scriptCompilationConfiguration)
+                    .valueOr { return@runBlocking COMPILATION_ERROR }
+            /*val evalResult = */
+            BasicJvmScriptEvaluator().invoke(compiledScript, evaluationConfiguration)
+                .valueOr {
+                    for (report in it.reports) {
+                        messageCollector.report(ERROR, report.toString())
+                    }
+                    return@runBlocking ExitCode.SCRIPT_EXECUTION_ERROR
+                }
+            ExitCode.OK
+        }
     }
 }

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptJvmCompilerProxy.kt

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2019 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
+
+import kotlin.script.experimental.api.CompiledScript
+import kotlin.script.experimental.api.ResultWithDiagnostics
+import kotlin.script.experimental.api.ScriptCompilationConfiguration
+import kotlin.script.experimental.api.SourceCode
+
+interface ScriptJvmCompilerProxy {
+    fun compile(
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration
+    ): ResultWithDiagnostics<CompiledScript<*>>
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/CompiledScriptClassLoader.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2018 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.impl
+
+import org.jetbrains.kotlin.codegen.BytesUrlUtils
+import java.io.ByteArrayInputStream
+import java.io.InputStream
+import java.net.URL
+import java.util.*
+
+class CompiledScriptClassLoader(parent: ClassLoader?, private val entries: Map<String, ByteArray>) : ClassLoader(parent) {
+
+    override fun findClass(name: String): Class<*>? {
+        val classPathName = name.replace('.', '/') + ".class"
+        val classBytes = entries[classPathName] ?: return null
+        return defineClass(name, classBytes, 0, classBytes.size)
+    }
+
+    override fun getResourceAsStream(name: String): InputStream? =
+        entries[name]?.let(::ByteArrayInputStream) ?: super.getResourceAsStream(name)
+
+    override fun findResources(name: String?): Enumeration<URL>? {
+        val fromParent = super.findResources(name)
+
+        val url = entries[name]?.let { BytesUrlUtils.createBytesUrl(it) } ?: return fromParent
+
+        return Collections.enumeration(listOf(url) + fromParent.asSequence())
+    }
+
+    override fun findResource(name: String?): URL? =
+        entries[name]?.let { BytesUrlUtils.createBytesUrl(it) } ?: super.findResource(name)
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/KJvmCompiledModuleInMemory.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2019 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.impl
+
+import java.io.Serializable
+import kotlin.script.experimental.jvm.impl.KJvmCompiledModule
+
+class KJvmCompiledModuleInMemory(val compilerOutputFiles: Map<String, ByteArray>) : KJvmCompiledModule,
+    Serializable {
+
+    override fun createClassLoader(baseClassLoader: ClassLoader?): ClassLoader =
+        CompiledScriptClassLoader(baseClassLoader, compilerOutputFiles)
+
+    companion object {
+        @JvmStatic
+        private val serialVersionUID = 0L
+    }
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/KJvmReplCompilerImpl.kt

@@ -0,0 +1,174 @@
+/*
+ * Copyright 2010-2019 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.impl
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.cli.common.environment.setIdeaIoUseFallback
+import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
+import org.jetbrains.kotlin.cli.common.messages.MessageCollectorBasedReporter
+import org.jetbrains.kotlin.cli.common.repl.IReplStageHistory
+import org.jetbrains.kotlin.cli.common.repl.LineId
+import org.jetbrains.kotlin.cli.common.repl.ReplCodeLine
+import org.jetbrains.kotlin.cli.common.repl.scriptResultFieldName
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+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.descriptors.ScriptDescriptor
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerState
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.KJvmReplCompilerProxy
+import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+import org.jetbrains.kotlin.scripting.repl.ReplCodeAnalyzer
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+
+class KJvmReplCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvmReplCompilerProxy {
+
+    override fun createReplCompilationState(scriptCompilationConfiguration: ScriptCompilationConfiguration): JvmReplCompilerState.Compilation {
+        val context = withMessageCollectorAndDisposable(disposeOnSuccess = false) { messageCollector, disposable ->
+            createIsolatedCompilationContext(
+                scriptCompilationConfiguration,
+                hostConfiguration,
+                messageCollector,
+                disposable
+            ).asSuccess()
+        }.valueOr { throw IllegalStateException("Unable to initialize repl compiler:\n  ${it.reports.joinToString("\n  ")}") }
+        return ReplCompilationState(context)
+    }
+
+    override fun checkSyntax(
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration,
+        project: Project
+    ): ResultWithDiagnostics<Boolean> =
+        withMessageCollector(script) { messageCollector ->
+            val ktFile = getScriptKtFile(
+                script,
+                scriptCompilationConfiguration,
+                project,
+                messageCollector
+            )
+                .valueOr { return it }
+            val errorHolder = object : MessageCollectorBasedReporter {
+                override val messageCollector = messageCollector
+            }
+            val syntaxErrorReport = AnalyzerWithCompilerReport.reportSyntaxErrors(ktFile, errorHolder)
+            when {
+                syntaxErrorReport.isHasErrors && syntaxErrorReport.isAllErrorsAtEof -> false.asSuccess(messageCollector.diagnostics)
+                syntaxErrorReport.isHasErrors -> failure(messageCollector)
+                else -> true.asSuccess()
+            }
+        }
+
+    override fun compileReplSnippet(
+        compilationState: JvmReplCompilerState.Compilation,
+        snippet: SourceCode,
+        snippetId: ReplSnippetId,
+        // TODO: replace history with some interface based on CompiledScript
+        history: IReplStageHistory<ScriptDescriptor>
+    ): ResultWithDiagnostics<CompiledScript<*>> =
+        withMessageCollector(snippet) { messageCollector ->
+
+            val context = (compilationState as? ReplCompilationState)?.context
+                ?: return failure(
+                    snippet, messageCollector, "Internal error: unknown parameter passed as compilationState: $compilationState"
+                )
+
+            setIdeaIoUseFallback()
+
+            // NOTE: converting between REPL entities from compiler and "new" scripting entities
+            // TODO: (big) move REPL API from compiler to the new scripting infrastructure and streamline ops
+            val codeLine = makeReplCodeLine(snippetId, snippet)
+
+            val errorHolder = object : MessageCollectorBasedReporter {
+                override val messageCollector = messageCollector
+            }
+
+            val snippetKtFile =
+                getScriptKtFile(
+                    snippet,
+                    context.baseScriptCompilationConfiguration,
+                    context.environment.project,
+                    messageCollector
+                )
+                    .valueOr { return it }
+
+            val (sourceFiles, sourceDependencies) = collectRefinedSourcesAndUpdateEnvironment(
+                context,
+                snippetKtFile,
+                messageCollector
+            )
+
+            val analysisResult =
+                compilationState.analyzerEngine.analyzeReplLineWithImportedScripts(snippetKtFile, sourceFiles.drop(1), codeLine)
+            AnalyzerWithCompilerReport.reportDiagnostics(analysisResult.diagnostics, errorHolder)
+
+            val scriptDescriptor = when (analysisResult) {
+                is ReplCodeAnalyzer.ReplLineAnalysisResult.WithErrors -> return failure(
+                    messageCollector
+                )
+                is ReplCodeAnalyzer.ReplLineAnalysisResult.Successful -> analysisResult.scriptDescriptor
+                else -> return failure(
+                    snippet,
+                    messageCollector,
+                    "Unexpected result ${analysisResult::class.java}"
+                )
+            }
+
+            val type = (scriptDescriptor as ScriptDescriptor).resultValue?.returnType
+
+            val generationState = GenerationState.Builder(
+                snippetKtFile.project,
+                ClassBuilderFactories.BINARIES,
+                compilationState.analyzerEngine.module,
+                compilationState.analyzerEngine.trace.bindingContext,
+                sourceFiles,
+                compilationState.environment.configuration
+            ).build().apply {
+                replSpecific.resultType = type
+                replSpecific.scriptResultFieldName = scriptResultFieldName(codeLine.no)
+                replSpecific.earlierScriptsForReplInterpreter = history.map { it.item }
+                beforeCompile()
+            }
+            KotlinCodegenFacade.generatePackage(
+                generationState,
+                snippetKtFile.script!!.containingKtFile.packageFqName,
+                setOf(snippetKtFile.script!!.containingKtFile),
+                CompilationErrorHandler.THROW_EXCEPTION
+            )
+
+            history.push(LineId(codeLine), scriptDescriptor)
+
+            val dependenciesProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
+            val compiledScript =
+                makeCompiledScript(
+                    generationState,
+                    snippet,
+                    sourceFiles.first(),
+                    sourceDependencies
+                ) { ktFile ->
+                    dependenciesProvider?.getScriptConfigurationResult(ktFile)?.valueOrNull()?.configuration
+                        ?: context.baseScriptCompilationConfiguration
+                }
+
+            ResultWithDiagnostics.Success(compiledScript, messageCollector.diagnostics)
+        }
+}
+
+internal class ReplCompilationState(val context: SharedScriptCompilationContext) : JvmReplCompilerState.Compilation {
+    override val disposable: Disposable? get() = context.disposable
+    override val baseScriptCompilationConfiguration: ScriptCompilationConfiguration get() = context.baseScriptCompilationConfiguration
+    override val environment: KotlinCoreEnvironment get() = context.environment
+    override val analyzerEngine: ReplCodeAnalyzer by lazy {
+        ReplCodeAnalyzer(context.environment)
+    }
+}
+
+internal fun makeReplCodeLine(id: ReplSnippetId, code: SourceCode): ReplCodeLine =
+    ReplCodeLine(id.no, id.generation, code.text)
+

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/ScriptJvmCompilerImpls.kt

@@ -0,0 +1,146 @@
+/*
+ * Copyright 2010-2019 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.impl
+
+import org.jetbrains.kotlin.analyzer.AnalysisResult
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
+import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
+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.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.CompilerConfiguration
+import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.scripting.compiler.plugin.ScriptJvmCompilerProxy
+import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+
+class ScriptJvmCompilerIsolated(val hostConfiguration: ScriptingHostConfiguration) : ScriptJvmCompilerProxy {
+
+    override fun compile(
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration
+    ): ResultWithDiagnostics<CompiledScript<*>> =
+        withMessageCollectorAndDisposable(script = script) { messageCollector, disposable ->
+
+            val context = createIsolatedCompilationContext(
+                scriptCompilationConfiguration, hostConfiguration, messageCollector, disposable
+            )
+
+            compileImpl(script, context, messageCollector)
+        }
+}
+
+class ScriptJvmCompilerFromEnvironment(val environment: KotlinCoreEnvironment) : ScriptJvmCompilerProxy {
+
+    override fun compile(
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration
+    ): ResultWithDiagnostics<CompiledScript<*>> {
+        val parentMessageCollector = environment.configuration[CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY]
+        return withMessageCollector(script = script, parentMessageCollector = parentMessageCollector) { messageCollector ->
+
+            val context = createCompilationContextFromEnvironment(scriptCompilationConfiguration, environment, messageCollector)
+
+            try {
+                environment.configuration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, messageCollector)
+
+                compileImpl(script, context, messageCollector)
+
+            } finally {
+                if (parentMessageCollector != null)
+                    environment.configuration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, parentMessageCollector)
+            }
+        }
+    }
+}
+
+private fun compileImpl(
+    script: SourceCode,
+    context: SharedScriptCompilationContext,
+    messageCollector: ScriptDiagnosticsMessageCollector
+): ResultWithDiagnostics<CompiledScript<*>> {
+    val mainKtFile =
+        getScriptKtFile(
+            script,
+            context.baseScriptCompilationConfiguration,
+            context.environment.project,
+            messageCollector
+        )
+            .valueOr { return it }
+
+    val (sourceFiles, sourceDependencies) = collectRefinedSourcesAndUpdateEnvironment(
+        context,
+        mainKtFile,
+        messageCollector
+    )
+
+    val dependenciesProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
+    val analysisResult = analyze(sourceFiles, context.environment)
+
+    if (!analysisResult.shouldGenerateCode) return failure(
+        script,
+        messageCollector,
+        "no code to generate"
+    )
+    if (analysisResult.isError() || messageCollector.hasErrors()) return failure(
+        messageCollector
+    )
+
+    val generationState =
+        generate(analysisResult, sourceFiles, context.environment.configuration)
+
+    val compiledScript =
+        makeCompiledScript(
+            generationState,
+            script,
+            sourceFiles.first(),
+            sourceDependencies
+        ) { ktFile ->
+            dependenciesProvider?.getScriptConfigurationResult(ktFile)?.valueOrNull()?.configuration
+                ?: context.baseScriptCompilationConfiguration
+        }
+
+    return ResultWithDiagnostics.Success(compiledScript, messageCollector.diagnostics)
+}
+
+private fun analyze(sourceFiles: Collection<KtFile>, environment: KotlinCoreEnvironment): AnalysisResult {
+    val messageCollector = environment.configuration[CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY]!!
+
+    val analyzerWithCompilerReport = AnalyzerWithCompilerReport(messageCollector, environment.configuration.languageVersionSettings)
+
+    analyzerWithCompilerReport.analyzeAndReport(sourceFiles) {
+        val project = environment.project
+        TopDownAnalyzerFacadeForJVM.analyzeFilesWithJavaIntegration(
+            project,
+            sourceFiles,
+            NoScopeRecordCliBindingTrace(),
+            environment.configuration,
+            environment::createPackagePartProvider
+        )
+    }
+    return analyzerWithCompilerReport.analysisResult
+}
+
+private fun generate(
+    analysisResult: AnalysisResult, sourceFiles: List<KtFile>, kotlinCompilerConfiguration: CompilerConfiguration
+): GenerationState {
+    val generationState = GenerationState.Builder(
+        sourceFiles.first().project,
+        ClassBuilderFactories.BINARIES,
+        analysisResult.moduleDescriptor,
+        analysisResult.bindingContext,
+        sourceFiles,
+        kotlinCompilerConfiguration
+    ).build()
+
+    KotlinCodegenFacade.compileCorrectFiles(generationState, CompilationErrorHandler.THROW_EXCEPTION)
+    return generationState
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/compilationContext.kt

@@ -0,0 +1,256 @@
+/*
+ * Copyright 2010-2019 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.impl
+
+import com.intellij.openapi.Disposable
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
+import org.jetbrains.kotlin.cli.common.arguments.K2JVMCompilerArguments
+import org.jetbrains.kotlin.cli.common.arguments.parseCommandLineArguments
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+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.JvmClasspathRoot
+import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
+import org.jetbrains.kotlin.cli.jvm.config.jvmClasspathRoots
+import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
+import org.jetbrains.kotlin.config.CommonConfigurationKeys
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.config.JvmTarget
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.scripting.compiler.plugin.ScriptingCompilerConfigurationComponentRegistrar
+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.dependencies.ScriptsCompilationDependencies
+import org.jetbrains.kotlin.scripting.dependencies.collectScriptsCompilationDependencies
+import kotlin.script.experimental.api.ScriptCompilationConfiguration
+import kotlin.script.experimental.api.compilerOptions
+import kotlin.script.experimental.api.dependencies
+import kotlin.script.experimental.api.valueOrNull
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+import kotlin.script.experimental.jvm.JvmDependency
+import kotlin.script.experimental.jvm.jdkHome
+import kotlin.script.experimental.jvm.jvm
+import kotlin.script.experimental.jvm.util.KotlinJars
+import kotlin.script.experimental.jvm.withUpdatedClasspath
+
+internal class SharedScriptCompilationContext(
+    val disposable: Disposable?,
+    val baseScriptCompilationConfiguration: ScriptCompilationConfiguration,
+    val environment: KotlinCoreEnvironment,
+    val ignoredOptionsReportingState: IgnoredOptionsReportingState
+)
+
+internal fun createIsolatedCompilationContext(
+    baseScriptCompilationConfiguration: ScriptCompilationConfiguration,
+    hostConfiguration: ScriptingHostConfiguration,
+    messageCollector: ScriptDiagnosticsMessageCollector,
+    disposable: Disposable
+): SharedScriptCompilationContext {
+    val ignoredOptionsReportingState = IgnoredOptionsReportingState()
+
+    val (initialScriptCompilationConfiguration, kotlinCompilerConfiguration) =
+        createInitialConfigurations(baseScriptCompilationConfiguration, hostConfiguration, messageCollector, ignoredOptionsReportingState)
+    val environment =
+        KotlinCoreEnvironment.createForProduction(
+            disposable, kotlinCompilerConfiguration, EnvironmentConfigFiles.JVM_CONFIG_FILES
+        )
+
+    return SharedScriptCompilationContext(
+        disposable, initialScriptCompilationConfiguration, environment, ignoredOptionsReportingState
+    )
+}
+
+internal fun createCompilationContextFromEnvironment(
+    baseScriptCompilationConfiguration: ScriptCompilationConfiguration,
+    environment: KotlinCoreEnvironment,
+    messageCollector: ScriptDiagnosticsMessageCollector
+): SharedScriptCompilationContext {
+    val ignoredOptionsReportingState = IgnoredOptionsReportingState()
+
+    val initialScriptCompilationConfiguration =
+        baseScriptCompilationConfiguration.withUpdatesFromCompilerConfiguration(environment.configuration)
+
+    initialScriptCompilationConfiguration[ScriptCompilationConfiguration.compilerOptions]?.let { compilerOptions ->
+        environment.configuration.updateWithCompilerOptions(compilerOptions, messageCollector, ignoredOptionsReportingState, false)
+    }
+
+    return SharedScriptCompilationContext(
+        null, initialScriptCompilationConfiguration, environment, ignoredOptionsReportingState
+    )
+}
+
+internal fun createInitialConfigurations(
+    scriptCompilationConfiguration: ScriptCompilationConfiguration,
+    hostConfiguration: ScriptingHostConfiguration,
+    messageCollector: ScriptDiagnosticsMessageCollector,
+    ignoredOptionsReportingState: IgnoredOptionsReportingState
+): Pair<ScriptCompilationConfiguration, CompilerConfiguration> {
+    val kotlinCompilerConfiguration =
+        createInitialCompilerConfiguration(
+            scriptCompilationConfiguration, hostConfiguration, messageCollector, ignoredOptionsReportingState
+        )
+
+    val initialScriptCompilationConfiguration =
+        scriptCompilationConfiguration.withUpdatesFromCompilerConfiguration(kotlinCompilerConfiguration)
+
+    kotlinCompilerConfiguration.add(
+        ScriptingConfigurationKeys.SCRIPT_DEFINITIONS,
+        ScriptDefinition.FromConfigurations(hostConfiguration, scriptCompilationConfiguration, null)
+    )
+
+    initialScriptCompilationConfiguration[ScriptCompilationConfiguration.compilerOptions]?.let { compilerOptions ->
+        kotlinCompilerConfiguration.updateWithCompilerOptions(compilerOptions, messageCollector, ignoredOptionsReportingState, false)
+    }
+
+    return Pair(initialScriptCompilationConfiguration, kotlinCompilerConfiguration)
+}
+
+private fun CompilerConfiguration.updateWithCompilerOptions(
+    compilerOptions: List<String>,
+    messageCollector: ScriptDiagnosticsMessageCollector,
+    ignoredOptionsReportingState: IgnoredOptionsReportingState,
+    isRefinement: Boolean
+) {
+    val compilerArguments = K2JVMCompilerArguments()
+    parseCommandLineArguments(compilerOptions, compilerArguments)
+
+    reportArgumentsIgnoredGenerally(
+        compilerArguments,
+        messageCollector,
+        ignoredOptionsReportingState
+    )
+    if (isRefinement) {
+        reportArgumentsIgnoredFromRefinement(
+            compilerArguments,
+            messageCollector,
+            ignoredOptionsReportingState
+        )
+    }
+
+    setupCommonArguments(compilerArguments)
+
+    setupJvmSpecificArguments(compilerArguments)
+
+    configureAdvancedJvmOptions(compilerArguments)
+}
+
+private fun ScriptCompilationConfiguration.withUpdatesFromCompilerConfiguration(kotlinCompilerConfiguration: CompilerConfiguration) =
+    withUpdatedClasspath(kotlinCompilerConfiguration.jvmClasspathRoots)
+
+private fun createInitialCompilerConfiguration(
+    scriptCompilationConfiguration: ScriptCompilationConfiguration,
+    hostConfiguration: ScriptingHostConfiguration,
+    messageCollector: MessageCollector,
+    reportingState: IgnoredOptionsReportingState
+): CompilerConfiguration {
+
+    val baseArguments = K2JVMCompilerArguments()
+    parseCommandLineArguments(
+        scriptCompilationConfiguration[ScriptCompilationConfiguration.compilerOptions] ?: emptyList(),
+        baseArguments
+    )
+
+    reportArgumentsIgnoredGenerally(baseArguments, messageCollector, reportingState)
+    reportingState.currentArguments = baseArguments
+
+    return CompilerConfiguration().apply {
+        put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, messageCollector)
+        setupCommonArguments(baseArguments)
+
+        setupJvmSpecificArguments(baseArguments)
+
+        // Default value differs from the argument's default (see #KT-29405 and #KT-29319)
+        put(JVMConfigurationKeys.JVM_TARGET, JvmTarget.JVM_1_8)
+
+        val jdkHomeFromConfigurations = scriptCompilationConfiguration.getNoDefault(ScriptCompilationConfiguration.jvm.jdkHome)
+            ?: hostConfiguration[ScriptingHostConfiguration.jvm.jdkHome]
+        if (jdkHomeFromConfigurations != null) {
+            messageCollector.report(CompilerMessageSeverity.LOGGING, "Using JDK home directory $jdkHomeFromConfigurations")
+            put(JVMConfigurationKeys.JDK_HOME, jdkHomeFromConfigurations)
+        } else {
+            configureJdkHome(baseArguments)
+        }
+
+        put(JVMConfigurationKeys.RETAIN_OUTPUT_IN_MEMORY, true)
+
+        val isModularJava = isModularJava()
+
+        scriptCompilationConfiguration[ScriptCompilationConfiguration.dependencies]?.let { dependencies ->
+            addJvmClasspathRoots(
+                dependencies.flatMap {
+                    (it as JvmDependency).classpath
+                }
+            )
+        }
+
+        add(
+            ComponentRegistrar.PLUGIN_COMPONENT_REGISTRARS,
+            ScriptingCompilerConfigurationComponentRegistrar()
+        )
+
+        configureExplicitContentRoots(baseArguments)
+
+        if (!baseArguments.noStdlib) {
+            addModularRootIfNotNull(isModularJava, "kotlin.stdlib", KotlinJars.stdlib)
+            addModularRootIfNotNull(isModularJava, "kotlin.script.runtime", KotlinJars.scriptRuntimeOrNull)
+        }
+        // see comments about logic in CompilerConfiguration.configureStandardLibs
+        if (!baseArguments.noReflect && !baseArguments.noStdlib) {
+            addModularRootIfNotNull(isModularJava, "kotlin.reflect", KotlinJars.reflectOrNull)
+        }
+
+        put(CommonConfigurationKeys.MODULE_NAME, baseArguments.moduleName ?: "kotlin-script")
+
+        configureAdvancedJvmOptions(baseArguments)
+    }
+}
+
+internal fun collectRefinedSourcesAndUpdateEnvironment(
+    context: SharedScriptCompilationContext,
+    mainKtFile: KtFile,
+    messageCollector: ScriptDiagnosticsMessageCollector
+): Pair<List<KtFile>, List<ScriptsCompilationDependencies.SourceDependencies>> {
+    val sourceFiles = arrayListOf(mainKtFile)
+    val (classpath, newSources, sourceDependencies) =
+        collectScriptsCompilationDependencies(
+            context.environment.configuration,
+            context.environment.project,
+            sourceFiles
+        )
+
+    // TODO: consider removing, it is probably redundant: the actual index update is performed with environment.updateClasspath
+    context.environment.configuration.addJvmClasspathRoots(classpath)
+    context.environment.updateClasspath(classpath.map(::JvmClasspathRoot))
+
+    sourceFiles.addAll(newSources)
+
+    // collectScriptsCompilationDependencies calls resolver for every file, so at this point all updated configurations are collected in the ScriptDependenciesProvider
+    context.environment.configuration.updateWithRefinedConfigurations(context, sourceFiles, messageCollector)
+    return sourceFiles to sourceDependencies
+}
+
+private fun CompilerConfiguration.updateWithRefinedConfigurations(
+    context: SharedScriptCompilationContext,
+    sourceFiles: List<KtFile>,
+    messageCollector: ScriptDiagnosticsMessageCollector
+) {
+    val dependenciesProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
+    val updatedCompilerOptions = sourceFiles.flatMap {
+        dependenciesProvider?.getScriptConfigurationResult(it)?.valueOrNull()?.configuration?.get(
+            ScriptCompilationConfiguration.compilerOptions
+        ) ?: emptyList()
+    }
+    if (updatedCompilerOptions.isNotEmpty() &&
+        updatedCompilerOptions != context.baseScriptCompilationConfiguration[ScriptCompilationConfiguration.compilerOptions]
+    ) {
+        updateWithCompilerOptions(updatedCompilerOptions, messageCollector, context.ignoredOptionsReportingState, true)
+    }
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/errorReporting.kt

@@ -0,0 +1,136 @@
+/*
+ * Copyright 2010-2019 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.impl
+
+import org.jetbrains.kotlin.cli.common.arguments.Argument
+import org.jetbrains.kotlin.cli.common.arguments.K2JVMCompilerArguments
+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 kotlin.reflect.KMutableProperty1
+import kotlin.reflect.full.findAnnotation
+import kotlin.script.experimental.api.ResultWithDiagnostics
+import kotlin.script.experimental.api.ScriptDiagnostic
+import kotlin.script.experimental.api.SourceCode
+import kotlin.script.experimental.api.asErrorDiagnostics
+
+internal class ScriptDiagnosticsMessageCollector(private val parentMessageCollector: MessageCollector?) : MessageCollector {
+
+    private val _diagnostics = arrayListOf<ScriptDiagnostic>()
+
+    val diagnostics: List<ScriptDiagnostic> get() = _diagnostics
+
+    override fun clear() {
+        _diagnostics.clear()
+        parentMessageCollector?.clear()
+    }
+
+    override fun hasErrors(): Boolean =
+        _diagnostics.any { it.severity == ScriptDiagnostic.Severity.ERROR } || parentMessageCollector?.hasErrors() == true
+
+    override fun report(severity: CompilerMessageSeverity, message: String, location: CompilerMessageLocation?) {
+        val mappedSeverity = when (severity) {
+            CompilerMessageSeverity.EXCEPTION,
+            CompilerMessageSeverity.ERROR -> ScriptDiagnostic.Severity.ERROR
+            CompilerMessageSeverity.STRONG_WARNING,
+            CompilerMessageSeverity.WARNING -> ScriptDiagnostic.Severity.WARNING
+            CompilerMessageSeverity.INFO -> ScriptDiagnostic.Severity.INFO
+            CompilerMessageSeverity.LOGGING -> ScriptDiagnostic.Severity.DEBUG
+            else -> null
+        }
+        if (mappedSeverity != null) {
+            val mappedLocation = location?.let {
+                if (it.line < 0 && it.column < 0) null // special location created by CompilerMessageLocation.create
+                else SourceCode.Location(
+                    SourceCode.Position(
+                        it.line,
+                        it.column
+                    )
+                )
+            }
+            _diagnostics.add(ScriptDiagnostic(message, mappedSeverity, location?.path, mappedLocation))
+        }
+        parentMessageCollector?.report(severity, message, location)
+    }
+}
+
+internal fun failure(
+    messageCollector: ScriptDiagnosticsMessageCollector, vararg diagnostics: ScriptDiagnostic
+): ResultWithDiagnostics.Failure =
+    ResultWithDiagnostics.Failure(*messageCollector.diagnostics.toTypedArray(), *diagnostics)
+
+internal fun failure(
+    script: SourceCode, messageCollector: ScriptDiagnosticsMessageCollector, message: String
+): ResultWithDiagnostics.Failure =
+    failure(messageCollector, message.asErrorDiagnostics(path = script.locationId))
+
+internal class IgnoredOptionsReportingState {
+    var currentArguments = K2JVMCompilerArguments()
+}
+
+internal fun reportArgumentsIgnoredGenerally(
+    arguments: K2JVMCompilerArguments,
+    messageCollector: MessageCollector,
+    reportingState: IgnoredOptionsReportingState
+) {
+
+    reportIgnoredArguments(
+        arguments,
+        "The following compiler arguments are ignored on script compilation: ",
+        messageCollector,
+        reportingState,
+        K2JVMCompilerArguments::version,
+        K2JVMCompilerArguments::destination,
+        K2JVMCompilerArguments::buildFile,
+        K2JVMCompilerArguments::commonSources,
+        K2JVMCompilerArguments::allWarningsAsErrors,
+        K2JVMCompilerArguments::script,
+        K2JVMCompilerArguments::scriptTemplates,
+        K2JVMCompilerArguments::scriptResolverEnvironment,
+        K2JVMCompilerArguments::disableStandardScript,
+        K2JVMCompilerArguments::disableDefaultScriptingPlugin,
+        K2JVMCompilerArguments::pluginClasspaths,
+        K2JVMCompilerArguments::pluginOptions,
+        K2JVMCompilerArguments::useJavac,
+        K2JVMCompilerArguments::compileJava,
+        K2JVMCompilerArguments::reportPerf,
+        K2JVMCompilerArguments::dumpPerf
+    )
+}
+
+internal fun reportArgumentsIgnoredFromRefinement(
+    arguments: K2JVMCompilerArguments, messageCollector: MessageCollector, reportingState: IgnoredOptionsReportingState
+) {
+    reportIgnoredArguments(
+        arguments,
+        "The following compiler arguments are ignored when configured from refinement callbacks: ",
+        messageCollector,
+        reportingState,
+        K2JVMCompilerArguments::noJdk,
+        K2JVMCompilerArguments::jdkHome,
+        K2JVMCompilerArguments::javaModulePath,
+        K2JVMCompilerArguments::classpath,
+        K2JVMCompilerArguments::noStdlib,
+        K2JVMCompilerArguments::noReflect
+    )
+}
+
+private fun reportIgnoredArguments(
+    arguments: K2JVMCompilerArguments, message: String,
+    messageCollector: MessageCollector, reportingState: IgnoredOptionsReportingState,
+    vararg toIgnore: KMutableProperty1<K2JVMCompilerArguments, *>
+) {
+    val ignoredArgKeys = toIgnore.mapNotNull { argProperty ->
+        if (argProperty.get(arguments) != argProperty.get(reportingState.currentArguments)) {
+            argProperty.findAnnotation<Argument>()?.value
+                ?: throw IllegalStateException("unknown compiler argument property: $argProperty: no Argument annotation found")
+        } else null
+    }
+
+    if (ignoredArgKeys.isNotEmpty()) {
+        messageCollector.report(CompilerMessageSeverity.STRONG_WARNING, "$message${ignoredArgKeys.joinToString(", ")}")
+    }
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/jvmCompilationUtil.kt

@@ -0,0 +1,165 @@
+/*
+ * Copyright 2010-2019 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.impl
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
+import com.intellij.psi.PsiFileFactory
+import com.intellij.psi.impl.PsiFileFactoryImpl
+import org.jetbrains.kotlin.cli.common.environment.setIdeaIoUseFallback
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtScript
+import org.jetbrains.kotlin.renderer.DescriptorRenderer
+import org.jetbrains.kotlin.scripting.dependencies.ScriptsCompilationDependencies
+import org.jetbrains.kotlin.scripting.resolve.ScriptLightVirtualFile
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+import java.util.*
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.FileBasedScriptSource
+import kotlin.script.experimental.host.getMergedScriptText
+import kotlin.script.experimental.jvm.impl.KJvmCompiledScript
+
+internal fun makeCompiledModule(generationState: GenerationState) =
+    KJvmCompiledModuleInMemory(
+        generationState.factory.asList()
+            .associateTo(sortedMapOf<String, ByteArray>()) { it.relativePath to it.asByteArray() }
+    )
+
+internal fun SourceCode.scriptFileName(
+    mainScript: SourceCode,
+    scriptCompilationConfiguration: ScriptCompilationConfiguration
+): String =
+    when {
+        name != null -> name!!
+        mainScript == this -> "script.${scriptCompilationConfiguration[ScriptCompilationConfiguration.fileExtension]}"
+        else -> throw Exception("Unexpected script without name: $this")
+    }
+
+internal inline fun <T> withMessageCollectorAndDisposable(
+    script: SourceCode? = null,
+    parentMessageCollector: MessageCollector? = null,
+    disposable: Disposable = Disposer.newDisposable(),
+    disposeOnSuccess: Boolean = true,
+    body: (ScriptDiagnosticsMessageCollector, Disposable) -> ResultWithDiagnostics<T>
+): ResultWithDiagnostics<T> {
+    var failed = false
+    val messageCollector = ScriptDiagnosticsMessageCollector(parentMessageCollector)
+    return try {
+        setIdeaIoUseFallback()
+        body(messageCollector, disposable).also {
+            failed = it is ResultWithDiagnostics.Failure
+        }
+    } catch (ex: Throwable) {
+        failed = true
+        failure(messageCollector, ex.asDiagnostics(path = script?.locationId))
+    } finally {
+        if (disposeOnSuccess || failed) {
+            Disposer.dispose(disposable)
+        }
+    }
+}
+
+internal inline fun <T> withMessageCollector(
+    script: SourceCode? = null,
+    parentMessageCollector: MessageCollector? = null,
+    body: (ScriptDiagnosticsMessageCollector) -> ResultWithDiagnostics<T>
+): ResultWithDiagnostics<T> {
+    val messageCollector = ScriptDiagnosticsMessageCollector(parentMessageCollector)
+    return try {
+        body(messageCollector)
+    } catch (ex: Throwable) {
+        failure(messageCollector, ex.asDiagnostics(path = script?.locationId))
+    }
+}
+
+internal fun getScriptKtFile(
+    script: SourceCode,
+    scriptCompilationConfiguration: ScriptCompilationConfiguration,
+    project: Project,
+    messageCollector: ScriptDiagnosticsMessageCollector
+): ResultWithDiagnostics<KtFile> {
+    val psiFileFactory: PsiFileFactoryImpl = PsiFileFactory.getInstance(project) as PsiFileFactoryImpl
+    val scriptText = getMergedScriptText(script, scriptCompilationConfiguration)
+    val virtualFile = ScriptLightVirtualFile(
+        script.scriptFileName(script, scriptCompilationConfiguration),
+        (script as? FileBasedScriptSource)?.file?.path,
+        scriptText
+    )
+    val ktFile = psiFileFactory.trySetupPsiForFile(virtualFile, KotlinLanguage.INSTANCE, true, false) as KtFile?
+    return when {
+        ktFile == null -> failure(
+            script,
+            messageCollector,
+            "Unable to make PSI file from script"
+        )
+        ktFile.declarations.firstIsInstanceOrNull<KtScript>() == null -> failure(
+            script,
+            messageCollector,
+            "Not a script file"
+        )
+        else -> ktFile.asSuccess()
+    }
+}
+
+internal fun makeCompiledScript(
+    generationState: GenerationState,
+    script: SourceCode,
+    ktFile: KtFile,
+    sourceDependencies: List<ScriptsCompilationDependencies.SourceDependencies>,
+    getScriptConfiguration: (KtFile) -> ScriptCompilationConfiguration
+): KJvmCompiledScript<Any> {
+    val scriptDependenciesStack = ArrayDeque<KtScript>()
+    val ktScript = ktFile.declarations.firstIsInstanceOrNull<KtScript>()
+        ?: throw IllegalStateException("Expecting script file: KtScript is not found in ${ktFile.name}")
+
+    fun makeOtherScripts(script: KtScript): List<KJvmCompiledScript<*>> {
+
+        // 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>(
+                        containingKtFile.virtualFile?.path,
+                        getScriptConfiguration(sourceFile),
+                        it.fqName.asString(),
+                        null,
+                        makeOtherScripts(it),
+                        null
+                    )
+                }
+            } ?: emptyList()
+
+        scriptDependenciesStack.pop()
+        return otherScripts
+    }
+
+    val module = makeCompiledModule(generationState)
+
+    val resultField = with(generationState.replSpecific) {
+        // TODO: pass it in the configuration instead
+        if (!hasResult || resultType == null || scriptResultFieldName == null) null
+        else scriptResultFieldName!! to KotlinType(DescriptorRenderer.FQ_NAMES_IN_TYPES.renderType(resultType!!))
+    }
+
+    return KJvmCompiledScript(
+        script.locationId,
+        getScriptConfiguration(ktScript.containingKtFile),
+        ktScript.fqName.asString(),
+        resultField,
+        makeOtherScripts(ktScript),
+        module
+    )
+}
+

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/jvmReplCompilation.kt

@@ -0,0 +1,109 @@
+/*
+ * Copyright 2010-2019 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.repl
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
+import org.jetbrains.kotlin.cli.common.repl.*
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.descriptors.ScriptDescriptor
+import org.jetbrains.kotlin.scripting.repl.ReplCodeAnalyzer
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import kotlin.concurrent.write
+import kotlin.script.experimental.api.*
+
+interface KJvmReplCompilerProxy {
+    fun createReplCompilationState(scriptCompilationConfiguration: ScriptCompilationConfiguration): JvmReplCompilerState.Compilation
+
+    fun checkSyntax(
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration,
+        project: Project
+    ): ResultWithDiagnostics<Boolean>
+
+
+    fun compileReplSnippet(
+        compilationState: JvmReplCompilerState.Compilation,
+        snippet: SourceCode,
+        snippetId: ReplSnippetId,
+        history: IReplStageHistory<ScriptDescriptor>
+    ): ResultWithDiagnostics<CompiledScript<*>>
+}
+
+class JvmReplCompilerStageHistory(private val state: JvmReplCompilerState) :
+    BasicReplStageHistory<ScriptDescriptor>(state.lock) {
+
+    override fun reset(): Iterable<ILineId> {
+        val removedCompiledLines = super.reset()
+        val removedAnalyzedLines = state.compilation.analyzerEngine.reset()
+
+        checkConsistent(removedCompiledLines, removedAnalyzedLines)
+        return removedCompiledLines
+    }
+
+    override fun resetTo(id: ILineId): Iterable<ILineId> {
+        val removedCompiledLines = super.resetTo(id)
+        val removedAnalyzedLines = state.compilation.analyzerEngine.resetToLine(id)
+
+        checkConsistent(removedCompiledLines, removedAnalyzedLines)
+        return removedCompiledLines
+    }
+
+    private fun checkConsistent(removedCompiledLines: Iterable<ILineId>, removedAnalyzedLines: List<ReplCodeLine>) {
+        removedCompiledLines.zip(removedAnalyzedLines).forEach { (removedCompiledLine, removedAnalyzedLine) ->
+            if (removedCompiledLine != LineId(removedAnalyzedLine)) {
+                throw IllegalStateException("History mismatch when resetting lines: ${removedCompiledLine.no} != $removedAnalyzedLine")
+            }
+        }
+    }
+}
+
+class JvmReplCompilerState(
+    val replCompilerProxy: KJvmReplCompilerProxy,
+    override val lock: ReentrantReadWriteLock = ReentrantReadWriteLock()
+) : IReplStageState<ScriptDescriptor> {
+
+    override val history = JvmReplCompilerStageHistory(this)
+
+    override val currentGeneration: Int get() = (history as BasicReplStageHistory<*>).currentGeneration.get()
+
+    override fun dispose() {
+        lock.write {
+            _compilation?.disposable?.let {
+                Disposer.dispose(it)
+            }
+            _compilation = null
+            super.dispose()
+        }
+    }
+
+    fun getCompilationState(scriptCompilationConfiguration: ScriptCompilationConfiguration): Compilation = lock.write {
+        if (_compilation == null) {
+            initializeCompilation(scriptCompilationConfiguration)
+        }
+        _compilation!!
+    }
+
+    internal val compilation: Compilation
+        get() = _compilation ?: throw IllegalStateException("Compilation state is either not initializad or already destroyed")
+
+    private var _compilation: Compilation? = null
+
+    val isCompilationInitialized get() = _compilation != null
+
+    private fun initializeCompilation(scriptCompilationConfiguration: ScriptCompilationConfiguration) {
+        if (_compilation != null) throw IllegalStateException("Compilation state is already initialized")
+        _compilation = replCompilerProxy.createReplCompilationState(scriptCompilationConfiguration)
+    }
+
+    interface Compilation {
+        val disposable: Disposable?
+        val baseScriptCompilationConfiguration: ScriptCompilationConfiguration
+        val environment: KotlinCoreEnvironment
+        val analyzerEngine: ReplCodeAnalyzer
+    }
+}