Fix import resolution in Repl

#KT-11428 Fixed

 Refactor ReplState to be a component

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/CliReplAnalyzerEngine.kt

@@ -32,14 +32,15 @@ import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfoFactory
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import org.jetbrains.kotlin.resolve.jvm.TopDownAnalyzerFacadeForJVM
-import org.jetbrains.kotlin.resolve.lazy.FileScopeFactory
-import org.jetbrains.kotlin.resolve.lazy.FileScopeProvider
-import org.jetbrains.kotlin.resolve.lazy.ResolveSession
+import org.jetbrains.kotlin.resolve.lazy.*
 import org.jetbrains.kotlin.resolve.lazy.data.KtClassLikeInfo
 import org.jetbrains.kotlin.resolve.lazy.declarations.*
 import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyScriptDescriptor
-import org.jetbrains.kotlin.resolve.scopes.LexicalScope
+import org.jetbrains.kotlin.resolve.scopes.ImportingScope
+import org.jetbrains.kotlin.resolve.scopes.utils.parentsWithSelf
+import org.jetbrains.kotlin.resolve.scopes.utils.replaceImportingScopes
 import org.jetbrains.kotlin.script.ScriptPriorities
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 
 class CliReplAnalyzerEngine(private val environment: KotlinCoreEnvironment) {
     private val topDownAnalysisContext: TopDownAnalysisContext
@@ -47,8 +48,8 @@ class CliReplAnalyzerEngine(private val environment: KotlinCoreEnvironment) {
     private val resolveSession: ResolveSession
     private val scriptDeclarationFactory: ScriptMutableDeclarationProviderFactory
     val module: ModuleDescriptorImpl
-    val trace: BindingTraceContext = CliLightClassGenerationSupport.NoScopeRecordCliBindingTrace()
     private val replState = ReplState()
+    val trace: BindingTraceContext = CliLightClassGenerationSupport.NoScopeRecordCliBindingTrace()
 
     init {
         val moduleContext = TopDownAnalyzerFacadeForJVM.createContextWithSealedModule(environment.project, environment.getModuleName())
@@ -61,7 +62,6 @@ class CliReplAnalyzerEngine(private val environment: KotlinCoreEnvironment) {
                 trace,
                 scriptDeclarationFactory,
                 ProjectScope.getAllScope(environment.project),
-                replState,
                 JvmPackagePartProvider(environment)
         )
 
@@ -101,8 +101,8 @@ class CliReplAnalyzerEngine(private val environment: KotlinCoreEnvironment) {
     }
 
     private fun doAnalyze(linePsi: KtFile): ReplLineAnalysisResult {
-        replState.submitLine(linePsi)
         scriptDeclarationFactory.setDelegateFactory(FileBasedDeclarationProviderFactory(resolveSession.storageManager, listOf(linePsi)))
+        replState.submitLine(linePsi)
 
         val context = topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi))
 
@@ -172,8 +172,17 @@ class ReplState {
     private val lines = hashMapOf<KtFile, LineInfo>()
     private val successfulLines = arrayListOf<SuccessfulLine>()
 
+    val successfulLinesCount: Int
+        get() = successfulLines.size
+
     fun submitLine(ktFile: KtFile) {
-        lines[ktFile] = SubmittedLine(ktFile, successfulLines.lastOrNull())
+        val line = SubmittedLine(ktFile, successfulLines.lastOrNull())
+        lines[ktFile] = line
+        ktFile.fileScopesCustomizer = object : FileScopesCustomizer {
+            override fun createFileScopes(fileScopeFactory: FileScopeFactory): FileScopes {
+                return lineInfo(ktFile)?.computeFileScopes(fileScopeFactory) ?: fileScopeFactory.createScopesForFile(ktFile)
+            }
+        }
     }
 
     fun lineSuccess(ktFile: KtFile, scriptDescriptor: LazyScriptDescriptor) {
@@ -188,24 +197,22 @@ class ReplState {
 
     fun lineInfo(ktFile: KtFile) = lines[ktFile]
 
-    interface LineInfo {
-        val linePsi: KtFile
-        val parentLine: SuccessfulLine?
+    inner abstract class LineInfo {
+        abstract val linePsi: KtFile
+        abstract val parentLine: SuccessfulLine?
 
-        val lexicalScopeBeforeThisLine: LexicalScope? get() = parentLine?.lineDescriptor?.scopeForInitializerResolution
+        fun computeFileScopes(fileScopeFactory: FileScopeFactory): FileScopes {
+            // create scope that wraps previous line lexical scope and adds imports from this line
+            val lexicalScopeAfterLastLine = parentLine?.lineDescriptor?.scopeForInitializerResolution
+                                            ?: return fileScopeFactory.createScopesForFile(linePsi)
+            val lastLineImports = lexicalScopeAfterLastLine.parentsWithSelf.firstIsInstance<ImportingScope>()
+            val scopesForThisLine = fileScopeFactory.createScopesForFile(linePsi, lastLineImports)
+            val combinedLexicalScopes = lexicalScopeAfterLastLine.replaceImportingScopes(scopesForThisLine.importingScope)
+            return FileScopes(combinedLexicalScopes, scopesForThisLine.importingScope, scopesForThisLine.importResolver)
+        }
     }
 
-    data class SubmittedLine(override val linePsi: KtFile, override val parentLine: SuccessfulLine?): LineInfo
-    data class SuccessfulLine(override val linePsi: KtFile, override val parentLine: SuccessfulLine?, val lineDescriptor: LazyScriptDescriptor): LineInfo
-    data class FailedLine(override val linePsi: KtFile, override val parentLine: SuccessfulLine?): LineInfo
+    inner class SubmittedLine(override val linePsi: KtFile, override val parentLine: SuccessfulLine?) : LineInfo()
+    inner class SuccessfulLine(override val linePsi: KtFile, override val parentLine: SuccessfulLine?, val lineDescriptor: LazyScriptDescriptor) : LineInfo()
+    inner class FailedLine(override val linePsi: KtFile, override val parentLine: SuccessfulLine?) : LineInfo()
 }
-
-class ReplFileScopeProvider(
-        private val replState: ReplState,
-        private val fileScopeFactory: FileScopeFactory
-) : FileScopeProvider {
-    override fun getFileResolutionScope(file: KtFile)
-            = replState.lineInfo(file)?.lexicalScopeBeforeThisLine ?: fileScopeFactory.getLexicalScopeAndImportResolver(file).scope
-
-    override fun getImportResolver(file: KtFile) = fileScopeFactory.getLexicalScopeAndImportResolver(file).importResolver
-}

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/di/injection.kt

@@ -17,8 +17,6 @@
 package org.jetbrains.kotlin.cli.jvm.repl.di
 
 import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.cli.jvm.repl.ReplFileScopeProvider
-import org.jetbrains.kotlin.cli.jvm.repl.ReplState
 import org.jetbrains.kotlin.container.*
 import org.jetbrains.kotlin.context.ModuleContext
 import org.jetbrains.kotlin.descriptors.PackagePartProvider
@@ -41,15 +39,12 @@ fun createContainerForReplWithJava(
         bindingTrace: BindingTrace,
         declarationProviderFactory: DeclarationProviderFactory,
         moduleContentScope: GlobalSearchScope,
-        replState: ReplState,
         packagePartProvider: PackagePartProvider
 ): ContainerForReplWithJava = createContainer("ReplWithJava") {
     useInstance(packagePartProvider)
     configureModule(moduleContext, JvmPlatform, bindingTrace)
     configureJavaTopDownAnalysis(moduleContentScope, moduleContext.project, LookupTracker.DO_NOTHING)
 
-    useInstance(replState)
-    useImpl<ReplFileScopeProvider>()
     useInstance(declarationProviderFactory)
 
     CompilerEnvironment.configure(this)