From f5df30df90b0aefa1384ae67f32d2655c93e9c5f Mon Sep 17 00:00:00 2001
From: Ilya Muradyan <Ilya.Muradyan@jetbrains.com>
Date: Thu, 28 Apr 2022 22:41:09 +0300
Subject: [PATCH] Fix duplicating class loaders

---
 .../jupyter/repl/impl/JupyterCompilerImpl.kt      | 15 ++++++++-------
 .../test/repl/ReplWithStandardResolverTests.kt    | 11 +++++++++++
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/JupyterCompilerImpl.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/JupyterCompilerImpl.kt
index b3fd10f9f..2812c989e 100644
--- a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/JupyterCompilerImpl.kt
+++ b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/JupyterCompilerImpl.kt
@@ -45,6 +45,7 @@ import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
 import kotlin.script.experimental.jvm.impl.KJvmCompiledScript
 import kotlin.script.experimental.jvm.impl.getOrCreateActualClassloader
 import kotlin.script.experimental.jvm.jvm
+import kotlin.script.experimental.jvm.lastSnippetClassLoader
 import kotlin.script.experimental.jvm.util.toSourceCodePosition
 import kotlin.script.experimental.util.LinkedSnippet
 
@@ -122,9 +123,11 @@ open class JupyterCompilerImpl<CompilerT : ReplCompiler<KJvmCompiledScript>>(
     override val lastKClass: KClass<*>
         get() = classes.last()
 
+    private val _baseClassLoader: ClassLoader
+        get() = basicEvaluationConfiguration[ScriptEvaluationConfiguration.jvm.baseClassLoader]!!
+
     override val lastClassLoader: ClassLoader
-        get() = classes.lastOrNull()?.java?.classLoader
-            ?: basicEvaluationConfiguration[ScriptEvaluationConfiguration.jvm.baseClassLoader]!!
+        get() = classes.lastOrNull()?.java?.classLoader ?: _baseClassLoader
 
     override fun nextCounter() = executionCounter.getAndIncrement()
 
@@ -174,11 +177,9 @@ open class JupyterCompilerImpl<CompilerT : ReplCompiler<KJvmCompiledScript>>(
                 val compiledScript = result.get()
 
                 val configWithClassloader = basicEvaluationConfiguration.with {
-                    val lastClassOrNull = classes.lastOrNull()
-                    if (lastClassOrNull != null) {
-                        jvm {
-                            baseClassLoader(lastClassOrNull.java.classLoader)
-                        }
+                    jvm {
+                        lastSnippetClassLoader(lastClassLoader)
+                        baseClassLoader(_baseClassLoader.parent)
                     }
                 }
                 val classLoader = compiledScript.getOrCreateActualClassloader(configWithClassloader)
diff --git a/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplWithStandardResolverTests.kt b/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplWithStandardResolverTests.kt
index e084d1649..501e63e6f 100644
--- a/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplWithStandardResolverTests.kt
+++ b/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplWithStandardResolverTests.kt
@@ -18,6 +18,7 @@ import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.parallel.Execution
 import org.junit.jupiter.api.parallel.ExecutionMode
+import java.net.URLClassLoader
 import kotlin.concurrent.thread
 import kotlin.test.assertEquals
 
@@ -46,6 +47,12 @@ class ReplWithStandardResolverTests : AbstractSingleReplTest() {
 
     @Test
     fun testStandardLibraryResolver() {
+        val baseClassLoader = repl.currentClassLoader.parent
+        fun urlClassLoadersCount() = generateSequence(repl.currentClassLoader) { classLoader ->
+            classLoader.parent?.takeIf { it != baseClassLoader }
+        }.filter { it is URLClassLoader }.count()
+        urlClassLoadersCount() shouldBe 1
+
         val res = eval(
             """
             %use krangl@d91d045946f59(0.16.2)
@@ -54,6 +61,10 @@ class ReplWithStandardResolverTests : AbstractSingleReplTest() {
             """.trimIndent()
         )
         assertEquals("John Smith", res.resultValue)
+        urlClassLoadersCount() shouldBe 2
+
+        eval("val x = 2 + 2")
+        urlClassLoadersCount() shouldBe 2
     }
 
     @Test