Make state of the default REPL compiler explicit and replaceable

restores proper object hierarchy in case the legacy REPL infrastructure

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

@@ -8,7 +8,10 @@ package org.jetbrains.kotlin.scripting.compiler.plugin.impl
 
 import com.intellij.openapi.Disposable
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
-import org.jetbrains.kotlin.backend.jvm.*
+import org.jetbrains.kotlin.backend.jvm.JvmGeneratorExtensionsImpl
+import org.jetbrains.kotlin.backend.jvm.JvmIrCodegenFactory
+import org.jetbrains.kotlin.backend.jvm.JvmNameProvider
+import org.jetbrains.kotlin.backend.jvm.jvmPhases
 import org.jetbrains.kotlin.backend.jvm.serialization.JvmIdSignatureDescriptor
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.environment.setIdeaIoUseFallback
@@ -30,11 +33,12 @@ import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade
-import org.jetbrains.kotlin.scripting.compiler.plugin.repl.*
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerState
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplCodeAnalyzerBase
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicits
 import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicitsExceptInnermost
-import java.util.concurrent.atomic.AtomicInteger
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
 import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
@@ -43,49 +47,30 @@ import kotlin.script.experimental.util.LinkedSnippet
 import kotlin.script.experimental.util.LinkedSnippetImpl
 import kotlin.script.experimental.util.add
 
