Scripting: Fix caching of packages in JvmPackagePartProvider

since in scripting scenarios the loadedModules could be updated
later in the pipeline, previous caching approach led to problems.

#KT-76009 fixed

Author: ligee

compiler/cli/cli-base/src/org/jetbrains/kotlin/cli/jvm/compiler/JvmPackagePartProvider.kt

@@ -31,6 +31,18 @@ class JvmPackagePartProvider(
 
     override val loadedModules: MutableList<ModuleMappingInfo<VirtualFile>> = SmartList()
 
+    private var allPackageNamesCache: Set<String>? = null
+
+    override val allPackageNames: Set<String>
+        get() {
+            // assuming that the modifications of loadedModules happen in predictable moments now, so nos syncronization is used
+            if (allPackageNamesCache == null) {
+                allPackageNamesCache = loadedModules.flatMapTo(mutableSetOf<String>()) { it.mapping.packageFqName2Parts.keys }
+            }
+            return allPackageNamesCache!!
+        }
+
+    // TODO: redesign to avoid cache-unfriendly usages, see KT-76516
     fun addRoots(roots: List<JavaRoot>, messageCollector: MessageCollector) {
         for ((root, type) in roots) {
             if (type != JavaRoot.RootType.BINARY) continue
@@ -45,6 +57,7 @@ class JvmPackagePartProvider(
                     deserializationConfiguration, messageCollector
                 )?.let {
                     loadedModules.add(ModuleMappingInfo(root, it, moduleFile.nameWithoutExtension))
+                    allPackageNamesCache = null
                 }
             }
         }

core/descriptors.jvm/src/org/jetbrains/kotlin/load/kotlin/JvmPackagePartProviderBase.kt

@@ -38,7 +38,7 @@ abstract class JvmPackagePartProviderBase<MappingsKey> : PackageAndMetadataPartP
         return result.toList()
     }
 
-    private val allPackageNames: Set<String> by lazy {
+    protected open val allPackageNames: Set<String> by lazy {
         loadedModules.flatMapTo(mutableSetOf()) { it.mapping.packageFqName2Parts.keys }
     }
 

plugins/scripting/scripting-compiler/build.gradle.kts

@@ -44,6 +44,8 @@ dependencies {
     testApi(libs.junit.platform.launcher)
     testApi(kotlinTest("junit5"))
 
+    testApi(project(":kotlin-scripting-dependencies-maven"))
+
     testImplementation(intellijCore())
     testImplementation(libs.kotlinx.coroutines.core)
     testImplementation(commonDependency("org.jetbrains.kotlin:kotlin-reflect")) { isTransitive = false }

plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/test/CustomK2ReplTest.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.scripting.compiler.test
 
+import kotlinx.coroutines.runBlocking
 import org.jetbrains.kotlin.codegen.forTestCompile.ForTestCompileRuntime
 import org.jetbrains.kotlin.scripting.compiler.plugin.SCRIPT_TEST_BASE_COMPILER_ARGUMENTS_PROPERTY
 import org.jetbrains.kotlin.scripting.compiler.plugin.impl.K2ReplCompiler
@@ -14,12 +15,18 @@ import java.io.File
 import kotlin.reflect.full.declaredMemberFunctions
 import kotlin.reflect.full.declaredMemberProperties
 import kotlin.script.experimental.api.*
+import kotlin.script.experimental.dependencies.CompoundDependenciesResolver
+import kotlin.script.experimental.dependencies.FileSystemDependenciesResolver
+import kotlin.script.experimental.dependencies.maven.MavenDependenciesResolver
 import kotlin.script.experimental.host.toScriptSource
 import kotlin.script.experimental.impl.internalScriptingRunSuspend
 import kotlin.script.experimental.jvm.KJvmEvaluatedSnippet
 import kotlin.script.experimental.jvm.updateClasspath
 import kotlin.script.experimental.util.LinkedSnippet
-import kotlin.test.*
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+import kotlin.test.fail
 
 class ReplReceiver1 {
     val ok = "OK"
@@ -28,6 +35,7 @@ class ReplReceiver1 {
 @Suppress("unused") // Used in snippets
 class TestReplReceiver1() { fun checkReceiver(block: ReplReceiver1.() -> Any) = block(ReplReceiver1()) }
 
+private val dependenciesResolver = CompoundDependenciesResolver(FileSystemDependenciesResolver(), MavenDependenciesResolver())
 
 class CustomK2ReplTest {
 
@@ -118,6 +126,38 @@ class CustomK2ReplTest {
         )
     }
 
+    @Test
+    fun testWithUpdatingDependeciesAndImportKotlinDeclarations() {
+        evalAndCheckSnippetsResultVals(
+            sequenceOf(
+                "println(\"firstLine\")",
+                "import org.jetbrains.kotlinx.dataframe.jupyter.KotlinNotebookPluginUtils",
+                "KotlinNotebookPluginUtils.getKotlinNotebookIDEBuildNumber().toString()",
+                "import org.jetbrains.kotlinx.dataframe.jupyter.importDataSchema",
+                "importDataSchema(\"ftp://xx\").url.toString()",
+            ),
+            sequenceOf(null, null, "null", null, "ftp://xx"),
+            baseCompilationConfiguration.with {
+                refineConfiguration {
+                    beforeCompiling { (script, config, _) ->
+                        if (!script.text.contains("firstLine")) {
+                            val resolveResults = runBlocking {
+                                dependenciesResolver.resolve("org.jetbrains.kotlinx:dataframe-core:0.15.0")
+                            }
+                            if (resolveResults is ResultWithDiagnostics.Failure)
+                                resolveResults
+                            else
+                                config.with {
+                                    updateClasspath(resolveResults.valueOrThrow())
+                                    defaultImports("kotlin.random.Random")
+                                }.asSuccess()
+                        } else config.asSuccess()
+                    }
+                }
+            }
+        )
+    }
+
     @Test
     fun testBasicReflection() {
         evalAndCheckSnippets(