Scripting unified cache: update sdks synchronously on changes

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

@@ -137,7 +137,7 @@ interface ScriptConfigurationManager {
         @TestOnly
         fun clearCaches(project: Project) {
             (getInstance(project) as CompositeScriptConfigurationManager).default
-                .updateScriptDefinitions()
+                .updateScriptDefinitionsReferences()
         }
 
         fun clearManualConfigurationLoadingIfNeeded(file: VirtualFile) {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/CompositeScriptConfigurationManager.kt

@@ -7,9 +7,10 @@ package org.jetbrains.kotlin.idea.core.script.configuration
 
 import com.intellij.ProjectTopics
 import com.intellij.codeInsight.daemon.OutsidersPsiFileSupport
-import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.projectRoots.ProjectJdkTable
+import com.intellij.openapi.projectRoots.ProjectJdkTable.JDK_TABLE_TOPIC
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.roots.ModuleRootEvent
 import com.intellij.openapi.roots.ModuleRootListener
@@ -18,9 +19,9 @@ import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager
 import org.jetbrains.kotlin.idea.core.script.ScriptDependenciesModificationTracker
 import org.jetbrains.kotlin.idea.core.script.configuration.listener.ScriptChangesNotifier
-import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsCache
-import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsUpdater
 import org.jetbrains.kotlin.idea.core.script.configuration.utils.getKtFile
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsCache
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsUpdater
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
@@ -44,12 +45,18 @@ class CompositeScriptConfigurationManager(val project: Project) : ScriptConfigur
     @Suppress("unused")
     private val notifier = ScriptChangesNotifier(project)
 
-    val updater = ScriptClassRootsUpdater(project, this)
+    private val classpathRoots: ScriptClassRootsCache
+        get() = updater.classpathRoots
 
     private val plugins = ScriptingSupport.EPN.getPoint(project).extensionList
 
     val default = DefaultScriptingSupport(this)
 
+    val updater = ScriptClassRootsUpdater(project, this) { builder ->
+        default.collectConfigurations(builder)
+        plugins.forEach { it.collectConfigurations(builder) }
+    }
+
     fun tryGetScriptDefinitionFast(locationId: String): ScriptDefinition? {
         return classpathRoots.getLightScriptInfo(locationId)?.definition
     }
@@ -80,30 +87,13 @@ class CompositeScriptConfigurationManager(val project: Project) : ScriptConfigur
         plugins.firstOrNull { it.isApplicable(file.originalFile.virtualFile) }?.isConfigurationLoadingInProgress(file)
             ?: default.isConfigurationLoadingInProgress(file)
 
-    @Volatile
-    private var classpathRoots: ScriptClassRootsCache = recreateRootsCache()
-
     fun getLightScriptInfo(file: String): ScriptClassRootsCache.LightScriptInfo? =
         classpathRoots.getLightScriptInfo(file)
 
-    private fun recreateRootsCache(): ScriptClassRootsCache {
-        val builder = ScriptClassRootsCache.Builder(project)
-        default.collectConfigurations(builder)
-        plugins.forEach { it.collectConfigurations(builder) }
-        return builder.build()
-    }
-
-    fun collectRootsAndCheckNew(): ScriptClassRootsCache.Updates {
-        val old = classpathRoots
-        val new = recreateRootsCache()
-        classpathRoots = new
-        return new.diff(old)
-    }
-
     override fun updateScriptDefinitionReferences() {
         ScriptDependenciesModificationTracker.getInstance(project).incModificationCount()
 
-        default.updateScriptDefinitions()
+        default.updateScriptDefinitionsReferences()
 
         if (classpathRoots.customDefinitionsUsed) {
             updater.ensureUpdateScheduled()
@@ -112,15 +102,19 @@ class CompositeScriptConfigurationManager(val project: Project) : ScriptConfigur
 
     init {
         val connection = project.messageBus.connect(project)
+
         connection.subscribe(ProjectTopics.PROJECT_ROOTS, object : ModuleRootListener {
             override fun rootsChanged(event: ModuleRootEvent) {
                 if (event.isCausedByFileTypesChange) return
-
-                if (classpathRoots.hasInvalidSdk(project)) {
-                    updater.ensureUpdateScheduled()
-                }
+                updater.checkInvalidSdks()
             }
         })
+
+        connection.subscribe(JDK_TABLE_TOPIC, object : ProjectJdkTable.Listener {
+            override fun jdkAdded(jdk: Sdk) = updater.checkInvalidSdks()
+            override fun jdkNameChanged(jdk: Sdk, previousName: String) = updater.checkInvalidSdks()
+            override fun jdkRemoved(jdk: Sdk) = updater.checkInvalidSdks(remove = jdk)
+        })
     }
 
     /**

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/DefaultScriptingSupport.kt

@@ -24,6 +24,8 @@ import org.jetbrains.kotlin.idea.core.script.configuration.loader.DefaultScriptC
 import org.jetbrains.kotlin.idea.core.script.configuration.loader.ScriptConfigurationLoader
 import org.jetbrains.kotlin.idea.core.script.configuration.loader.ScriptConfigurationLoadingContext
 import org.jetbrains.kotlin.idea.core.script.configuration.loader.ScriptOutsiderFileConfigurationLoader
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsBuilder
+import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsStorage
 import org.jetbrains.kotlin.idea.core.script.configuration.utils.*
 import org.jetbrains.kotlin.idea.core.script.settings.KotlinScriptingSettings
 import org.jetbrains.kotlin.idea.core.util.EDT
@@ -91,8 +93,8 @@ import kotlin.script.experimental.api.ScriptDiagnostic
 class DefaultScriptingSupport(manager: CompositeScriptConfigurationManager) : DefaultScriptingSupportBase(manager) {
     // TODO public for tests
     val backgroundExecutor: BackgroundExecutor =
-        if (ApplicationManager.getApplication().isUnitTestMode) TestingBackgroundExecutor(manager.updater)
-        else DefaultBackgroundExecutor(project, manager.updater)
+        if (ApplicationManager.getApplication().isUnitTestMode) TestingBackgroundExecutor(manager)
+        else DefaultBackgroundExecutor(project, manager)
 
     private val outsiderLoader = ScriptOutsiderFileConfigurationLoader(project)
     private val fileAttributeCache = ScriptConfigurationFileAttributeCache(project)
@@ -121,7 +123,7 @@ class DefaultScriptingSupport(manager: CompositeScriptConfigurationManager) : De
      *
      * Each files may be in on of the states described below:
      * - scriptDefinition is not ready. `ScriptDefinitionsManager.getInstance(project).isReady() == false`.
-     * [updateScriptDefinitions] will be called when [ScriptDefinitionsManager] will be ready
+     * [updateScriptDefinitionsReferences] will be called when [ScriptDefinitionsManager] will be ready
      * which will call [reloadOutOfDateConfiguration] for opened editors.
      * - unknown. When [isFirstLoad] true (`cache[file] == null`).
      * - up-to-date. `cache[file]?.upToDate == true`.
@@ -509,22 +511,20 @@ abstract class DefaultScriptingSupportBase(val manager: CompositeScriptConfigura
         }
     }
 
-    fun updateScriptDefinitions() {
+    fun updateScriptDefinitionsReferences() {
         cache.clear()
     }
 
-    fun collectConfigurations(builder: ScriptClassRootsCache.Builder) {
+    fun collectConfigurations(builder: ScriptClassRootsBuilder) {
         // todo: drop the hell below
         // keep this one only:
         // cache.allApplied().forEach { (vFile, configuration) -> builder.add(vFile, configuration) }
 
         // own builder for saving to storage
-        val ownBuilder = ScriptClassRootsCache.Builder(project)
         val rootsStorage = ScriptClassRootsStorage.getInstance(project)
-
-        rootsStorage.load(ownBuilder)
+        val ownBuilder = ScriptClassRootsBuilder.fromStorage(rootsStorage)
         cache.allApplied().forEach { (vFile, configuration) -> ownBuilder.add(vFile, configuration) }
-        rootsStorage.save(ownBuilder)
+        ownBuilder.toStorage(rootsStorage)
 
         builder.add(ownBuilder)
     }

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/ScriptingSupport.kt

@@ -7,8 +7,9 @@ package org.jetbrains.kotlin.idea.core.script.configuration
 
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.vfs.VirtualFile
-import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsCache
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsCache
 import org.jetbrains.kotlin.idea.core.script.configuration.listener.ScriptChangeListener
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsBuilder
 import org.jetbrains.kotlin.psi.KtFile
 
 /**
@@ -33,7 +34,7 @@ import org.jetbrains.kotlin.psi.KtFile
 abstract class ScriptingSupport {
     abstract fun isApplicable(file: VirtualFile): Boolean
     abstract fun isConfigurationLoadingInProgress(file: KtFile): Boolean
-    abstract fun collectConfigurations(builder: ScriptClassRootsCache.Builder)
+    abstract fun collectConfigurations(builder: ScriptClassRootsBuilder)
 
     companion object {
         val EPN: ExtensionPointName<ScriptingSupport> =

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/utils/DefaultBackgroundExecutor.kt

@@ -12,6 +12,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.util.Alarm
 import com.intellij.util.containers.HashSetQueue
+import org.jetbrains.kotlin.idea.core.script.configuration.CompositeScriptConfigurationManager
 import org.jetbrains.kotlin.idea.core.script.debug
 import org.jetbrains.kotlin.idea.core.util.KotlinIdeaCoreBundle
 import java.util.*
@@ -33,13 +34,14 @@ import javax.swing.SwingUtilities
  */
 internal class DefaultBackgroundExecutor(
     val project: Project,
-    val rootsManager: ScriptClassRootsUpdater
+    val manager: CompositeScriptConfigurationManager
 ) : BackgroundExecutor {
     companion object {
         const val PROGRESS_INDICATOR_DELAY = 1000
         const val PROGRESS_INDICATOR_MIN_QUEUE = 3
     }
 
+    val rootsManager get() = manager.updater
     private val work = Any()
     private val queue: Queue<LoadTask> = HashSetQueue()
 

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/utils/ScriptClassRootsCache.kt

@@ -1,257 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.idea.core.script.configuration.utils
-
-import com.intellij.openapi.module.ModuleManager
-import com.intellij.openapi.project.Project
-import com.intellij.openapi.projectRoots.JavaSdkType
-import com.intellij.openapi.projectRoots.Sdk
-import com.intellij.openapi.roots.ModuleRootManager
-import com.intellij.openapi.roots.OrderRootType
-import com.intellij.openapi.roots.ProjectRootManager
-import com.intellij.openapi.util.io.FileUtil.toSystemIndependentName
-import com.intellij.openapi.vfs.VfsUtil
-import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.psi.search.GlobalSearchScope
-import com.intellij.psi.search.NonClasspathDirectoriesScope.compose
-import org.jetbrains.kotlin.idea.caches.project.getAllProjectSdks
-import org.jetbrains.kotlin.idea.core.script.LOG
-import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager.Companion.classpathEntryToVfs
-import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager.Companion.toVfsRoots
-import org.jetbrains.kotlin.idea.util.getProjectJdkTableSafe
-import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
-import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
-import java.io.File
-import java.lang.ref.Reference
-import java.lang.ref.SoftReference
-
-class ScriptClassRootsCache(
-    val scripts: Map<String, LightScriptInfo>,
-    val classes: Set<String>,
-    val sources: Set<String>,
-    val sdks: Map<String?, Sdk?>,
-    private val nonModulesSdks: Collection<Sdk>,
-    val customDefinitionsUsed: Boolean
-) {
-    abstract class LightScriptInfo(val definition: ScriptDefinition?) {
-        @Volatile
-        var heavyCache: Reference<HeavyScriptInfo>? = null
-
-        abstract fun buildConfiguration(): ScriptCompilationConfigurationWrapper?
-    }
-
-    class DirectScriptInfo(val result: ScriptCompilationConfigurationWrapper) : LightScriptInfo(null) {
-        override fun buildConfiguration(): ScriptCompilationConfigurationWrapper = result
-    }
-
-    class HeavyScriptInfo(
-        val scriptConfiguration: ScriptCompilationConfigurationWrapper,
-        val classFilesScope: GlobalSearchScope,
-        val sdk: Sdk?
-    )
-
-    class Builder(val project: Project) {
-        val scripts = mutableMapOf<String, LightScriptInfo>()
-
-        private val defaultSdk = getScriptDefaultSdk()
-        val sdks = mutableMapOf<String?, Sdk?>(null to defaultSdk)
-
-        val classes = mutableSetOf<String>()
-        val sources = mutableSetOf<String>()
-        private var customDefinitionsUsed: Boolean = false
-
-        fun build(): ScriptClassRootsCache {
-            val nonIndexedSdks = collectNonIndexedSdks()
-            return ScriptClassRootsCache(scripts, classes, sources, sdks, nonIndexedSdks, customDefinitionsUsed)
-        }
-
-        fun collectNonIndexedSdks(): MutableSet<Sdk> {
-            val nonIndexedSdks = sdks.values.filterNotNullTo(mutableSetOf())
-            ModuleManager.getInstance(project).modules.map {
-                nonIndexedSdks.remove(ModuleRootManager.getInstance(it).sdk)
-            }
-            return nonIndexedSdks
-        }
-
-        fun useCustomScriptDefinition() {
-            customDefinitionsUsed = true
-        }
-
-        fun add(
-            vFile: VirtualFile,
-            configuration: ScriptCompilationConfigurationWrapper
-        ) {
-            addSdk(configuration.javaHome)
-
-            configuration.dependenciesClassPath.forEach { classes.add(it.absolutePath) }
-            configuration.dependenciesSources.forEach { sources.add(it.absolutePath) }
-
-            scripts[vFile.path] = DirectScriptInfo(configuration)
-
-            useCustomScriptDefinition()
-        }
-
-        fun add(other: Builder) {
-            classes.addAll(other.classes)
-            sources.addAll(other.sources)
-            sdks.putAll(other.sdks)
-            scripts.putAll(other.scripts)
-        }
-
-        fun addSdk(javaHome: File?): Sdk? {
-            if (javaHome == null) return defaultSdk
-            val canonicalPath = toSystemIndependentName(javaHome.canonicalPath)
-            return sdks.getOrPut(canonicalPath) {
-                getScriptSdkByJavaHome(javaHome) ?: defaultSdk
-            }
-        }
-
-        private fun getScriptSdkByJavaHome(javaHome: File): Sdk? {
-            // workaround for mismatched gradle wrapper and plugin version
-            val javaHomeVF = try {
-                VfsUtil.findFileByIoFile(javaHome, true)
-            } catch (e: Throwable) {
-                null
-            } ?: return null
-
-            return getProjectJdkTableSafe().allJdks.find { it.homeDirectory == javaHomeVF }
-        }
-
-        fun addSdkByName(sdkName: String) {
-            val sdk = getProjectJdkTableSafe().allJdks.find { it.name == sdkName } ?: defaultSdk ?: return
-            val homePath = sdk.homePath ?: return
-            sdks[homePath] = sdk
-        }
-
-        private fun getScriptDefaultSdk(): Sdk? {
-            val projectSdk = ProjectRootManager.getInstance(project).projectSdk?.takeIf { it.canBeUsedForScript() }
-            if (projectSdk != null) return projectSdk
-
-            val anyJavaSdk = getAllProjectSdks().find { it.canBeUsedForScript() }
-            if (anyJavaSdk != null) {
-                return anyJavaSdk
-            }
-
-            LOG.warn(
-                "Default Script SDK is null: " +
-                        "projectSdk = ${ProjectRootManager.getInstance(project).projectSdk}, " +
-                        "all sdks = ${getAllProjectSdks().joinToString("\n")}"
-            )
-
-            return null
-        }
-
-        private fun Sdk.canBeUsedForScript() = sdkType is JavaSdkType
-    }
-
-    fun getLightScriptInfo(file: String) = scripts[file]
-
-    fun contains(file: VirtualFile): Boolean = file.path in scripts
-
-    private fun getHeavyScriptInfo(file: String): HeavyScriptInfo? {
-        val lightScriptInfo = getLightScriptInfo(file) ?: return null
-        val heavy0 = lightScriptInfo.heavyCache?.get()
-        if (heavy0 != null) return heavy0
-        synchronized(lightScriptInfo) {
-            val heavy1 = lightScriptInfo.heavyCache?.get()
-            if (heavy1 != null) return heavy1
-            val heavy2 = computeHeavy(lightScriptInfo)
-            lightScriptInfo.heavyCache = SoftReference(heavy2)
-            return heavy2
-        }
-    }
-
-    private fun computeHeavy(lightScriptInfo: LightScriptInfo): HeavyScriptInfo? {
-        val configuration = lightScriptInfo.buildConfiguration() ?: return null
-
-        val roots = configuration.dependenciesClassPath
-        val sdk = sdks[configuration.javaHome?.canonicalPath]
-
-        return if (sdk == null) {
-            HeavyScriptInfo(configuration, compose(toVfsRoots(roots)), null)
-        } else {
-            val sdkClasses = sdk.rootProvider.getFiles(OrderRootType.CLASSES).toList()
-            HeavyScriptInfo(configuration, compose(sdkClasses + toVfsRoots(roots)), sdk)
-        }
-    }
-
-    val firstScriptSdk: Sdk? = sdks.values.firstOrNull()
-
-    val allDependenciesClassFiles: List<VirtualFile>
-
-    val allDependenciesSources: List<VirtualFile>
-
-    init {
-        allDependenciesClassFiles = mutableSetOf<VirtualFile>().also { result ->
-            nonModulesSdks.forEach { result.addAll(it.rootProvider.getFiles(OrderRootType.CLASSES)) }
-            classes.mapNotNullTo(result) { classpathEntryToVfs(File(it)) }
-        }.toList()
-
-        allDependenciesSources = mutableSetOf<VirtualFile>().also { result ->
-            nonModulesSdks.forEach { result.addAll(it.rootProvider.getFiles(OrderRootType.SOURCES)) }
-            sources.mapNotNullTo(result) { classpathEntryToVfs(File(it)) }
-        }.toList()
-    }
-
-    val allDependenciesClassFilesScope = compose(allDependenciesClassFiles)
-
-    val allDependenciesSourcesScope = compose(allDependenciesSources)
-
-    fun getScriptConfiguration(file: VirtualFile): ScriptCompilationConfigurationWrapper? =
-        getHeavyScriptInfo(file.path)?.scriptConfiguration
-
-    fun getScriptSdk(file: VirtualFile): Sdk? =
-        getHeavyScriptInfo(file.path)?.sdk
-
-    fun getScriptDependenciesClassFilesScope(file: VirtualFile): GlobalSearchScope =
-        getHeavyScriptInfo(file.path)?.classFilesScope ?: GlobalSearchScope.EMPTY_SCOPE
-
-    fun hasInvalidSdk(project: Project): Boolean {
-        val builder = Builder(project)
-        if (sdks.any { (home, sdk) -> builder.addSdk(home?.let(::File)) != sdk }) return true
-        if (builder.collectNonIndexedSdks() != nonModulesSdks) return true
-        return false
-    }
-
-    fun diff(old: ScriptClassRootsCache): Updates =
-        Updates(
-            hasNewRoots(old),
-            getChangedScripts(old)
-        )
-
-    private fun hasNewRoots(old: ScriptClassRootsCache): Boolean {
-        return classes.any { it !in old.classes }
-                || sources.any { it !in old.sources }
-                || sdks.any { it.key !in old.sdks }
-    }
-
-    private fun getChangedScripts(old: ScriptClassRootsCache): Set<String> {
-        val changed = mutableSetOf<String>()
-
-        scripts.forEach {
-            if (old.scripts[it.key] != it.value) {
-                changed.add(it.key)
-            }
-        }
-
-        old.scripts.forEach {
-            if (it.key !in scripts) {
-                changed.add(it.key)
-            }
-        }
-
-        return changed
-    }
-
-    class Updates(
-        val hasNewRoots: Boolean,
-        val updatedScripts: Set<String>
-    ) {
-        val changed: Boolean
-            get() = hasNewRoots || updatedScripts.isNotEmpty()
-    }
-}
-

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/utils/ScriptClassRootsStorage.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -19,9 +19,9 @@ import com.intellij.util.xmlb.XmlSerializerUtil
     storages = [Storage(StoragePathMacros.CACHE_FILE)]
 )
 class ScriptClassRootsStorage(val project: Project) : PersistentStateComponent<ScriptClassRootsStorage> {
-    private var classpath: Set<String> = hashSetOf()
-    private var sources: Set<String> = hashSetOf()
-    private var sdks: Set<String> = hashSetOf()
+    var classpath: Set<String> = hashSetOf()
+    var sources: Set<String> = hashSetOf()
+    var sdks: Set<String> = hashSetOf()
 
     override fun getState(): ScriptClassRootsStorage? {
         return this
@@ -31,18 +31,6 @@ class ScriptClassRootsStorage(val project: Project) : PersistentStateComponent<S
         XmlSerializerUtil.copyBean(state, this)
     }
 
-    fun save(roots: ScriptClassRootsCache.Builder) {
-        classpath = roots.classes
-        sources = roots.sources
-        sdks = roots.sdks.values.mapNotNullTo(mutableSetOf()) { it?.name }
-    }
-
-    fun load(builder: ScriptClassRootsCache.Builder) {
-        builder.sources.addAll(sources)
-        builder.classes.addAll(classpath)
-        sdks.forEach(builder::addSdkByName)
-    }
-
     companion object {
         fun getInstance(project: Project): ScriptClassRootsStorage =
             ServiceManager.getService(project, ScriptClassRootsStorage::class.java)

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/utils/TestingBackgroundExecutor.kt

@@ -8,10 +8,12 @@ package org.jetbrains.kotlin.idea.core.script.configuration.utils
 import com.intellij.openapi.application.impl.LaterInvocator
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.util.containers.HashSetQueue
+import org.jetbrains.kotlin.idea.core.script.configuration.CompositeScriptConfigurationManager
 
 class TestingBackgroundExecutor internal constructor(
-    private val rootsManager: ScriptClassRootsUpdater
+    private val manager: CompositeScriptConfigurationManager
 ) : BackgroundExecutor {
+    val rootsManager get() = manager.updater
     val backgroundQueue = HashSetQueue<BackgroundTask>()
 
     class BackgroundTask(val file: VirtualFile, val actions: () -> Unit) {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/uсache/ScriptClassRootsBuilder.kt

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.core.script.uсache
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsStorage
+import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
+
+class ScriptClassRootsBuilder(
+    val project: Project,
+    private val classes: MutableSet<String> = mutableSetOf(),
+    private val sources: MutableSet<String> = mutableSetOf(),
+    private val scripts: MutableMap<String, ScriptClassRootsCache.LightScriptInfo> = mutableMapOf()
+) {
+    val sdks = ScriptSdksBuilder(project)
+
+    private var customDefinitionsUsed: Boolean = false
+
+    fun build(): ScriptClassRootsCache =
+        ScriptClassRootsCache(
+            project, scripts, classes, sources,
+            customDefinitionsUsed, sdks.build()
+        )
+
+    fun useCustomScriptDefinition() {
+        customDefinitionsUsed = true
+    }
+
+    fun add(
+        vFile: VirtualFile,
+        configuration: ScriptCompilationConfigurationWrapper
+    ) {
+        sdks.addSdk(configuration.javaHome)
+
+        configuration.dependenciesClassPath.forEach { classes.add(it.absolutePath) }
+        configuration.dependenciesSources.forEach { sources.add(it.absolutePath) }
+
+        scripts[vFile.path] = ScriptClassRootsCache.DirectScriptInfo(configuration)
+
+        useCustomScriptDefinition()
+    }
+
+    fun addCustom(
+        path: String,
+        scriptClassesRoots: Collection<String>,
+        sourceSourcesRoots: Collection<String>,
+        info: ScriptClassRootsCache.LightScriptInfo
+    ) {
+        classes.addAll(scriptClassesRoots)
+        sources.addAll(sourceSourcesRoots)
+        scripts[path] = info
+    }
+
+    fun addTemplateClassesRoots(classesRoots: Collection<String>) {
+        classes.addAll(classesRoots)
+    }
+
+    fun add(other: ScriptClassRootsBuilder) {
+        classes.addAll(other.classes)
+        sources.addAll(other.sources)
+        sdks.addAll(other.sdks)
+        scripts.putAll(other.scripts)
+    }
+
+    fun toStorage(storage: ScriptClassRootsStorage) {
+        storage.classpath = classes
+        storage.sources = sources
+        sdks.toStorage(storage)
+    }
+
+    companion object {
+        fun fromStorage(storage: ScriptClassRootsStorage) = ScriptClassRootsBuilder(
+            storage.project,
+            storage.classpath.toMutableSet(),
+            storage.sources.toMutableSet(),
+        ).also { roots ->
+            storage.sdks.forEach {
+                roots.sdks.addSdkByName(it)
+            }
+        }
+    }
+}

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/uсache/ScriptClassRootsCache.kt

@@ -0,0 +1,180 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.core.script.uсache
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.roots.OrderRootType
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.NonClasspathDirectoriesScope.compose
+import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager.Companion.classpathEntryToVfs
+import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager.Companion.toVfsRoots
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
+import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
+import java.io.File
+import java.lang.ref.Reference
+import java.lang.ref.SoftReference
+
+class ScriptClassRootsCache(
+    val project: Project,
+    private val scripts: Map<String, LightScriptInfo>,
+    private val classes: Set<String>,
+    private val sources: Set<String>,
+    val customDefinitionsUsed: Boolean,
+    val sdks: ScriptSdks
+) {
+    fun withUpdatedSdks(newSdks: ScriptSdks) =
+        ScriptClassRootsCache(project, scripts, classes, sources, customDefinitionsUsed, newSdks)
+
+    abstract class LightScriptInfo(val definition: ScriptDefinition?) {
+        @Volatile
+        var heavyCache: Reference<HeavyScriptInfo>? = null
+
+        abstract fun buildConfiguration(): ScriptCompilationConfigurationWrapper?
+    }
+
+    class DirectScriptInfo(val result: ScriptCompilationConfigurationWrapper) : LightScriptInfo(null) {
+        override fun buildConfiguration(): ScriptCompilationConfigurationWrapper = result
+    }
+
+    class HeavyScriptInfo(
+        val scriptConfiguration: ScriptCompilationConfigurationWrapper,
+        val classFilesScope: GlobalSearchScope,
+        val sdk: Sdk?
+    )
+
+    fun getLightScriptInfo(file: String) = scripts[file]
+
+    fun contains(file: VirtualFile): Boolean = file.path in scripts
+
+    private fun getHeavyScriptInfo(file: String): HeavyScriptInfo? {
+        val lightScriptInfo = getLightScriptInfo(file) ?: return null
+        val heavy0 = lightScriptInfo.heavyCache?.get()
+        if (heavy0 != null) return heavy0
+        synchronized(lightScriptInfo) {
+            val heavy1 = lightScriptInfo.heavyCache?.get()
+            if (heavy1 != null) return heavy1
+            val heavy2 = computeHeavy(lightScriptInfo)
+            lightScriptInfo.heavyCache = SoftReference(heavy2)
+            return heavy2
+        }
+    }
+
+    private fun computeHeavy(lightScriptInfo: LightScriptInfo): HeavyScriptInfo? {
+        val configuration = lightScriptInfo.buildConfiguration() ?: return null
+
+        val roots = configuration.dependenciesClassPath
+        val sdk = sdks[configuration.javaHome?.canonicalPath]
+
+        return if (sdk == null) {
+            HeavyScriptInfo(configuration, compose(toVfsRoots(roots)), null)
+        } else {
+            val sdkClasses = sdk.rootProvider.getFiles(OrderRootType.CLASSES).toList()
+            HeavyScriptInfo(configuration, compose(sdkClasses + toVfsRoots(roots)), sdk)
+        }
+    }
+
+    val firstScriptSdk: Sdk?
+        get() = sdks.first
+
+    val allDependenciesClassFiles: List<VirtualFile>
+
+    val allDependenciesSources: List<VirtualFile>
+
+    init {
+        allDependenciesClassFiles = mutableSetOf<VirtualFile>().also { result ->
+            result.addAll(sdks.nonIndexedClassRoots)
+            classes.mapNotNullTo(result) { classpathEntryToVfs(File(it)) }
+        }.toList()
+
+        allDependenciesSources = mutableSetOf<VirtualFile>().also { result ->
+            result.addAll(sdks.nonIndexedSourceRoots)
+            sources.mapNotNullTo(result) { classpathEntryToVfs(File(it)) }
+        }.toList()
+    }
+
+    val allDependenciesClassFilesScope = compose(allDependenciesClassFiles)
+
+    val allDependenciesSourcesScope = compose(allDependenciesSources)
+
+    fun getScriptConfiguration(file: VirtualFile): ScriptCompilationConfigurationWrapper? =
+        getHeavyScriptInfo(file.path)?.scriptConfiguration
+
+    fun getScriptSdk(file: VirtualFile): Sdk? =
+        getHeavyScriptInfo(file.path)?.sdk
+
+    fun getScriptDependenciesClassFilesScope(file: VirtualFile): GlobalSearchScope =
+        getHeavyScriptInfo(file.path)?.classFilesScope ?: GlobalSearchScope.EMPTY_SCOPE
+
+    fun diff(old: ScriptClassRootsCache?): Updates =
+        when (old) {
+            null -> FullUpdate(this)
+            this -> NotChanged(this)
+            else -> IncrementalUpdates(this, hasNewRoots(old), getChangedScripts(old))
+        }
+
+    private fun hasNewRoots(old: ScriptClassRootsCache): Boolean {
+        val oldClassRoots = old.allDependenciesClassFiles.toSet()
+        val oldSourceRoots = old.allDependenciesSources.toSet()
+
+        return allDependenciesClassFiles.any { it !in oldClassRoots }
+                || allDependenciesSources.any { it !in oldSourceRoots }
+    }
+
+    private fun getChangedScripts(old: ScriptClassRootsCache): Set<String> {
+        val changed = mutableSetOf<String>()
+
+        scripts.forEach {
+            if (old.scripts[it.key] != it.value) {
+                changed.add(it.key)
+            }
+        }
+
+        old.scripts.forEach {
+            if (it.key !in scripts) {
+                changed.add(it.key)
+            }
+        }
+
+        return changed
+    }
+
+    interface Updates {
+        val cache: ScriptClassRootsCache
+        val changed: Boolean
+        val hasNewRoots: Boolean
+        val hasUpdatedScripts: Boolean
+        fun isScriptChanged(scriptPath: String): Boolean
+    }
+
+    class IncrementalUpdates(
+        override val cache: ScriptClassRootsCache,
+        override val hasNewRoots: Boolean,
+        private val updatedScripts: Set<String>
+    ) : Updates {
+        override val hasUpdatedScripts: Boolean get() = updatedScripts.isNotEmpty()
+        override fun isScriptChanged(scriptPath: String) = scriptPath in updatedScripts
+
+        override val changed: Boolean
+            get() = hasNewRoots || updatedScripts.isNotEmpty()
+    }
+
+    class FullUpdate(override val cache: ScriptClassRootsCache) : Updates {
+        override val changed: Boolean get() = true
+        override val hasNewRoots: Boolean get() = true
+        override val hasUpdatedScripts: Boolean get() = true
+        override fun isScriptChanged(scriptPath: String): Boolean = true
+    }
+
+    class NotChanged(override val cache: ScriptClassRootsCache) : Updates {
+        override val changed: Boolean get() = false
+        override val hasNewRoots: Boolean get() = false
+        override val hasUpdatedScripts: Boolean get() = false
+        override fun isScriptChanged(scriptPath: String): Boolean = false
+    }
+}
+

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/uсache/ScriptClassRootsUpdater.kt

@@ -1,20 +1,21 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package org.jetbrains.kotlin.idea.core.script.configuration.utils
+package org.jetbrains.kotlin.idea.core.script.uсache
 
 import com.intellij.codeInsight.daemon.DaemonCodeAnalyzer
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.TransactionGuard
 import com.intellij.openapi.fileEditor.FileEditorManager
-import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.progress.util.BackgroundTaskUtil
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.roots.ex.ProjectRootManagerEx
+import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.EmptyRunnable
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElementFinder
@@ -26,6 +27,7 @@ import org.jetbrains.kotlin.idea.core.script.ScriptDependenciesModificationTrack
 import org.jetbrains.kotlin.idea.core.script.configuration.CompositeScriptConfigurationManager
 import org.jetbrains.kotlin.idea.core.script.debug
 import org.jetbrains.kotlin.idea.core.util.EDT
+import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.locks.ReentrantLock
@@ -36,12 +38,23 @@ import kotlin.concurrent.withLock
  */
 class ScriptClassRootsUpdater(
     val project: Project,
-    val manager: CompositeScriptConfigurationManager
+    val manager: CompositeScriptConfigurationManager,
+    private val buildNewRoots: (ScriptClassRootsBuilder) -> Unit
 ) {
+    private var lastSeen: ScriptClassRootsCache? = null
     private var invalidated: Boolean = false
     private var syncUpdateRequired: Boolean = false
     private val concurrentUpdates = AtomicInteger()
 
+    private fun recreateRootsCache(): ScriptClassRootsCache {
+        val builder = ScriptClassRootsBuilder(project)
+        buildNewRoots(builder)
+        return builder.build()
+    }
+
+    var classpathRoots: ScriptClassRootsCache = recreateRootsCache()
+        private set
+
     /**
      * @param synchronous Used from legacy FS cache only, don't use
      */
@@ -108,8 +121,12 @@ class ScriptClassRootsUpdater(
     @Synchronized
     fun ensureUpdateScheduled() {
         scheduledUpdate?.cancel()
-        scheduledUpdate = BackgroundTaskUtil.executeOnPooledThread(project) {
-            doUpdate()
+        runReadAction {
+            if (project.isDisposed && !Disposer.isDisposing(project)) {
+                scheduledUpdate = BackgroundTaskUtil.executeOnPooledThread(project) {
+                    doUpdate()
+                }
+            }
         }
     }
 
@@ -119,10 +136,22 @@ class ScriptClassRootsUpdater(
         doUpdate(false)
     }
 
+    fun checkInvalidSdks(remove: Sdk? = null) {
+        // sdks should be updated synchronously to avoid disposed roots usage
+        syncLock.withLock {
+            val current = classpathRoots
+            val actualSdks = current.sdks.rebuild(remove = remove)
+            if (actualSdks != current.sdks) {
+                classpathRoots = current.withUpdatedSdks(actualSdks)
+                ensureUpdateScheduled()
+            }
+        }
+    }
+
     private fun doUpdate(underProgressManager: Boolean = true) {
         syncLock.withLock {
             try {
-                val updates = manager.collectRootsAndCheckNew()
+                val updates = recreateRootsCacheAndDiff()
 
                 if (!updates.changed) return
 
@@ -130,6 +159,8 @@ class ScriptClassRootsUpdater(
                     ProgressManager.checkCanceled()
                 }
 
+                if (project.isDisposed) return
+
                 if (updates.hasNewRoots) {
                     notifyRootsChanged()
                 }
@@ -139,13 +170,13 @@ class ScriptClassRootsUpdater(
 
                 ScriptDependenciesModificationTracker.getInstance(project).incModificationCount()
 
-                if (updates.updatedScripts.isNotEmpty()) {
+                if (updates.hasUpdatedScripts) {
                     updateHighlighting(project) {
-                        it.path in updates.updatedScripts
+                        updates.isScriptChanged(it.path)
                     }
                 }
-            } catch (cancel: ProcessCanceledException) {
-                if (underProgressManager) throw cancel
+
+                lastSeen = updates.cache
             } finally {
                 synchronized(this) {
                     scheduledUpdate = null
@@ -154,6 +185,12 @@ class ScriptClassRootsUpdater(
         }
     }
 
+    private fun recreateRootsCacheAndDiff(): ScriptClassRootsCache.Updates {
+        val new = recreateRootsCache()
+        classpathRoots = new
+        return new.diff(lastSeen)
+    }
+
     private fun notifyRootsChanged() {
         val doNotifyRootsChanged = Runnable {
             runWriteAction {
@@ -173,7 +210,7 @@ class ScriptClassRootsUpdater(
         }
     }
 
-    fun updateHighlighting(project: Project, filter: (VirtualFile) -> Boolean) {
+    private fun updateHighlighting(project: Project, filter: (VirtualFile) -> Boolean) {
         if (!project.isOpen) return
 
         val openFiles = FileEditorManager.getInstance(project).openFiles

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/uсache/ScriptSdks.kt

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.core.script.uсache
+
+import com.intellij.openapi.module.ModuleManager
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.roots.ModuleRootManager
+import com.intellij.openapi.roots.OrderRootType
+import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.idea.util.application.runReadAction
+import java.io.File
+
+class ScriptSdks(val project: Project, private val sdks: Map<String?, Sdk?>) {
+    fun rebuild(remove: Sdk?): ScriptSdks {
+        val builder = ScriptSdksBuilder(project, remove = remove)
+        sdks.keys.forEach { home ->
+            if (home != null) {
+                builder.addSdk(File(home))
+            }
+        }
+        return builder.build()
+    }
+
+    val nonIndexedClassRoots = mutableSetOf<VirtualFile>()
+    val nonIndexedSourceRoots = mutableSetOf<VirtualFile>()
+
+    val first: Sdk? = sdks.values.firstOrNull()
+    operator fun get(sdkHome: String?) = sdks[sdkHome]
+
+    init {
+        val nonIndexedSdks = sdks.values.filterNotNullTo(mutableSetOf())
+
+        runReadAction {
+            ModuleManager.getInstance(project).modules.map {
+                nonIndexedSdks.remove(ModuleRootManager.getInstance(it).sdk)
+            }
+
+            nonIndexedSdks.forEach {
+                nonIndexedClassRoots.addAll(it.rootProvider.getFiles(OrderRootType.CLASSES))
+                nonIndexedSourceRoots.addAll(it.rootProvider.getFiles(OrderRootType.SOURCES))
+            }
+        }
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as ScriptSdks
+
+        if (nonIndexedClassRoots != other.nonIndexedClassRoots) return false
+        if (nonIndexedSourceRoots != other.nonIndexedSourceRoots) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = nonIndexedClassRoots.hashCode()
+        result = 31 * result + nonIndexedSourceRoots.hashCode()
+        return result
+    }
+}

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/uсache/ScriptSdksBuilder.kt

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.core.script.uсache
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.projectRoots.JavaSdkType
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.roots.ProjectRootManager
+import com.intellij.openapi.util.io.FileUtil
+import com.intellij.openapi.vfs.VfsUtil
+import org.jetbrains.kotlin.idea.caches.project.getAllProjectSdks
+import org.jetbrains.kotlin.idea.core.script.LOG
+import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsStorage
+import org.jetbrains.kotlin.idea.util.getProjectJdkTableSafe
+import java.io.File
+
+class ScriptSdksBuilder(
+    val project: Project,
+    private val sdks: MutableMap<String?, Sdk?> = mutableMapOf(),
+    private val remove: Sdk? = null
+) {
+    private val defaultSdk = getScriptDefaultSdk()
+
+    init {
+        sdks[null] = defaultSdk
+    }
+
+    fun build() = ScriptSdks(project, sdks)
+
+    fun addAll(other: ScriptSdksBuilder) {
+        sdks.putAll(other.sdks)
+    }
+
+    fun addSdk(javaHome: File?): Sdk? {
+        if (javaHome == null) return defaultSdk
+        val canonicalPath = FileUtil.toSystemIndependentName(javaHome.canonicalPath)
+        return sdks.getOrPut(canonicalPath) {
+            getScriptSdkByJavaHome(javaHome) ?: defaultSdk
+        }
+    }
+
+    private fun getScriptSdkByJavaHome(javaHome: File): Sdk? {
+        // workaround for mismatched gradle wrapper and plugin version
+        val javaHomeVF = try {
+            VfsUtil.findFileByIoFile(javaHome, true)
+        } catch (e: Throwable) {
+            null
+        } ?: return null
+
+        return getProjectJdkTableSafe().allJdks.find { it.homeDirectory == javaHomeVF }
+            ?.takeIf { it.canBeUsedForScript() }
+    }
+
+    fun addSdkByName(sdkName: String) {
+        val sdk = getProjectJdkTableSafe().allJdks
+            .find { it.name == sdkName }
+            ?.takeIf { it.canBeUsedForScript() }
+            ?: defaultSdk ?: return
+        val homePath = sdk.homePath ?: return
+        sdks[homePath] = sdk
+    }
+
+    private fun getScriptDefaultSdk(): Sdk? {
+        val projectSdk = ProjectRootManager.getInstance(project).projectSdk?.takeIf { it.canBeUsedForScript() }
+        if (projectSdk != null) return projectSdk
+
+        val anyJavaSdk = getAllProjectSdks().find { it.canBeUsedForScript() }
+        if (anyJavaSdk != null) {
+            return anyJavaSdk
+        }
+
+        LOG.warn(
+            "Default Script SDK is null: " +
+                    "projectSdk = ${ProjectRootManager.getInstance(project).projectSdk}, " +
+                    "all sdks = ${getAllProjectSdks().joinToString("\n")}"
+        )
+
+        return null
+    }
+
+    private fun Sdk.canBeUsedForScript() = this != remove && sdkType is JavaSdkType
+
+    fun toStorage(storage: ScriptClassRootsStorage) {
+        storage.sdks = sdks.values.mapNotNullTo(mutableSetOf()) { it?.name }
+    }
+}

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/roots/GradleBuildRoot.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.idea.scripting.gradle.roots
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VirtualFile
-import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsCache
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsBuilder
 import org.jetbrains.kotlin.idea.scripting.gradle.GradleScriptDefinitionsContributor
 import org.jetbrains.kotlin.idea.scripting.gradle.LastModifiedFiles
 import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslScriptModel
@@ -121,21 +121,24 @@ sealed class GradleBuildRoot {
             loadLastModifiedFiles()
         }
 
-        fun collectConfigurations(builder: ScriptClassRootsCache.Builder) {
+        fun collectConfigurations(builder: ScriptClassRootsBuilder) {
             if (javaHome != null) {
-                builder.addSdk(javaHome)
+                builder.sdks.addSdk(javaHome)
             }
 
             val definitions = GradleScriptDefinitionsContributor.getDefinitions(project)
 
-            builder.classes.addAll(data.templateClasspath)
+            builder.addTemplateClassesRoots(data.templateClasspath)
+
             data.models.forEach { script ->
                 val definition = selectScriptDefinition(script, definitions)
 
-                builder.scripts[script.file] = GradleScriptInfo(this, definition, script)
-
-                builder.classes.addAll(script.classPath)
-                builder.sources.addAll(script.sourcePath)
+                builder.addCustom(
+                    script.file,
+                    script.classPath,
+                    script.sourcePath,
+                    GradleScriptInfo(this, definition, script)
+                )
             }
         }
 

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/roots/GradleBuildRootsManager.kt

@@ -20,7 +20,7 @@ import kotlinx.coroutines.launch
 import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager
 import org.jetbrains.kotlin.idea.core.script.configuration.CompositeScriptConfigurationManager
 import org.jetbrains.kotlin.idea.core.script.configuration.ScriptingSupport
-import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsCache
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsBuilder
 import org.jetbrains.kotlin.idea.core.util.EDT
 import org.jetbrains.kotlin.idea.scripting.gradle.*
 import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslGradleBuildSync
@@ -88,7 +88,7 @@ class GradleBuildRootsManager(val project: Project) : ScriptingSupport() {
         return !script.model.inputs.isUpToDate(project, file)
     }
 
-    override fun collectConfigurations(builder: ScriptClassRootsCache.Builder) {
+    override fun collectConfigurations(builder: ScriptClassRootsBuilder) {
         roots.list.forEach { root ->
             if (root is GradleBuildRoot.Imported) {
                 root.collectConfigurations(builder)

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/roots/GradleScriptInfo.kt

@@ -6,7 +6,7 @@
 package org.jetbrains.kotlin.idea.scripting.gradle.roots
 
 import com.intellij.openapi.vfs.VfsUtil
-import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsCache
+import org.jetbrains.kotlin.idea.core.script.uсache.ScriptClassRootsCache
 import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslScriptModel
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper

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

@@ -226,7 +226,10 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
         }
         runWriteAction {
             val jdk = PluginTestCaseBase.jdk(jdkKind)
-            getProjectJdkTableSafe().addJdk(jdk, testRootDisposable)
+            val projectJdkTable = getProjectJdkTableSafe()
+            if (jdk !in projectJdkTable.allJdks) {
+                projectJdkTable.addJdk(jdk, testRootDisposable)
+            }
             env["javaHome"] = File(jdk.homePath)
         }