Sdk passed in script dependencies should be used during script resolve

PsiElementFinder doesn't find classes from sdks that aren't Project SDK, so when Script Sdk differs from it we need to additionally search classes in script sdk using KotlinScriptDependenciesClassFinder (as we do for classes from ScriptDependencies.classpath)

^KT-31152 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/ScriptModuleInfos.kt

@@ -73,8 +73,8 @@ sealed class ScriptDependenciesInfo(val project: Project) : IdeaModuleInfo, Bina
     ) : ScriptDependenciesInfo(project) {
         override val sdk: Sdk?
             get() {
-                val manager = ScriptDependenciesManager.getInstance(project)
+                return ScriptDependenciesManager.getInstance(project).getScriptSdk(scriptFile)
-                return manager.getScriptSdk(scriptFile) ?: ScriptDependenciesManager.getScriptDefaultSdk(project)
+                    ?: ScriptDependenciesManager.getScriptDefaultSdk(project)
             }
 
         override fun contentScope(): GlobalSearchScope {
@@ -85,14 +85,17 @@ sealed class ScriptDependenciesInfo(val project: Project) : IdeaModuleInfo, Bina
         }
     }
 
+    // we do not know which scripts these dependencies are
     class ForProject(project: Project) : ScriptDependenciesInfo(project) {
         override val sdk: Sdk?
-            get() = ScriptDependenciesManager.getScriptDefaultSdk(project)
+            get() {
+                return ScriptDependenciesManager.getInstance(project).getAllScriptsSdks().firstOrNull()
+                    ?: ScriptDependenciesManager.getScriptDefaultSdk(project)
+            }
 
         override fun contentScope(): GlobalSearchScope {
-            // we do not know which scripts these dependencies are
             return KotlinSourceFilterScope.libraryClassFiles(
-                ScriptDependenciesManager.getInstance(project).getAllScriptsClasspathScope(), project
+                ScriptDependenciesManager.getInstance(project).getAllScriptsDependenciesClassFilesScope(), project
             )
         }
     }
@@ -105,7 +108,7 @@ sealed class ScriptDependenciesSourceInfo(val project: Project) : IdeaModuleInfo
         get() = ScriptDependenciesInfo.ForProject(project)
 
     override fun sourceScope(): GlobalSearchScope = KotlinSourceFilterScope.librarySources(
-        ScriptDependenciesManager.getInstance(project).getAllLibrarySourcesScope(), project
+        ScriptDependenciesManager.getInstance(project).getAllScriptDependenciesSourcesScope(), project
     )
 
     override fun hashCode() = project.hashCode()

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt

