KTIJ-36997 Refactor Gradle script storage update logic

- Consolidate `updateStorage` methods into `getUpdatedStorage`.
- Refactor related logic for improved clarity and immutability.
- Optimize module-related lookups in the workspace model.

GitOrigin-RevId: 178a445fe7599a21e98da431a85d2817b4bf13de

plugins/kotlin/base/fir/scripting/src/org/jetbrains/kotlin/base/fir/scripting/projectStructure/FirKaScriptingModuleFactory.kt

@@ -5,6 +5,7 @@ import com.intellij.injected.editor.VirtualFileWindow
 import com.intellij.openapi.project.Project
 import com.intellij.platform.backend.workspace.WorkspaceModel
 import com.intellij.platform.backend.workspace.toVirtualFileUrl
+import com.intellij.platform.backend.workspace.workspaceModel
 import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaLibraryModule
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
@@ -59,9 +60,9 @@ internal class FirKaScriptingModuleFactory : FirKaModuleFactory {
         val project = ktFile.project
         val virtualFile = file.originalFile.virtualFile
 
+        val snapshot = project.workspaceModel.currentSnapshot
         if (!nameSequence.endsWith(KotlinFileType.DOT_SCRIPT_EXTENSION)) {
             val workspaceModel = WorkspaceModel.getInstance(project)
-            val snapshot = workspaceModel.currentSnapshot
 
             val url = virtualFile.toVirtualFileUrl(workspaceModel.getVirtualFileUrlManager())
             val entitiesByUrl = snapshot.getVirtualFileUrlIndex().findEntitiesByUrl(url)
@@ -70,6 +71,6 @@ internal class FirKaScriptingModuleFactory : FirKaModuleFactory {
             }
         }
 
-        return KaScriptModuleImpl(project, virtualFile)
+        return KaScriptModuleImpl(project, virtualFile, snapshot)
     }
 }

plugins/kotlin/base/fir/scripting/src/org/jetbrains/kotlin/base/fir/scripting/projectStructure/modules/KaScriptModuleBase.kt

@@ -24,23 +24,19 @@ import java.util.*
 
 abstract class KaScriptModuleBase(
     override val project: Project,
-    open val virtualFile: VirtualFile,
+    protected open val virtualFile: VirtualFile,
 ) : KaScriptModule, KaModuleBase() {
+    protected open val snapshot: ImmutableEntityStorage by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        project.workspaceModel.currentSnapshot
+    }
+
     protected val virtualFileUrlManager: VirtualFileUrlManager
         get() = project.workspaceModel.getVirtualFileUrlManager()
 
-    protected val currentSnapshot: ImmutableEntityStorage
-        get() = project.workspaceModel.currentSnapshot
-
     private val scriptDefinition: ScriptDefinition by lazy {
         findScriptDefinition(project, KtFileScriptSource(file))
     }
 
-    protected val kotlinScriptEntity: KotlinScriptEntity?
-        get() = currentSnapshot.getVirtualFileUrlIndex()
-            .findEntitiesByUrl(virtualFile.toVirtualFileUrl(virtualFileUrlManager))
-            .filterIsInstance<KotlinScriptEntity>().firstOrNull()
-
     override val directDependsOnDependencies: List<KaModule> get() = emptyList()
 
     override val transitiveDependsOnDependencies: List<KaModule> get() = emptyList()

plugins/kotlin/base/fir/scripting/src/org/jetbrains/kotlin/base/fir/scripting/projectStructure/modules/KaScriptModuleImpl.kt

@@ -3,9 +3,13 @@ package org.jetbrains.kotlin.base.fir.scripting.projectStructure.modules
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.findPsiFile
+import com.intellij.platform.backend.workspace.toVirtualFileUrl
+import com.intellij.platform.backend.workspace.workspaceModel
+import com.intellij.platform.workspace.storage.ImmutableEntityStorage
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
 import org.jetbrains.kotlin.idea.base.projectStructure.*
 import org.jetbrains.kotlin.idea.core.script.k2.modules.K2IdeScriptAdditionalIdeaDependenciesProvider
