Fix Optimize import for scripts not under source roots

KT-17231 Fixed
KT-21981 Fixed

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

@@ -17,8 +17,6 @@
 package org.jetbrains.kotlin.resolve.lazy
 
 import gnu.trove.THashSet
-import org.jetbrains.kotlin.builtins.PlatformToKotlinClassMap
-import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.incremental.components.LookupLocation
@@ -27,11 +25,13 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtImportDirective
 import org.jetbrains.kotlin.psi.KtImportInfo
-import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.ImportPath
+import org.jetbrains.kotlin.resolve.TargetPlatform
+import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
 import org.jetbrains.kotlin.resolve.scopes.*
 import org.jetbrains.kotlin.resolve.source.KotlinSourceElement
-import org.jetbrains.kotlin.script.getScriptExternalDependencies
+import org.jetbrains.kotlin.script.ScriptDependenciesProvider
-import org.jetbrains.kotlin.storage.StorageManager
 import org.jetbrains.kotlin.storage.getValue
 import org.jetbrains.kotlin.utils.Printer
 
@@ -178,9 +178,9 @@ class FileScopeFactory(
         val result = FileScopes(lexicalScope, lazyImportingScope, importResolver)
 
         private fun createDefaultImportResolversForFile(): DefaultImportResolvers {
-            val extraImports = file.takeIf { it.isScript() }?.originalFile?.virtualFile?.let { vFile ->
+            val extraImports = file.takeIf { it.isScript() }?.let { ktFile ->
-                val scriptExternalDependencies = getScriptExternalDependencies(vFile, file.project)
+                val scriptDependencies = ScriptDependenciesProvider.getInstance(ktFile.project).getScriptDependencies(ktFile.originalFile)
-                scriptExternalDependencies?.imports?.map { DefaultImportImpl(ImportPath.fromString(it)) }.orEmpty()
+                scriptDependencies?.imports?.map { DefaultImportImpl(ImportPath.fromString(it)) }
             }.orEmpty()
 
             if (extraImports.isEmpty() && aliasImportNames.isEmpty()) {

compiler/psi/src/org/jetbrains/kotlin/script/ScriptDependenciesProvider.kt

@@ -31,6 +31,3 @@ interface ScriptDependenciesProvider {
             ServiceManager.getService(project, ScriptDependenciesProvider::class.java)
     }
 }
-
-fun getScriptExternalDependencies(file: VirtualFile, project: Project): ScriptDependencies? =
-    ScriptDependenciesProvider.getInstance(project).getScriptDependencies(file)

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt

@@ -601,7 +601,7 @@ fun main(args: Array<String>) {
         }
 
         testClass<AbstractJvmOptimizeImportsTest> {
-            model("editor/optimizeImports/jvm", pattern = KT_WITHOUT_DOTS_IN_NAME)
+            model("editor/optimizeImports/jvm", pattern = KT_OR_KTS_WITHOUT_DOTS_IN_NAME)
             model("editor/optimizeImports/common", pattern = KT_WITHOUT_DOTS_IN_NAME)
         }
         testClass<AbstractJsOptimizeImportsTest> {

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt.as31

@@ -599,7 +599,7 @@ fun main(args: Array<String>) {
         }
 
         testClass<AbstractJvmOptimizeImportsTest> {
-            model("editor/optimizeImports/jvm", pattern = KT_WITHOUT_DOTS_IN_NAME)
+            model("editor/optimizeImports/jvm", pattern = KT_OR_KTS_WITHOUT_DOTS_IN_NAME)
             model("editor/optimizeImports/common", pattern = KT_WITHOUT_DOTS_IN_NAME)
         }
         testClass<AbstractJsOptimizeImportsTest> {

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt.as32

@@ -599,7 +599,7 @@ fun main(args: Array<String>) {
         }
 
         testClass<AbstractJvmOptimizeImportsTest> {
-            model("editor/optimizeImports/jvm", pattern = KT_WITHOUT_DOTS_IN_NAME)
+            model("editor/optimizeImports/jvm", pattern = KT_OR_KTS_WITHOUT_DOTS_IN_NAME)
             model("editor/optimizeImports/common", pattern = KT_WITHOUT_DOTS_IN_NAME)
         }
         testClass<AbstractJsOptimizeImportsTest> {

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt.as33

@@ -591,7 +591,7 @@ fun main(args: Array<String>) {
         }
 
         testClass<AbstractJvmOptimizeImportsTest> {
-            model("editor/optimizeImports/jvm", pattern = KT_WITHOUT_DOTS_IN_NAME)
+            model("editor/optimizeImports/jvm", pattern = KT_OR_KTS_WITHOUT_DOTS_IN_NAME)
             model("editor/optimizeImports/common", pattern = KT_WITHOUT_DOTS_IN_NAME)
         }
         testClass<AbstractJsOptimizeImportsTest> {

idea/ide-common/src/org/jetbrains/kotlin/idea/util/OptimizedImportsBuilder.kt

@@ -266,7 +266,10 @@ class OptimizedImportsBuilder(
                 append("\n")
             }
         }
-        val fileWithImports = KtPsiFactory(originalFile).createAnalyzableFile("Dummy.kt", fileText, originalFile)
+        val fileWithImports = KtPsiFactory(originalFile).createAnalyzableFile("Dummy_" + originalFile.name, fileText, originalFile)
+        if (file.isScript()) {
+            fileWithImports.originalFile = originalFile
+        }
         return fileWithImports.getFileResolutionScope()
     }
 

idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt

@@ -22,6 +22,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.project.ScriptModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getNullableModuleInfo
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
@@ -52,7 +53,8 @@ class KotlinImportOptimizer : ImportOptimizer {
 
     private class OptimizeProcess(private val file: KtFile) {
         fun execute() {
-            if (file.getNullableModuleInfo() !is ModuleSourceInfo) return
+            val moduleInfo = file.getNullableModuleInfo()
+            if (moduleInfo !is ModuleSourceInfo && moduleInfo !is ScriptModuleInfo) return
 
             val oldImports = file.importDirectives
             if (oldImports.isEmpty()) return

idea/src/org/jetbrains/kotlin/idea/util/ImportInsertHelperImpl.kt

@@ -47,6 +47,7 @@ import org.jetbrains.kotlin.resolve.scopes.utils.findClassifier
 import org.jetbrains.kotlin.resolve.scopes.utils.findFunction
 import org.jetbrains.kotlin.resolve.scopes.utils.findPackage
 import org.jetbrains.kotlin.resolve.scopes.utils.findVariable
+import org.jetbrains.kotlin.script.ScriptDependenciesProvider
 import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.*
 
@@ -62,7 +63,13 @@ class ImportInsertHelperImpl(private val project: Project) : ImportInsertHelper(
         val languageVersionSettings = contextFile.getResolutionFacade().frontendService<LanguageVersionSettings>()
         val platform = TargetPlatformDetector.getPlatform(contextFile)
         val allDefaultImports = platform.getDefaultImports(languageVersionSettings, includeLowPriorityImports = true)
-        return importPath.isImported(allDefaultImports, platform.excludedImports)
+
+        val scriptExtraImports = contextFile.takeIf { it.isScript() }?.let { ktFile ->
+            val scriptDependencies = ScriptDependenciesProvider.getInstance(ktFile.project).getScriptDependencies(ktFile.originalFile)
+            scriptDependencies?.imports?.map { ImportPath.fromString(it) }
+        }.orEmpty()
+
+        return importPath.isImported(allDefaultImports + scriptExtraImports, platform.excludedImports)
     }
 
     override fun isImportedWithLowPriorityDefaultImport(importPath: ImportPath, contextFile: KtFile): Boolean {

idea/testData/editor/optimizeImports/jvm/UnusedImportsInScript.kts

@@ -0,0 +1,24 @@
+import com.sun.corba.se.impl.util.JDKBridge
+import com.sun.org.apache.xpath.internal.operations.And
+import java.lang.StringBuilder
+import java.net.HttpRetryException
+import java.util.ArrayList
+import java.util.HashMap
+import kotlin.test.asserter
+import kotlin.test.Asserter
+import kotlin.Charsets
+import kotlin.Assertions
+import kotlin.system.measureTimeMillis
+
+class Action {
+    fun test() {
+        val chs = Charsets.UTF8
+        // val traait : Asserter = asserter
+        val traait : Asserter? = null
+        val objectImport : Assertions? = null
+        measureTimeMillis({ println(HashMap<String, Int>().size()) })
+        val test : ArrayList<Int>? = null
+    }
+}
+
+// RUNTIME

idea/testData/editor/optimizeImports/jvm/UnusedImportsInScript.kts.after

@@ -0,0 +1,16 @@
+import java.util.ArrayList
+import java.util.HashMap
+import kotlin.system.measureTimeMillis
+
+class Action {
+    fun test() {
+        val chs = Charsets.UTF8
+        // val traait : Asserter = asserter
+        val traait : Asserter? = null
+        val objectImport : Assertions? = null
+        measureTimeMillis({ println(HashMap<String, Int>().size()) })
+        val test : ArrayList<Int>? = null
+    }
+}
+
+// RUNTIME

idea/tests/org/jetbrains/kotlin/idea/imports/JvmOptimizeImportsTestGenerated.java

@@ -28,7 +28,7 @@ public class JvmOptimizeImportsTestGenerated extends AbstractJvmOptimizeImportsT
         }
 
         public void testAllFilesPresentInJvm() throws Exception {
-            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/editor/optimizeImports/jvm"), Pattern.compile("^([^.]+)\\.kt$"), TargetBackend.ANY, true);
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/editor/optimizeImports/jvm"), Pattern.compile("^([^.]+)\\.(kt|kts)$"), TargetBackend.ANY, true);
         }
 
         @TestMetadata("AlreadyOptimized.kt")
@@ -161,6 +161,11 @@ public class JvmOptimizeImportsTestGenerated extends AbstractJvmOptimizeImportsT
             runTest("idea/testData/editor/optimizeImports/jvm/UnusedImports.kt");
         }
 
+        @TestMetadata("UnusedImportsInScript.kts")
+        public void testUnusedImportsInScript() throws Exception {
+            runTest("idea/testData/editor/optimizeImports/jvm/UnusedImportsInScript.kts");
+        }
+
         @TestMetadata("WithAliases.kt")
         public void testWithAliases() throws Exception {
             runTest("idea/testData/editor/optimizeImports/jvm/WithAliases.kt");
@@ -175,7 +180,7 @@ public class JvmOptimizeImportsTestGenerated extends AbstractJvmOptimizeImportsT
             }
 
             public void testAllFilesPresentInAllUnderImports() throws Exception {
-                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/editor/optimizeImports/jvm/allUnderImports"), Pattern.compile("^([^.]+)\\.kt$"), TargetBackend.ANY, true);
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/editor/optimizeImports/jvm/allUnderImports"), Pattern.compile("^([^.]+)\\.(kt|kts)$"), TargetBackend.ANY, true);
             }
 
             @TestMetadata("ClassNameConflict.kt")