@@ -225,14 +225,14 @@ private inline fun <T> collectInfosByVirtualFile(
 
     val isBinary = virtualFile.fileType.isKotlinBinary()
     val scriptConfigurationManager = ScriptDependenciesManager.getInstance(project)
-    if (isBinary && virtualFile in scriptConfigurationManager.getAllScriptsClasspathScope()) {
+    if (isBinary && virtualFile in scriptConfigurationManager.getAllScriptsDependenciesClassFilesScope()) {
         if (treatAsLibrarySource) {
             onOccurrence(ScriptDependenciesSourceInfo.ForProject(project))
         } else {
             onOccurrence(ScriptDependenciesInfo.ForProject(project))
         }
     }
-    if (!isBinary && virtualFile in scriptConfigurationManager.getAllLibrarySourcesScope()) {
+    if (!isBinary && virtualFile in scriptConfigurationManager.getAllScriptDependenciesSourcesScope()) {
         onOccurrence(ScriptDependenciesSourceInfo.ForProject(project))
     }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/ProjectRootsUtil.kt

@@ -129,11 +129,11 @@ object ProjectRootsUtil {
 
         if (includeLibraryClasses && (isBinary || canContainClassFiles)) {
             if (fileIndex.isInLibraryClasses(file)) return true
-            if (scriptConfigurationManager?.getAllScriptsClasspathScope()?.contains(file) == true) return true
+            if (scriptConfigurationManager?.getAllScriptsDependenciesClassFilesScope()?.contains(file) == true) return true
         }
         if (includeLibrarySource && !isBinary) {
             if (fileIndex.isInLibrarySource(file)) return true
-            if (scriptConfigurationManager?.getAllLibrarySourcesScope()?.contains(file) == true &&
+            if (scriptConfigurationManager?.getAllScriptDependenciesSourcesScope()?.contains(file) == true &&
                 !fileIndex.isInSourceContentWithoutInjected(file)
             ) {
                 return true

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/KotlinScriptDependenciesClassFinder.kt

@@ -31,7 +31,7 @@ class KotlinScriptDependenciesClassFinder(
     private val scriptDependenciesManager: ScriptDependenciesManager
 ) : NonClasspathClassFinder(project), KotlinSafeClassFinder {
 
-    override fun calcClassRoots(): List<VirtualFile> = scriptDependenciesManager.getAllScriptsClasspath().toList()
+    override fun calcClassRoots(): List<VirtualFile> = scriptDependenciesManager.getAllScriptsDependenciesClassFiles().toList()
 
     override fun findClass(qualifiedName: String, scope: GlobalSearchScope): PsiClass? {
         tailrec fun findClassInner(parentQualifier: String, inners: List<String> = emptyList()): PsiClass? {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesCache.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.core.script
 import com.intellij.codeInsight.daemon.DaemonCodeAnalyzer
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.roots.OrderRootType
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElementFinder
 import com.intellij.psi.PsiManager
@@ -55,36 +56,66 @@ class ScriptDependenciesCache(private val project: Project) {
 
     private val scriptsDependenciesClasspathScopeCache = SLRUCacheWithLock<GlobalSearchScope>()
 
-    fun scriptDependenciesClasspathScope(file: VirtualFile): GlobalSearchScope {
+    fun scriptDependenciesClassFilesScope(file: VirtualFile): GlobalSearchScope {
         return scriptsDependenciesClasspathScopeCache.getOrPut(file) {
-            val dependencies = scriptDependenciesCache.get(file)
+            val dependencies = scriptDependenciesCache.get(file) ?: return@getOrPut GlobalSearchScope.EMPTY_SCOPE
-            val roots = dependencies?.classpath ?: emptyList()
+            val roots = dependencies.classpath
-            NonClasspathDirectoriesScope.compose(ScriptDependenciesManager.toVfsRoots(roots))
+
+            val sdk = ScriptDependenciesManager.getScriptSdk(dependencies)
+
+            @Suppress("FoldInitializerAndIfToElvis")
+            if (sdk == null) {
+                return@getOrPut NonClasspathDirectoriesScope.compose(ScriptDependenciesManager.toVfsRoots(roots))
+            }
+
+            return@getOrPut NonClasspathDirectoriesScope.compose(
+                sdk.rootProvider.getFiles(OrderRootType.CLASSES).toList() +
+                        ScriptDependenciesManager.toVfsRoots(roots)
+            )
         }
     }
 
-    val allScriptsClasspath by ClearableLazyValue(cacheLock) {
+    val allSdks by ClearableLazyValue(cacheLock) {
-        val files = scriptDependenciesCache.getAll().flatMap { it.value.classpath }.distinct()
+        scriptDependenciesCache.getAll()
-        ScriptDependenciesManager.toVfsRoots(files)
+            .mapNotNull { ScriptDependenciesManager.getInstance(project).getScriptSdk(it.key) }
+            .distinct()
     }
 
-    val allScriptsClasspathScope by ClearableLazyValue(cacheLock) {
+    val allDependenciesClassFiles by ClearableLazyValue(cacheLock) {
-        NonClasspathDirectoriesScope(allScriptsClasspath)
+        val sdkFiles = allSdks
+            .filter { it != ScriptDependenciesManager.getProjectSdk(project) }
+            .flatMap { it.rootProvider.getFiles(OrderRootType.CLASSES).toList() }
+
+        val scriptDependenciesClasspath = scriptDependenciesCache.getAll().flatMap { it.value.classpath }.distinct()
+
+        sdkFiles + ScriptDependenciesManager.toVfsRoots(scriptDependenciesClasspath)
     }
 
-    val allLibrarySources by ClearableLazyValue(cacheLock) {
+    val allDependenciesSources by ClearableLazyValue(cacheLock) {
-        ScriptDependenciesManager.toVfsRoots(scriptDependenciesCache.getAll().flatMap { it.value.sources }.distinct())
+        val sdkSources = allSdks
+            .filter { it != ScriptDependenciesManager.getProjectSdk(project) }
+            .flatMap { it.rootProvider.getFiles(OrderRootType.SOURCES).toList() }
+
+        val scriptDependenciesSources = scriptDependenciesCache.getAll().flatMap { it.value.sources }.distinct()
+        sdkSources + ScriptDependenciesManager.toVfsRoots(scriptDependenciesSources)
     }
 
-    val allLibrarySourcesScope by ClearableLazyValue(cacheLock) {
+    val allDependenciesClassFilesScope by ClearableLazyValue(cacheLock) {
-        NonClasspathDirectoriesScope(allLibrarySources)
+        NonClasspathDirectoriesScope.compose(allDependenciesClassFiles)
+    }
+
+    val allDependenciesSourcesScope by ClearableLazyValue(cacheLock) {
+        NonClasspathDirectoriesScope.compose(allDependenciesSources)
     }
 
     private fun onChange(files: List<VirtualFile>) {
-        this::allScriptsClasspath.clearValue()
+        this::allSdks.clearValue()
-        this::allScriptsClasspathScope.clearValue()
+
-        this::allLibrarySources.clearValue()
+        this::allDependenciesClassFiles.clearValue()
-        this::allLibrarySourcesScope.clearValue()
+        this::allDependenciesClassFilesScope.clearValue()
+
+        this::allDependenciesSources.clearValue()
+        this::allDependenciesSourcesScope.clearValue()
 
         scriptsDependenciesClasspathScopeCache.clear()
 
@@ -110,8 +141,10 @@ class ScriptDependenciesCache(private val project: Project) {
     }
 
     fun hasNotCachedRoots(scriptDependencies: ScriptDependencies): Boolean {
-        return !allScriptsClasspath.containsAll(ScriptDependenciesManager.toVfsRoots(scriptDependencies.classpath)) ||
+        return !allSdks.contains(ScriptDependenciesManager.getScriptSdk(scriptDependencies)) ||
-                !allLibrarySources.containsAll(ScriptDependenciesManager.toVfsRoots(scriptDependencies.sources))
+                !allDependenciesClassFiles.containsAll(ScriptDependenciesManager.toVfsRoots(scriptDependencies.classpath)) ||
+                !allDependenciesSources.containsAll(ScriptDependenciesManager.toVfsRoots(scriptDependencies.sources))
+
     }
 
     fun clear() {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesManager.kt

@@ -52,12 +52,15 @@ class ScriptDependenciesManager internal constructor(
     fun getScriptDependencies(file: VirtualFile): ScriptDependencies = cacheUpdater.getCurrentDependencies(file)
     fun getScriptSdk(file: VirtualFile): Sdk? = getScriptSdk(getScriptDependencies(file))
 
-    fun getScriptDependenciesClassFilesScope(file: VirtualFile) = cache.scriptDependenciesClasspathScope(file)
+    fun getScriptDependenciesClassFilesScope(file: VirtualFile) = cache.scriptDependenciesClassFilesScope(file)
 
-    fun getAllScriptsClasspathScope() = cache.allScriptsClasspathScope
+    fun getAllScriptsSdks() = cache.allSdks
-    fun getAllLibrarySourcesScope() = cache.allLibrarySourcesScope
+
-    fun getAllLibrarySources() = cache.allLibrarySources
+    fun getAllScriptsDependenciesClassFilesScope() = cache.allDependenciesClassFilesScope
-    fun getAllScriptsClasspath() = cache.allScriptsClasspath
+    fun getAllScriptDependenciesSourcesScope() = cache.allDependenciesSourcesScope
+
+    fun getAllScriptsDependenciesClassFiles() = cache.allDependenciesClassFiles
+    fun getAllScriptDependenciesSources() = cache.allDependenciesSources
 
     companion object {
         @JvmStatic

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/dependencies/JavaClassesInScriptDependenciesShortNameCache.kt

@@ -35,7 +35,7 @@ class JavaClassesInScriptDependenciesShortNameCache(private val project: Project
     override fun getAllClassNames(dest: HashSet<String>) {}
 
     override fun getClassesByName(name: String, scope: GlobalSearchScope): Array<out PsiClass> {
-        val classpathScope = ScriptDependenciesManager.getInstance(project).getAllScriptsClasspathScope()
+        val classpathScope = ScriptDependenciesManager.getInstance(project).getAllScriptsDependenciesClassFilesScope()
         val classes = StubIndex.getElements(
                 JavaShortClassNameIndex.getInstance().key, name, project, classpathScope.intersectWith(scope), PsiClass::class.java
         )

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/dependencies/KotlinScriptDependenciesIndexableSetContributor.kt

@@ -25,7 +25,7 @@ class KotlinScriptDependenciesIndexableSetContributor : IndexableSetContributor(
 
     override fun getAdditionalProjectRootsToIndex(project: Project): Set<VirtualFile> {
         val manager = ScriptDependenciesManager.getInstance(project)
-        return (manager.getAllScriptsClasspath() + manager.getAllLibrarySources()).filterTo(LinkedHashSet()) { it.isValid }
+        return (manager.getAllScriptsDependenciesClassFiles() + manager.getAllScriptDependenciesSources()).filterTo(LinkedHashSet()) { it.isValid }
     }
 
     override fun getAdditionalRootsToIndex(): Set<VirtualFile> = emptySet()

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/dependencies/KotlinScriptResolveScopeProvider.kt

@@ -40,7 +40,7 @@ class KotlinScriptResolveScopeProvider : ResolveScopeProvider() {
             scriptDefinition.template == Any::class -> null
             scriptDefinition is StandardIdeScriptDefinition -> null
             scriptDefinition is KotlinScriptDefinitionFromAnnotatedTemplate -> // TODO: should include the file itself
-                ScriptDependenciesManager.getInstance(project).getAllScriptsClasspathScope()
+                ScriptDependenciesManager.getInstance(project).getAllScriptsDependenciesClassFilesScope()
             else -> null
         }
     }

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/dependencies/ScriptDependencySourceNavigationPolicyForJavaClasses.kt

@@ -50,13 +50,13 @@ class ScriptDependencySourceNavigationPolicyForJavaClasses : ClsCustomNavigation
         val project = file.project
 
         val kotlinScriptConfigurationManager = ScriptDependenciesManager.getInstance(project)
-        if (virtualFile !in kotlinScriptConfigurationManager.getAllScriptsClasspathScope()) return null
+        if (virtualFile !in kotlinScriptConfigurationManager.getAllScriptsDependenciesClassFilesScope()) return null
 
         val sourceFileName = (file.classes.first() as ClsClassImpl).sourceFileName
         val packageName = file.packageName
         val relativePath = if (packageName.isEmpty()) sourceFileName else packageName.replace('.', '/') + '/' + sourceFileName
 
-        for (root in kotlinScriptConfigurationManager.getAllLibrarySources().filter { it.isValid }) {
+        for (root in kotlinScriptConfigurationManager.getAllScriptDependenciesSources().filter { it.isValid }) {
             val sourceFile = root.findFileByRelativePath(relativePath)
             if (sourceFile != null && sourceFile.isValid) {
                 val sourcePsi = file.manager.findFile(sourceFile)

idea/testData/script/definition/defaultTemplate/template.kt

@@ -10,6 +10,7 @@ class FromTextDependenciesResolver : DependenciesResolver {
     @Suppress("UNCHECKED_CAST")
     override fun resolve(scriptContents: ScriptContents, environment: Environment): DependenciesResolver.ResolveResult {
         return ScriptDependencies(
+            javaHome = environment["javaHome"] as? File,
             classpath = (environment["classpath"] as? List<File>).orEmpty(),
             imports = (environment["imports"] as? List<String>).orEmpty(),
             sources = (environment["sources"] as? List<File>).orEmpty()

idea/testData/script/definition/highlighting/customJavaHome/script.kts

@@ -0,0 +1,3 @@
+import java.util.stream.IntStream
+
+// DEPENDENCIES: javaHome:9; classpath:runtime-classes

idea/tests/org/jetbrains/kotlin/idea/script/AbstractScriptConfigurationTest.kt

@@ -10,7 +10,9 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.module.JavaModuleType
 import com.intellij.openapi.module.Module
+import com.intellij.openapi.projectRoots.ProjectJdkTable
 import com.intellij.openapi.roots.ModuleRootModificationUtil
+import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.openapi.vfs.LocalFileSystem
@@ -31,10 +33,12 @@ import org.jetbrains.kotlin.idea.core.script.isScriptDependenciesUpdaterDisabled
 import org.jetbrains.kotlin.idea.core.script.settings.KotlinScriptingSettings
 import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingUtil
 import org.jetbrains.kotlin.idea.navigation.GotoCheck
+import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.MockLibraryUtil
+import org.jetbrains.kotlin.test.TestJdkKind
 import org.jetbrains.kotlin.test.util.addDependency
 import org.jetbrains.kotlin.test.util.projectLibrary
 import org.jetbrains.kotlin.test.util.renderAsGotoImplementation
@@ -136,7 +140,7 @@ abstract class AbstractScriptDefinitionsOrderTest : AbstractScriptConfigurationT
     }
 }
 
-private val validKeys = setOf("sources", "classpath", "imports", "template-classes-names")
+private val validKeys = setOf("javaHome", "sources", "classpath", "imports", "template-classes-names")
 private const val useDefaultTemplate = "// DEPENDENCIES:"
 private const val templatesSettings = "// TEMPLATES: "
 // some bugs can only be reproduced when some module and script have intersecting library dependencies
@@ -165,6 +169,15 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
             ?: error("Couldn't find $SCRIPT_NAME file in $testDir")
     }
 
+    private val sdk by lazy {
+        val jdk = PluginTestCaseBase.jdk(TestJdkKind.MOCK_JDK)
+        runWriteAction {
+            ProjectJdkTable.getInstance().addJdk(jdk, testRootDisposable)
+            ProjectRootManager.getInstance(project).projectSdk = jdk
+        }
+        jdk
+    }
+
     protected fun configureScriptFile(path: String) {
         val mainScriptFile = findMainScript(path)
         val environment = createScriptEnvironment(mainScriptFile)
@@ -189,6 +202,10 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
         }
 
         if (module != null) {
+            ModuleRootModificationUtil.updateModel(module) { model ->
+                model.sdk = sdk
+            }
+
             module.addDependency(
                 projectLibrary(
                     "script-runtime",
@@ -267,6 +284,18 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
             env["template-classes-names"] = listOf("custom.scriptDefinition.Template")
         }
 
+        if (env["javaHome"] != null) {
+            val jdkKind = when ((env["javaHome"] as? List<String>)?.singleOrNull()) {
+                "9" -> TestJdkKind.FULL_JDK_9
+                else -> TestJdkKind.MOCK_JDK
+            }
+            runWriteAction {
+                val jdk = PluginTestCaseBase.jdk(jdkKind)
+                ProjectJdkTable.getInstance().addJdk(jdk, testRootDisposable)
+                env["javaHome"] = File(jdk.homePath)
+            }
+        }
+
         env.putAll(defaultEnvironment)
         return env
     }

idea/tests/org/jetbrains/kotlin/idea/script/ScriptConfigurationHighlightingTestGenerated.java

@@ -61,6 +61,11 @@ public class ScriptConfigurationHighlightingTestGenerated extends AbstractScript
             runTest("idea/testData/script/definition/highlighting/customExtension/");
         }
 
+        @TestMetadata("customJavaHome")
+        public void testCustomJavaHome() throws Exception {
+            runTest("idea/testData/script/definition/highlighting/customJavaHome/");
+        }
+
         @TestMetadata("customLibrary")
         public void testCustomLibrary() throws Exception {
             runTest("idea/testData/script/definition/highlighting/customLibrary/");