+import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptEntity
 import org.jetbrains.kotlin.idea.core.script.v1.ScriptAdditionalIdeaDependenciesProvider
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -15,7 +19,14 @@ import org.jetbrains.kotlin.utils.exceptions.withVirtualFileEntry
 internal class KaScriptModuleImpl(
     override val project: Project,
     override val virtualFile: VirtualFile,
+    override val snapshot: ImmutableEntityStorage,
 ) : KaScriptModuleBase(project, virtualFile) {
+    val kotlinScriptEntity by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        snapshot.getVirtualFileUrlIndex()
+            .findEntitiesByUrl(virtualFile.toVirtualFileUrl(virtualFileUrlManager))
+            .filterIsInstance<KotlinScriptEntity>().singleOrNull()
+    }
+
     override val file: KtFile
         get() {
             (virtualFile.findPsiFile(project) as? KtFile)?.let { return it }
@@ -40,7 +51,7 @@ internal class KaScriptModuleImpl(
 
             addAll(
                 K2IdeScriptAdditionalIdeaDependenciesProvider.getRelatedScripts(virtualFile, project)
-                    .map { KaScriptModuleImpl(project, it) }
+                    .map { KaScriptModuleImpl(project, it, snapshot) }
             )
 
             addAll(
@@ -75,7 +86,7 @@ internal class KaScriptModuleImpl(
     }
 
     private fun MutableCollection<KaModule>.addRegularDependencies() {
-        val libraryDependencies = kotlinScriptEntity?.dependencies?.mapNotNull { currentSnapshot.resolve(it) }?.flatMap {
+        val libraryDependencies = kotlinScriptEntity?.dependencies?.mapNotNull { snapshot.resolve(it) }?.flatMap {
             project.ideProjectStructureProvider.getKaScriptLibraryModules(it)
         } ?: emptyList()
 

plugins/kotlin/gradle/scripting/kotlin.gradle.scripting.k2/src/org/jetbrains/kotlin/gradle/scripting/k2/GradleKotlinScriptEntityProvider.kt

@@ -4,22 +4,21 @@ package org.jetbrains.kotlin.gradle.scripting.k2
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.externalSystem.service.execution.ExternalSystemJdkUtil
-import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.openapi.options.advanced.AdvancedSettings
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
-import com.intellij.openapi.util.io.toNioPathOrNull
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.platform.backend.workspace.WorkspaceModelChangeListener
+import com.intellij.platform.backend.workspace.virtualFile
 import com.intellij.platform.backend.workspace.workspaceModel
+import com.intellij.platform.workspace.jps.entities.ContentRootEntity
 import com.intellij.platform.workspace.jps.entities.ModuleId
+import com.intellij.platform.workspace.storage.ImmutableEntityStorage
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import com.intellij.platform.workspace.storage.VersionedStorageChange
 import com.intellij.platform.workspace.storage.toBuilder
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
 import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
-import com.intellij.workspaceModel.ide.legacyBridge.findSnapshotModuleEntity
 import org.jetbrains.kotlin.gradle.scripting.k2.definition.withIdeKeys
 import org.jetbrains.kotlin.gradle.scripting.k2.importing.GradleScriptData
 import org.jetbrains.kotlin.gradle.scripting.k2.importing.GradleScriptModel
@@ -67,24 +66,21 @@ class GradleKotlinScriptEntityProvider(override val project: Project) : KotlinSc
         }
     }
 
-    fun updateStorage(
+    fun getUpdatedStorage(
         scriptData: GradleScriptData,
         storageToUpdate: MutableEntityStorage
-    ) {
+    ): ImmutableEntityStorage {
         val javaHome = scriptData.definitionsParams.javaHome
         val definitions = loadGradleDefinitions(scriptData.definitionsParams).map { it.withIdeKeys(project) }
-        val updatedStorage = MutableEntityStorage.create().apply {
-            updateStorage(this, scriptData.models, definitions, javaHome)
-        }
-        storageToUpdate.replaceBySource({ it is KotlinGradleScriptEntitySource }, updatedStorage)
+        return getUpdatedStorage(storageToUpdate, scriptData.models, definitions, javaHome)
     }
 
-    fun updateStorage(
+    fun getUpdatedStorage(
         storage: MutableEntityStorage,
         models: Collection<GradleScriptModel>,
         definitions: Collection<GradleScriptDefinition>,
         javaHome: String? = null
-    ) {
+    ): ImmutableEntityStorage {
         definitions.forEach {
             storage addEntity GradleScriptDefinitionEntity(
                 it.definitionId,
@@ -114,6 +110,8 @@ class GradleKotlinScriptEntityProvider(override val project: Project) : KotlinSc
             val configurationResult = refineScriptCompilationConfiguration(sourceCode, definition, project, configuration)
             storage.updateStorage(model.virtualFile, configurationResult, model.classpathModel)
         }
+
+        return storage.toSnapshot()
     }
 
     private fun String?.resolveSdk(): Sdk? {
@@ -177,17 +175,33 @@ class GradleKotlinScriptEntityProvider(override val project: Project) : KotlinSc
             this.configuration = configurationWrapper.configuration?.asEntity()
             this.reports = configurationResult.reports.map(ScriptDiagnostic::map).toMutableList()
             this.sdkId = configurationWrapper.configuration?.sdkId
-            this.relatedModuleIds = classpathModel?.getRelatedModules().orEmpty().toMutableList()
+            this.relatedModuleIds = classpathModel?.let { getRelatedModules(it) }.orEmpty().toMutableList()
         }
     }
 
-    private fun GradleBuildScriptClasspathModel.getRelatedModules(): Sequence<ModuleId> =
-        classpath.flatMap { it.classes + it.sources }.distinct().asSequence()
-            .mapNotNull { it.toNioPathOrNull() }
-            .mapNotNull { VirtualFileManager.getInstance().findFileByNioPath(it) }
-            .mapNotNull { ModuleUtilCore.findModuleForFile(it, project) }
-            .mapNotNull { it.findSnapshotModuleEntity()?.symbolicId }
-            .distinct()
+    private fun MutableEntityStorage.getRelatedModules(classpathModel: GradleBuildScriptClasspathModel): MutableSet<ModuleId> {
+
+        val virtualFileUrls = classpathModel.classpath.flatMap { it.sources }.mapNotNull {
+            it.toVirtualFileUrl(urlManager)
+        }.filter { it.virtualFile != null }
+
+        val result = mutableSetOf<ModuleId>()
+        for (url in virtualFileUrls) {
+            var current: VirtualFileUrl? = url
+            while (current != null && current.url != project.basePath) {
+                val moduleIds = getVirtualFileUrlIndex().findEntitiesByUrl(current)
+                        .filterIsInstance<ContentRootEntity>().map { it.module.symbolicId }.toSet()
+
+                if (result.addAll(moduleIds)) {
+                    break
+                }
+
+                current = current.parent
+            }
+        }
+
+        return result
+    }
 
     private fun MutableEntityStorage.getOrCreateScriptLibrary(
         jar: VirtualFileUrl, sources: Collection<VirtualFileUrl>

plugins/kotlin/gradle/scripting/kotlin.gradle.scripting.k2/src/org/jetbrains/kotlin/gradle/scripting/k2/importing/KotlinDslBaseScriptSyncContributor.kt

@@ -40,24 +40,20 @@ internal class KotlinDslBaseScriptSyncContributor : GradleSyncContributor {
             }
         }
 
-        val builder = storage.toBuilder()
-        builder.also { storage ->
-            val models = readAction {
-                FileEditorManager.getInstance(context.project).allEditors
-            }.filter {
-                it.file.name.endsWith(".gradle.kts")
-            }.map { fileEditor ->
-                GradleScriptModel(
-                    fileEditor.file, baseModel.compileClassPath.map { it.path }, listOf(), baseModel.implicitImports
-                )
-            }
-
-            GradleKotlinScriptEntityProvider.getInstance(context.project).updateStorage(
-                storage,
-                models,
-                definitions,
+        val models = readAction {
+            FileEditorManager.getInstance(context.project).allEditors
+        }.filter {
+            it.file.name.endsWith(".gradle.kts")
+        }.map { fileEditor ->
+            GradleScriptModel(
+                fileEditor.file, baseModel.compileClassPath.map { it.path }, listOf(), baseModel.implicitImports
             )
         }
-        return builder.toSnapshot()
+
+        return GradleKotlinScriptEntityProvider.getInstance(context.project).getUpdatedStorage(
+            storage.toBuilder(),
+            models,
+            definitions,
+        )
     }
 }

plugins/kotlin/gradle/scripting/kotlin.gradle.scripting.k2/src/org/jetbrains/kotlin/gradle/scripting/k2/importing/KotlinDslScriptSyncContributor.kt

@@ -81,9 +81,7 @@ internal class KotlinDslScriptSyncContributor : GradleSyncContributor {
             )
         )
 
-        GradleKotlinScriptEntityProvider.getInstance(project).updateStorage(scriptData, builder)
-
-        return builder.toSnapshot()
+        return GradleKotlinScriptEntityProvider.getInstance(project).getUpdatedStorage(scriptData, builder)
     }
 }