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.*
-import org.jetbrains.kotlin.resolve.lazy.FileScopeProvider
-import org.jetbrains.kotlin.resolve.lazy.ResolveSession
 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 {
+    inner abstract class LineInfo {
-        val linePsi: KtFile
+        abstract val linePsi: KtFile
-        val parentLine: SuccessfulLine?
+        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
+    inner 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
+    inner 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 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)

compiler/frontend.java/src/org/jetbrains/kotlin/frontend/java/di/injection.kt

@@ -53,6 +53,8 @@ fun StorageComponentContainer.configureJavaTopDownAnalysis(moduleContentScope: G
 
     useInstance(JvmVirtualFileFinderFactory.SERVICE.getInstance(project).create(moduleContentScope))
 
+    useImpl<FileScopeProviderImpl>()
+
     useImpl<JavaClassFinderImpl>()
     useImpl<SignaturePropagatorImpl>()
     useImpl<LazyResolveBasedCache>()
@@ -83,7 +85,6 @@ fun createContainerForLazyResolveWithJava(
 
     targetEnvironment.configure(this)
 
-    useImpl<FileScopeProviderImpl>()
     useImpl<LazyResolveToken>()
 }.apply {
     javaAnalysisInit()
@@ -108,7 +109,6 @@ fun createContainerForTopDownAnalyzerForJvm(
     CompilerEnvironment.configure(this)
 
     useImpl<SingleModuleClassResolver>()
-    useImpl<FileScopeProviderImpl>()
 }.let {
     it.javaAnalysisInit()
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/FileScopeFactory.kt

@@ -39,6 +39,8 @@ import org.jetbrains.kotlin.types.TypeSubstitutor
 import org.jetbrains.kotlin.utils.Printer
 import org.jetbrains.kotlin.utils.addToStdlib.check
 
+data class FileScopes(val lexicalScope: LexicalScope, val importingScope: ImportingScope, val importResolver: ImportResolver)
+
 class FileScopeFactory(
         private val topLevelDescriptorProvider: TopLevelDescriptorProvider,
         private val storageManager: StorageManager,
@@ -47,18 +49,11 @@ class FileScopeFactory(
         private val bindingTrace: BindingTrace,
         private val ktImportsFactory: KtImportsFactory
 ) {
-
-    class FileData(val scope: LexicalScope, val importResolver: ImportResolver)
-
-    fun getLexicalScopeAndImportResolver(file: KtFile): FileData = cache(file)
-
     private val defaultImports by storageManager.createLazyValue {
         ktImportsFactory.createImportDirectives(moduleDescriptor.defaultImports)
     }
 
-    private val cache = storageManager.createMemoizedFunction { file: KtFile -> createScopeChainAndImportResolver(file) }
+    fun createScopesForFile(file: KtFile, existingImports: ImportingScope? = null): FileScopes {
-
-    private fun createScopeChainAndImportResolver(file: KtFile): FileData {
         val debugName = "LazyFileScope for file " + file.name
         val tempTrace = TemporaryBindingTrace.create(bindingTrace, "Transient trace for default imports lazy resolve")
 
@@ -89,7 +84,7 @@ class FileScopeFactory(
 
         var scope: ImportingScope
 
-        scope = LazyImportScope(null, defaultAllUnderImportResolver, LazyImportScope.FilteringKind.INVISIBLE_CLASSES,
+        scope = LazyImportScope(existingImports, defaultAllUnderImportResolver, LazyImportScope.FilteringKind.INVISIBLE_CLASSES,
                                 "Default all under imports in $debugName (invisible classes only)")
 
         scope = LazyImportScope(scope, allUnderImportResolver, LazyImportScope.FilteringKind.INVISIBLE_CLASSES,
@@ -110,11 +105,9 @@ class FileScopeFactory(
 
         scope = currentPackageScope(packageView, aliasImportNames, dummyContainerDescriptor, FilteringKind.VISIBLE_CLASSES, scope)
 
-        scope = LazyImportScope(scope, explicitImportResolver, LazyImportScope.FilteringKind.ALL, "Explicit imports in $debugName")
+        val importingScope = LazyImportScope(scope, explicitImportResolver, LazyImportScope.FilteringKind.ALL, "Explicit imports in $debugName")
 
-        val lexicalScope = LexicalScope.empty(scope, packageFragment)
+        val lexicalScope = LexicalScope.empty(importingScope, packageFragment)
-
-        bindingTrace.recordScope(lexicalScope, file)
 
         val importResolver = object : ImportResolver {
             override fun forceResolveAllImports() {
@@ -132,7 +125,7 @@ class FileScopeFactory(
             }
         }
 
-        return FileData(lexicalScope, importResolver)
+        return FileScopes(lexicalScope, importingScope, importResolver)
     }
 
     private enum class FilteringKind {

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/FileScopeProvider.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2015 JetBrains s.r.o.
+ * Copyright 2010-2016 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,17 +16,42 @@
 
 package org.jetbrains.kotlin.resolve.lazy
 
+import com.intellij.openapi.util.Key
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.UserDataProperty
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.bindingContextUtil.recordScope
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
+import org.jetbrains.kotlin.storage.StorageManager
 
 interface FileScopeProvider {
-    fun getFileResolutionScope(file: KtFile): LexicalScope
+    fun getFileResolutionScope(file: KtFile): LexicalScope = getFileScopes(file).lexicalScope
-    fun getImportResolver(file: KtFile): ImportResolver
+    fun getImportResolver(file: KtFile): ImportResolver = getFileScopes(file).importResolver
+
+    fun getFileScopes(file: KtFile): FileScopes
 
     object ThrowException : FileScopeProvider {
-        override fun getFileResolutionScope(file: KtFile) = throw UnsupportedOperationException("Should not be called")
+        override fun getFileScopes(file: KtFile) = throw UnsupportedOperationException("Should not be called")
-        override fun getImportResolver(file: KtFile) = throw UnsupportedOperationException("Should not be called")
     }
 }
 
+class FileScopeProviderImpl(
+        private val fileScopeFactory: FileScopeFactory,
+        private val bindingTrace: BindingTrace,
+        private val storageManager: StorageManager
+) : FileScopeProvider {
 
+    private val cache = storageManager.createMemoizedFunction<KtFile, FileScopes> { file ->
+        val scopes = (file.originalFile as KtFile?)?.fileScopesCustomizer?.let { it.createFileScopes(fileScopeFactory) } ?: fileScopeFactory.createScopesForFile(file)
+        bindingTrace.recordScope(scopes.lexicalScope, file)
+        scopes
+    }
+
+    override fun getFileScopes(file: KtFile) = cache(file)
+}
+
+interface FileScopesCustomizer {
+    fun createFileScopes(fileScopeFactory: FileScopeFactory): FileScopes
+}
+
+var KtFile.fileScopesCustomizer: FileScopesCustomizer? by UserDataProperty(Key.create("FILE_SCOPES_CUSTOMIZER"))

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/FileScopeProviderImpl.kt

@@ -1,30 +0,0 @@
-/*
- * Copyright 2010-2016 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.jetbrains.kotlin.resolve.lazy
-
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.scopes.LexicalScope
-
-class FileScopeProviderImpl(private val fileScopeFactory: FileScopeFactory) : FileScopeProvider {
-    override fun getFileResolutionScope(file: KtFile): LexicalScope {
-        return fileScopeFactory.getLexicalScopeAndImportResolver(file).scope
-    }
-
-    override fun getImportResolver(file: KtFile): ImportResolver {
-        return fileScopeFactory.getLexicalScopeAndImportResolver(file).importResolver
-    }
-}

compiler/testData/repl/multipleImports.repl

@@ -0,0 +1,10 @@
+>>> val c = 3
+>>> import java.util.Date
+>>> Date(0).getTime()
+0
+>>> import java.io.ByteArrayOutputStream
+>>> val s = ByteArrayOutputStream()
+>>> s.javaClass!!.getName()
+java.io.ByteArrayOutputStream
+>>> c + 1
+4

compiler/tests/org/jetbrains/kotlin/repl/ReplInterpreterTestGenerated.java

@@ -89,6 +89,12 @@ public class ReplInterpreterTestGenerated extends AbstractReplInterpreterTest {
         doTest(fileName);
     }
 
+    @TestMetadata("multipleImports.repl")
+    public void testMultipleImports() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/repl/multipleImports.repl");
+        doTest(fileName);
+    }
+
     @TestMetadata("simple.repl")
     public void testSimple() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/repl/simple.repl");