-open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected constructor(
+// NOTE: this implementation, as it is used in the REPL infrastructure, may be created for every snippet and provided with the state
+// so it should not keep any compilation state outside of the stste field
+open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase>(
     protected val hostConfiguration: ScriptingHostConfiguration = defaultJvmScriptingHostConfiguration,
-    val initAnalyzer: (SharedScriptCompilationContext, ImplicitsExtensionsResolutionFilter) -> AnalyzerT
+    val state: JvmReplCompilerState<*> = JvmReplCompilerState({ createCompilationState(it, hostConfiguration) })
 ) : ReplCompiler<KJvmCompiledScript>, ScriptCompiler {
-    val state = JvmReplCompilerState({ createReplCompilationState(it, initAnalyzer) })
-    val history = JvmReplCompilerStageHistory(state)
-    protected val scriptPriority = AtomicInteger()
-    private val resolutionFilter = ReplImplicitsExtensionsResolutionFilter()
 
     override var lastCompiledSnippet: LinkedSnippetImpl<KJvmCompiledScript>? = null
         protected set
 
-    fun createReplCompilationState(
-        scriptCompilationConfiguration: ScriptCompilationConfiguration,
-        initAnalyzer: (SharedScriptCompilationContext, ImplicitsExtensionsResolutionFilter) -> AnalyzerT /* = { ReplCodeAnalyzer1(it.environment) } */
-    ): ReplCompilationState<AnalyzerT> {
-        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  ")}") }
-
-        updateResolutionFilter(scriptCompilationConfiguration)
-
-        return ReplCompilationState(context, initAnalyzer, resolutionFilter)
-    }
-
     override suspend fun compile(
         snippets: Iterable<SourceCode>,
         configuration: ScriptCompilationConfiguration
     ): ResultWithDiagnostics<LinkedSnippet<KJvmCompiledScript>> =
         snippets.map { snippet ->
             withMessageCollector(snippet) { messageCollector ->
-                updateResolutionFilter(configuration)
-
                 val initialConfiguration = configuration.refineBeforeParsing(snippet).valueOr {
                     return it
                 }
 
-                val compilationState = state.getCompilationState(initialConfiguration)
+                @Suppress("UNCHECKED_CAST")
+                val compilationState = state.getCompilationState(initialConfiguration) as ReplCompilationState<AnalyzerT>
+
+                updateResolutionFilterWithHistory(configuration)
 
                 val (context, errorHolder, snippetKtFile) = prepareForAnalyze(
                     snippet,
@@ -109,7 +94,7 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
 
                 if (messageCollector.hasErrors()) return failure(messageCollector)
 
-                if (history.isEmpty()) {
+                if (state.history.isEmpty()) {
                     val updatedConfiguration = ScriptDependenciesProvider.getInstance(context.environment.project)
                         ?.getScriptConfiguration(snippetKtFile)?.configuration
                         ?: context.baseScriptCompilationConfiguration
@@ -119,7 +104,7 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
                     )
                 }
 
-                val snippetNo = scriptPriority.getAndIncrement()
+                val snippetNo = state.getNextLineNo()
 
                 val analysisResult =
                     compilationState.analyzerEngine.analyzeReplLineWithImportedScripts(
@@ -147,7 +132,7 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
                     generateWithOldBackend(compilationState, snippetKtFile, sourceFiles)
                 }
 
-                history.push(LineId(snippetNo, 0, snippet.hashCode()), scriptDescriptor)
+                state.history.push(LineId(snippetNo, 0, snippet.hashCode()), scriptDescriptor)
 
                 val dependenciesProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
                 makeCompiledScript(
@@ -185,7 +170,7 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
             sourceFiles,
             compilationState.environment.configuration
         ).build().also { generationState ->
-            generationState.scriptSpecific.earlierScriptsForReplInterpreter = history.map { it.item }
+            generationState.scriptSpecific.earlierScriptsForReplInterpreter = state.history.map { it.item }
             generationState.beforeCompile()
         }
         KotlinCodegenFacade.generatePackage(generationState, snippetKtFile.script!!.containingKtFile.packageFqName, sourceFiles)
@@ -199,7 +184,7 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
         sourceFiles: List<KtFile>
     ): GenerationState {
         val generatorExtensions = object : JvmGeneratorExtensionsImpl(compilationState.environment.configuration) {
-            override fun getPreviousScripts() = history.map { compilationState.symbolTable.referenceScript(it.item) }
+            override fun getPreviousScripts() = state.history.map { compilationState.symbolTable.referenceScript(it.item) }
         }
         val codegenFactory = JvmIrCodegenFactory(
             compilationState.environment.configuration,
@@ -282,7 +267,7 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
             ).asSuccess()
         }
 
-    protected fun updateResolutionFilter(configuration: ScriptCompilationConfiguration) {
+    protected fun updateResolutionFilterWithHistory(configuration: ScriptCompilationConfiguration) {
         val updatedConfiguration = updateConfigurationWithPreviousScripts(configuration)
 
         val classesToSkip =
@@ -290,9 +275,10 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
         val classesToSkipAfterFirstTime =
             updatedConfiguration[ScriptCompilationConfiguration.skipExtensionsResolutionForImplicitsExceptInnermost]!!
 
-        resolutionFilter.update(classesToSkip, classesToSkipAfterFirstTime)
+        (state.compilation as ReplCompilationState<*>).implicitsResolutionFilter.update(classesToSkip, classesToSkipAfterFirstTime)
     }
 
+
     private fun updateConfigurationWithPreviousScripts(
         configuration: ScriptCompilationConfiguration
     ): ScriptCompilationConfiguration {
@@ -313,23 +299,51 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
     }
 
     companion object {
-        fun create(hostConfiguration: ScriptingHostConfiguration = defaultJvmScriptingHostConfiguration) =
-            KJvmReplCompilerBase(hostConfiguration) { context, resolutionFilter ->
-                ReplCodeAnalyzerBase(context.environment, implicitsResolutionFilter = resolutionFilter)
-            }
-    }
 
+        fun createCompilationState(
+            scriptCompilationConfiguration: ScriptCompilationConfiguration,
+            hostConfiguration: ScriptingHostConfiguration = defaultJvmScriptingHostConfiguration
+        ): ReplCompilationState<ReplCodeAnalyzerBase> =
+            createCompilationState(scriptCompilationConfiguration, hostConfiguration) { context1, resolutionFilter ->
+                ReplCodeAnalyzerBase(context1.environment, implicitsResolutionFilter = resolutionFilter)
+            }
+
+        fun <AnalyzerT : ReplCodeAnalyzerBase> createCompilationState(
+            scriptCompilationConfiguration: ScriptCompilationConfiguration,
+            hostConfiguration: ScriptingHostConfiguration = defaultJvmScriptingHostConfiguration,
+            initAnalyzer: (SharedScriptCompilationContext, ImplicitsExtensionsResolutionFilter) -> AnalyzerT
+        ): ReplCompilationState<AnalyzerT> {
+            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,
+                initAnalyzer,
+                ReplImplicitsExtensionsResolutionFilter(
+                    scriptCompilationConfiguration[ScriptCompilationConfiguration.skipExtensionsResolutionForImplicits].orEmpty(),
+                    scriptCompilationConfiguration[ScriptCompilationConfiguration.skipExtensionsResolutionForImplicitsExceptInnermost].orEmpty()
+                )
+            )
+        }
+    }
 }
 
 class ReplCompilationState<AnalyzerT : ReplCodeAnalyzerBase>(
     val context: SharedScriptCompilationContext,
     val analyzerInit: (context: SharedScriptCompilationContext, implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter) -> AnalyzerT,
-    override val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter
+    val implicitsResolutionFilter: ReplImplicitsExtensionsResolutionFilter
 ) : 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: AnalyzerT by lazy {
+
+    val analyzerEngine: AnalyzerT by lazy {
         val analyzer = analyzerInit(context, implicitsResolutionFilter)
         val psiFacade = KotlinJavaPsiFacade.getInstance(environment.project)
         psiFacade.setNotFoundPackagesCachingStrategy(ReplNotFoundPackagesCachingStrategy)
@@ -344,6 +358,6 @@ class ReplCompilationState<AnalyzerT : ReplCodeAnalyzerBase>(
         mangler to symbolTable
     }
 
-    override val mangler: JvmDescriptorMangler get() = manglerAndSymbolTable.first
-    override val symbolTable: SymbolTable get() = manglerAndSymbolTable.second
+    val mangler: JvmDescriptorMangler get() = manglerAndSymbolTable.first
+    val symbolTable: SymbolTable get() = manglerAndSymbolTable.second
 }

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

@@ -16,7 +16,10 @@ import kotlin.concurrent.read
 import kotlin.concurrent.write
 import kotlin.script.experimental.api.KotlinType
 
-class ReplImplicitsExtensionsResolutionFilter : ImplicitsExtensionsResolutionFilter {
+class ReplImplicitsExtensionsResolutionFilter(
+    classesToSkip: Collection<KotlinType> = emptyList(),
+    classesToSkipAfterFirstTime: Collection<KotlinType> = emptyList()
+) : ImplicitsExtensionsResolutionFilter {
     private val lock = ReentrantReadWriteLock()
     private var classesToSkipNames: Set<String> = emptySet()
     private var classesToSkipFirstTimeNames: Set<String> = emptySet()
@@ -29,6 +32,10 @@ class ReplImplicitsExtensionsResolutionFilter : ImplicitsExtensionsResolutionFil
         classesToSkipFirstTimeNames = classesToSkipAfterFirstTime.mapTo(hashSetOf()) { it.typeName }
     }
 
+    init {
+        update(classesToSkip, classesToSkipAfterFirstTime)
+    }
+
     override fun getScopesWithInfo(
         scopes: Sequence<HierarchicalScope>
     ): Sequence<ScopeWithImplicitsExtensionsResolutionInfo> {

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

@@ -11,14 +11,11 @@ import org.jetbrains.kotlin.cli.common.repl.BasicReplStageHistory
 import org.jetbrains.kotlin.cli.common.repl.IReplStageState
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
-import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmDescriptorMangler
-import org.jetbrains.kotlin.ir.util.SymbolTable
-import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.write
 import kotlin.script.experimental.api.ScriptCompilationConfiguration
 
-class JvmReplCompilerStageHistory<CompilationT : JvmReplCompilerState.Compilation>(private val state: JvmReplCompilerState<CompilationT>) :
+class JvmReplCompilerStageHistory<CompilationT : JvmReplCompilerState.Compilation>(state: JvmReplCompilerState<CompilationT>) :
     BasicReplStageHistory<ScriptDescriptor>(state.lock)
 
 class JvmReplCompilerState<CompilationT : JvmReplCompilerState.Compilation>(
@@ -63,9 +60,5 @@ class JvmReplCompilerState<CompilationT : JvmReplCompilerState.Compilation>(
         val disposable: Disposable?
         val baseScriptCompilationConfiguration: ScriptCompilationConfiguration
         val environment: KotlinCoreEnvironment
-        val analyzerEngine: ReplCodeAnalyzerBase
-        val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter
-        val mangler: JvmDescriptorMangler
-        val symbolTable: SymbolTable
     }
 }