From 89c8392255782c468f4f8c09371be0cdd57bfd1a Mon Sep 17 00:00:00 2001
From: Pavel Mikhailovskii <pavel.mikhailovskii@jetbrains.com>
Date: Tue, 13 Jun 2023 13:34:11 +0000
Subject: [PATCH] Make it possible to instantiate StandaloneAnalysisAPISession
 with a custom clasloader

---
 .../structure/StandaloneProjectFactory.kt     | 21 ++++++++++++++++---
 .../StandaloneAnalysisAPISessionBuilder.kt    |  5 +++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/StandaloneProjectFactory.kt b/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/StandaloneProjectFactory.kt
index 4ea19a828ed..2f89d621c17 100644
--- a/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/StandaloneProjectFactory.kt
+++ b/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/StandaloneProjectFactory.kt
@@ -14,6 +14,7 @@ import com.intellij.ide.highlighter.JavaFileType
 import com.intellij.mock.MockApplication
 import com.intellij.mock.MockProject
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.extensions.PluginDescriptor
 import com.intellij.openapi.roots.PackageIndex
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
@@ -54,6 +55,7 @@ import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.popLast
+import org.picocontainer.PicoContainer
 import java.nio.file.Path
 
 object StandaloneProjectFactory {
@@ -61,6 +63,7 @@ object StandaloneProjectFactory {
         projectDisposable: Disposable,
         applicationDisposable: Disposable,
         compilerConfiguration: CompilerConfiguration = CompilerConfiguration(),
+        classLoader: ClassLoader = MockProject::class.java.classLoader
     ): KotlinCoreProjectEnvironment {
         val applicationEnvironment =
             KotlinCoreEnvironment.getOrCreateApplicationEnvironmentForTests(applicationDisposable, compilerConfiguration)
@@ -68,9 +71,21 @@ object StandaloneProjectFactory {
 
         registerApplicationServices(applicationEnvironment.application)
 
-        return KotlinCoreProjectEnvironment(projectDisposable, applicationEnvironment).apply {
-            registerProjectServices(project)
-            registerJavaPsiFacade(project)
+        return object : KotlinCoreProjectEnvironment(projectDisposable, applicationEnvironment) {
+            init {
+                registerProjectServices(project)
+                registerJavaPsiFacade(project)
+            }
+
+            override fun createProject(parent: PicoContainer, parentDisposable: Disposable): MockProject {
+                return object : MockProject(parent, parentDisposable) {
+                    @Throws(ClassNotFoundException::class)
+                    override fun <T> loadClass(className: String, pluginDescriptor: PluginDescriptor): Class<T> {
+                        @Suppress("UNCHECKED_CAST")
+                        return Class.forName(className, true, classLoader) as Class<T>
+                    }
+                }
+            }
         }
     }
 
diff --git a/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt b/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt
index f0f4cf14856..1d43df53f7f 100644
--- a/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt
+++ b/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.standalone
 
+import com.intellij.mock.MockProject
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.Application
 import com.intellij.openapi.project.Project
@@ -43,6 +44,7 @@ import kotlin.contracts.contract
 public class StandaloneAnalysisAPISessionBuilder(
     applicationDisposable: Disposable,
     projectDisposable: Disposable,
+    classLoader: ClassLoader = MockProject::class.java.classLoader
 ) {
     init {
         // We depend on swing (indirectly through PSI or something), so we want to declare headless mode,
@@ -57,6 +59,7 @@ public class StandaloneAnalysisAPISessionBuilder(
         StandaloneProjectFactory.createProjectEnvironment(
             projectDisposable,
             applicationDisposable,
+            classLoader = classLoader
         )
 
 
@@ -215,6 +218,7 @@ public inline fun buildStandaloneAnalysisAPISession(
     applicationDisposable: Disposable = Disposer.newDisposable("StandaloneAnalysisAPISession.application"),
     projectDisposable: Disposable = Disposer.newDisposable("StandaloneAnalysisAPISession.project"),
     withPsiDeclarationFromBinaryModuleProvider: Boolean = false,
+    classLoader: ClassLoader = MockProject::class.java.classLoader,
     init: StandaloneAnalysisAPISessionBuilder.() -> Unit
 ): StandaloneAnalysisAPISession {
     contract {
@@ -223,6 +227,7 @@ public inline fun buildStandaloneAnalysisAPISession(
     return StandaloneAnalysisAPISessionBuilder(
         applicationDisposable,
         projectDisposable,
+        classLoader
     ).apply(init).build(
         withPsiDeclarationFromBinaryModuleProvider,
     )