More prep for 2.1.10+ (#1301)

Author: Bencodes

src/main/kotlin/io/bazel/kotlin/plugin/jdeps/k2/JdepsK2Utils.kt

@@ -1,12 +1,27 @@
 package io.bazel.kotlin.plugin.jdeps.k2
 
+import io.bazel.kotlin.plugin.jdeps.k2.RefCache.vbseClass
+import io.bazel.kotlin.plugin.jdeps.k2.RefCache.vbseGetVirtualFileMethod
 import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 
 private object RefCache {
+  val vbseClass: Class<*>? by lazy {
+    runCatching {
+      Class.forName("org.jetbrains.kotlin.fir.java.VirtualFileBasedSourceElement")
+    }.getOrNull()
+  }
+
+  val vbseGetVirtualFileMethod by lazy {
+    vbseClass
+      ?.runCatching {
+        getMethod("getVirtualFile")
+      }?.getOrNull()
+  }
+
   val jbseClass: Class<*>? by lazy {
     runCatching {
       Class.forName("org.jetbrains.kotlin.fir.java.JavaBinarySourceElement")
@@ -28,8 +43,7 @@ private object RefCache {
  * @return whether class is provided by JSM runtime or not
  */
 internal fun isJvmClass(className: String): Boolean =
-  className.startsWith("java") ||
-    className.startsWith("modules/java.base/java/")
+  className.startsWith("java") || className.startsWith("modules/java.base/java/")
 
 internal fun DeserializedContainerSource.classId(): ClassId? =
   when (this) {
@@ -43,6 +57,9 @@ internal fun SourceElement.binaryClass(): String? =
     binaryClass.location
   } else if (this is JvmPackagePartSource) {
     this.knownJvmBinaryClass?.location
+  } else if (vbseClass != null && vbseClass!!.isInstance(this)) {
+    val virtualFile = vbseGetVirtualFileMethod!!.invoke(this)
+    virtualFile?.javaClass!!.getMethod("getPath").invoke(virtualFile) as? String
   } else if (RefCache.jbseClass != null && RefCache.jbseClass!!.isInstance(this)) {
     val jClass = RefCache.jbseGetJavaClassMethod!!.invoke(this)
     val virtualFile = jClass!!.javaClass.getMethod("getVirtualFile").invoke(jClass)