Added support for source JARs with Maven

Author: daplf

server/src/main/kotlin/org/javacs/kt/CompilerClassPath.kt

@@ -1,9 +1,10 @@
 package org.javacs.kt
 
+import org.javacs.kt.classpath.ClassPathEntry
 import org.javacs.kt.classpath.defaultClassPathResolver
 import org.javacs.kt.compiler.Compiler
+import org.javacs.kt.util.AsyncExecutor
 import java.io.Closeable
-import java.nio.file.Files
 import java.nio.file.FileSystems
 import java.nio.file.Path
 
@@ -14,10 +15,11 @@ import java.nio.file.Path
 class CompilerClassPath(private val config: CompilerConfiguration) : Closeable {
     private val workspaceRoots = mutableSetOf<Path>()
     private val javaSourcePath = mutableSetOf<Path>()
-    private val classPath = mutableSetOf<Path>()
+    val classPath = mutableSetOf<ClassPathEntry>()
     private val buildScriptClassPath = mutableSetOf<Path>()
-    var compiler = Compiler(javaSourcePath, classPath, buildScriptClassPath)
+    var compiler = Compiler(javaSourcePath, classPath.map { it.compiledJar }.toSet(), buildScriptClassPath)
         private set
+    private val async = AsyncExecutor()
 
     init {
         compiler.updateConfiguration(config)
@@ -36,9 +38,18 @@ class CompilerClassPath(private val config: CompilerConfiguration) : Closeable {
         if (updateClassPath) {
             val newClassPath = resolver.classpathOrEmpty
             if (newClassPath != classPath) {
-                syncPaths(classPath, newClassPath, "class path")
+                synchronized(classPath) {
+                    syncClassPathEntries(classPath, newClassPath, "class path")
+                }
                 refreshCompiler = true
             }
+
+            async.compute {
+                val newClassPathWithSources = resolver.fetchClasspathWithSources()
+                synchronized(classPath) {
+                    syncClassPathEntries(classPath, newClassPathWithSources, "class path")
+                }
+            }
         }
 
         if (updateBuildScriptClassPath) {
@@ -53,13 +64,25 @@ class CompilerClassPath(private val config: CompilerConfiguration) : Closeable {
         if (refreshCompiler) {
             LOG.info("Reinstantiating compiler")
             compiler.close()
-            compiler = Compiler(javaSourcePath, classPath, buildScriptClassPath)
+            compiler = Compiler(javaSourcePath, classPath.map { it.compiledJar }.toSet(), buildScriptClassPath)
             updateCompilerConfiguration()
         }
 
         return refreshCompiler
     }
 
+    /** Synchronizes the given two class path entry sets and logs the differences. */
+    private fun syncClassPathEntries(dest: MutableSet<ClassPathEntry>, new: Set<ClassPathEntry>, name: String) {
+        val added = new - dest
+        val removed = dest - new
+
+        logAdded(added.map { it.compiledJar }, name)
+        logRemoved(removed.map { it.compiledJar }, name)
+
+        dest.removeAll(removed)
+        dest.addAll(added)
+    }
+
     /** Synchronizes the given two path sets and logs the differences. */
     private fun syncPaths(dest: MutableSet<Path>, new: Set<Path>, name: String) {
         val added = new - dest

server/src/main/kotlin/org/javacs/kt/KotlinLanguageServer.kt

@@ -28,7 +28,7 @@ class KotlinLanguageServer : LanguageServer, LanguageClientAware, Closeable {
     val sourcePath = SourcePath(classPath, uriContentProvider, config.indexing)
     val sourceFiles = SourceFiles(sourcePath, uriContentProvider)
 
-    private val textDocuments = KotlinTextDocumentService(sourceFiles, sourcePath, config, tempDirectory, uriContentProvider)
+    private val textDocuments = KotlinTextDocumentService(sourceFiles, sourcePath, config, tempDirectory, uriContentProvider, classPath)
     private val workspaces = KotlinWorkspaceService(sourceFiles, sourcePath, classPath, textDocuments, config)
     private val protocolExtensions = KotlinProtocolExtensionService(uriContentProvider)
 

server/src/main/kotlin/org/javacs/kt/KotlinTextDocumentService.kt

@@ -38,7 +38,8 @@ class KotlinTextDocumentService(
     private val sp: SourcePath,
     private val config: Configuration,
     private val tempDirectory: TemporaryDirectory,
-    private val uriContentProvider: URIContentProvider
+    private val uriContentProvider: URIContentProvider,
+    private val classPath: CompilerClassPath
 ) : TextDocumentService, Closeable {
     private lateinit var client: LanguageClient
     private val async = AsyncExecutor()
@@ -124,7 +125,7 @@ class KotlinTextDocumentService(
             LOG.info("Go-to-definition at {}", describePosition(position))
 
             val (file, cursor) = recover(position, Recompile.NEVER)
-            goToDefinition(file, cursor, uriContentProvider.jarClassContentProvider, tempDirectory, config.externalSources)
+            goToDefinition(file, cursor, uriContentProvider.jarClassContentProvider, tempDirectory, config.externalSources, classPath.classPath)
                 ?.let(::listOf)
                 ?.let { Either.forLeft<List<Location>, List<LocationLink>>(it) }
                 ?: noResult("Couldn't find definition at ${describePosition(position)}", Either.forLeft(emptyList()))

server/src/main/kotlin/org/javacs/kt/URIContentProvider.kt

@@ -7,6 +7,7 @@ import java.nio.file.Paths
 import org.javacs.kt.externalsources.JarClassContentProvider
 import org.javacs.kt.externalsources.toKlsURI
 import org.javacs.kt.util.KotlinLSException
+import org.javacs.kt.util.partitionAroundLast
 
 /**
  * Fetches the content of Kotlin files identified by a URI.
@@ -16,7 +17,7 @@ class URIContentProvider(
 ) {
     fun contentOf(uri: URI): String = when (uri.scheme) {
         "file" -> Paths.get(uri).toFile().readText()
-        "kls" -> uri.toKlsURI()?.let { jarClassContentProvider.contentOf(it).second }
+        "kls" -> uri.toKlsURI()?.let { jarClassContentProvider.contentOf(it, it.source).second }
             ?: throw KotlinLSException("Could not find ${uri}")
         else -> throw KotlinLSException("Unrecognized scheme ${uri.scheme}")
     }

server/src/main/kotlin/org/javacs/kt/definition/GoToDefinition.kt

@@ -8,6 +8,7 @@ import java.nio.file.Files
 import org.javacs.kt.CompiledFile
 import org.javacs.kt.LOG
 import org.javacs.kt.ExternalSourcesConfiguration
+import org.javacs.kt.classpath.ClassPathEntry
 import org.javacs.kt.externalsources.JarClassContentProvider
 import org.javacs.kt.externalsources.toKlsURI
 import org.javacs.kt.externalsources.KlsURI
@@ -20,6 +21,9 @@ import org.javacs.kt.util.parseURI
 import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
 import org.jetbrains.kotlin.psi.KtNamedDeclaration
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
+import org.jetbrains.kotlin.js.parser.parse
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import java.nio.file.Paths
 
 private val cachedTempFiles = mutableMapOf<KlsURI, Path>()
 private val definitionPattern = Regex("(?:class|interface|object|fun)\\s+(\\w+)")
@@ -29,7 +33,8 @@ fun goToDefinition(
     cursor: Int,
     jarClassContentProvider: JarClassContentProvider,
     tempDir: TemporaryDirectory,
-    config: ExternalSourcesConfiguration
+    config: ExternalSourcesConfiguration,
+    compilerClassPath: Set<ClassPathEntry>
 ): Location? {
     val (_, target) = file.referenceAtPoint(cursor) ?: return null
 
@@ -45,8 +50,10 @@ fun goToDefinition(
         val rawClassURI = destination.uri
 
         if (isInsideJar(rawClassURI)) {
-            parseURI(rawClassURI).toKlsURI()?.let { klsURI ->
-                val (klsSourceURI, content) = jarClassContentProvider.contentOf(klsURI)
+            val sourceURI = getSourceURI(rawClassURI, compilerClassPath)
+            val actualClassURI = sourceURI ?: rawClassURI
+            parseURI(actualClassURI).toKlsURI()?.let { klsURI ->
+                val (klsSourceURI, content) = jarClassContentProvider.contentOf(klsURI, sourceURI != null)
 
                 if (config.useKlsScheme) {
                     // Defer decompilation until a jarClassContents request is sent
@@ -86,4 +93,13 @@ fun goToDefinition(
     return destination
 }
 
+private fun getSourceURI(rawClassURI: String, compilerClassPath: Set<ClassPathEntry>): String? {
+    val rawClassPath = parseURI(rawClassURI).path.toString()
+    val jarUri = rawClassPath.substring(0, rawClassPath.indexOf("!"))
+    val classPartUri = rawClassPath.substring(rawClassPath.indexOf("!"))
+    val sourceJar = compilerClassPath.find { it.compiledJar.toUri().path == jarUri }?.sourceJar
+
+    return if (sourceJar != null) sourceJar.toUri().toString() + classPartUri else null
+}
+
 private fun isInsideJar(uri: String) = uri.contains(".jar!")

server/src/main/kotlin/org/javacs/kt/externalsources/JarClassContentProvider.kt

@@ -22,27 +22,30 @@ class JarClassContentProvider(
     private val tempDir: TemporaryDirectory,
     private val decompiler: Decompiler = FernflowerDecompiler()
 ) {
-    /** Maps recently used (source-)KLS-URIs to their source contents (e.g. decompiled code). */
-    private val cachedContents = object : LinkedHashMap<String, String>() {
-        override fun removeEldestEntry(eldest: MutableMap.MutableEntry<String, String>) = size > 5
+    /** Maps recently used (source-)KLS-URIs to their source contents (e.g. decompiled code) and the file extension. */
+    private val cachedContents = object : LinkedHashMap<String, Pair<String, String>>() {
+        override fun removeEldestEntry(eldest: MutableMap.MutableEntry<String, Pair<String, String>>) = size > 5
     }
 
     /**
      * Fetches the contents of a compiled class/source file in a JAR
      * and another URI which can be used to refer to these extracted
      * contents.
+     * If the file is inside a source JAR, the source code is returned as is.
      */
-    public fun contentOf(uri: KlsURI): Pair<KlsURI, String> {
-        val sourceURI = uri.withFileExtension(if (config.autoConvertToKotlin) "kt" else "java")
-        val key = sourceURI.toString()
-        val contents: String = cachedContents[key] ?: run {
-            LOG.info("Finding contents of {}", describeURI(uri.uri))
-            tryReadContentOf(uri)
-                ?: tryReadContentOf(uri.withFileExtension("class"))
-                ?: tryReadContentOf(uri.withFileExtension("java"))
-                ?: tryReadContentOf(uri.withFileExtension("kt"))
-                ?: throw KotlinLSException("Could not find $uri")
-        }.also { cachedContents[key] = it }
+    public fun contentOf(uri: KlsURI, source: Boolean): Pair<KlsURI, String> {
+        val key = uri.toString()
+        val (contents, extension) = cachedContents[key] ?: run {
+                LOG.info("Finding contents of {}", describeURI(uri.uri))
+                tryReadContentOf(uri, source)
+                    ?: tryReadContentOf(uri.withFileExtension("class"), source)
+                    ?: tryReadContentOf(uri.withFileExtension("java"), source)
+                    ?: tryReadContentOf(uri.withFileExtension("kt"), source)
+                    ?: throw KotlinLSException("Could not find $uri")
+            }.also {
+                cachedContents[key] = Pair(it.first, it.second)
+            }
+        val sourceURI = uri.withFileExtension(extension).withSource(source)
         return Pair(sourceURI, contents)
     }
 
@@ -52,16 +55,17 @@ class JarClassContentProvider(
         javaCode
     }
 
-    private fun tryReadContentOf(uri: KlsURI): String? = try {
-        when (uri.fileExtension) {
-            "class" -> uri.extractToTemporaryFile(tempDir)
+    private fun tryReadContentOf(uri: KlsURI, source: Boolean): Pair<String, String>? = try {
+        val actualUri = uri.withoutQuery()
+        when (actualUri.fileExtension) {
+            "class" -> Pair(actualUri.extractToTemporaryFile(tempDir)
                 .let(decompiler::decompileClass)
                 .let { Files.newInputStream(it) }
                 .bufferedReader()
                 .use(BufferedReader::readText)
-                .let(this::convertToKotlinIfNeeded)
-            "java" -> convertToKotlinIfNeeded(uri.readContents())
-            else -> uri.readContents() // e.g. for Kotlin source files
+                .let(this::convertToKotlinIfNeeded), if (config.autoConvertToKotlin) "kt" else "java")
+            "java" -> if (source) Pair(actualUri.readContents(), "java") else Pair(convertToKotlinIfNeeded(actualUri.readContents()), "kt")
+            else -> Pair(actualUri.readContents(), "kt") // e.g. for Kotlin source files
         }
     } catch (e: FileNotFoundException) { null }
 }

server/src/main/kotlin/org/javacs/kt/externalsources/KlsURI.kt

@@ -37,16 +37,28 @@ data class KlsURI(val uri: URI) {
             .split(".")
             .takeIf { it.size > 1 }
             ?.lastOrNull()
+    val source: Boolean get() = uri.schemeSpecificPart.split("?").getOrNull(1)?.split("&")?.find {
+        it.matches("source=(true|false)".toRegex())
+    }?.split("=")?.get(1)?.toBoolean() ?: false
     val isCompiled: Boolean
         get() = fileExtension == "class"
 
     fun withFileExtension(newExtension: String): KlsURI {
-        val (parentURI, fileName) = uri.toString().partitionAroundLast("/")
-        val newURI = "$parentURI${fileName.split(".").first()}.$newExtension"
+        val (parentURI, fileNamePlusQuery) = uri.toString().partitionAroundLast("/")
+        val (fileName, query) = if (fileNamePlusQuery.contains("?")) fileNamePlusQuery.partitionAroundLast("?") else Pair(fileNamePlusQuery, "")
+        val newURI = "$parentURI${fileName.split(".").first()}.$newExtension$query"
         return KlsURI(URI(newURI))
     }
 
-    private fun toFileURI(): URI = URI(uri.schemeSpecificPart)
+    fun withSource(source: Boolean): KlsURI {
+        return KlsURI(URI("${uri.toString()}?source=$source"))
+    }
+
+    fun withoutQuery(): KlsURI {
+        return KlsURI(URI(uri.toString().split("?")[0]))
+    }
+
+    fun toFileURI(): URI = URI(uri.schemeSpecificPart)
 
     private fun toJarURL(): URL = URL("jar:${uri.schemeSpecificPart}")
 

shared/src/main/kotlin/org/javacs/kt/classpath/BackupClassPathResolver.kt

@@ -12,7 +12,7 @@ import java.nio.file.Paths
 /** Backup classpath that find Kotlin in the user's Maven/Gradle home or kotlinc's libraries folder. */
 object BackupClassPathResolver : ClassPathResolver {
     override val resolverType: String = "Backup"
-    override val classpath: Set<Path> get() = findKotlinStdlib()?.let { setOf(it) }.orEmpty()
+    override val classpath: Set<ClassPathEntry> get() = findKotlinStdlib()?.let { setOf(it) }.orEmpty().map { ClassPathEntry(it, null) }.toSet()
 }
 
 fun findKotlinStdlib(): Path? =

shared/src/main/kotlin/org/javacs/kt/classpath/ClassPathResolver.kt

@@ -3,17 +3,19 @@ package org.javacs.kt.classpath
 import org.javacs.kt.LOG
 import java.nio.file.Path
 
+data class ClassPathEntry(val compiledJar: Path, val sourceJar: Path?)
+
 /** A source for creating class paths */
 interface ClassPathResolver {
     val resolverType: String
 
-    val classpath: Set<Path> // may throw exceptions
-    val classpathOrEmpty: Set<Path> // does not throw exceptions
+    val classpath: Set<ClassPathEntry> // may throw exceptions
+    val classpathOrEmpty: Set<ClassPathEntry> // does not throw exceptions
         get() = try {
             classpath
         } catch (e: Exception) {
             LOG.warn("Could not resolve classpath using {}: {}", resolverType, e.message)
-            emptySet<Path>()
+            emptySet<ClassPathEntry>()
         }
 
     val buildScriptClasspath: Set<Path>
@@ -26,11 +28,13 @@ interface ClassPathResolver {
             emptySet<Path>()
         }
 
+    fun fetchClasspathWithSources(): Set<ClassPathEntry> = classpath
+
     companion object {
         /** A default empty classpath implementation */
         val empty = object : ClassPathResolver {
             override val resolverType = "[]"
-            override val classpath = emptySet<Path>()
+            override val classpath = emptySet<ClassPathEntry>()
         }
     }
 }
@@ -52,6 +56,7 @@ internal class UnionClassPathResolver(val lhs: ClassPathResolver, val rhs: Class
     override val classpathOrEmpty get() = lhs.classpathOrEmpty + rhs.classpathOrEmpty
     override val buildScriptClasspath get() = lhs.buildScriptClasspath + rhs.buildScriptClasspath
     override val buildScriptClasspathOrEmpty get() = lhs.buildScriptClasspathOrEmpty + rhs.buildScriptClasspathOrEmpty
+    override fun fetchClasspathWithSources() = lhs.fetchClasspathWithSources() + rhs.fetchClasspathWithSources()
 }
 
 internal class FirstNonEmptyClassPathResolver(val lhs: ClassPathResolver, val rhs: ClassPathResolver) : ClassPathResolver {
@@ -60,4 +65,5 @@ internal class FirstNonEmptyClassPathResolver(val lhs: ClassPathResolver, val rh
     override val classpathOrEmpty get() = lhs.classpathOrEmpty.takeIf { it.isNotEmpty() } ?: rhs.classpathOrEmpty
     override val buildScriptClasspath get() = lhs.buildScriptClasspath.takeIf { it.isNotEmpty() } ?: rhs.buildScriptClasspath
     override val buildScriptClasspathOrEmpty get() = lhs.buildScriptClasspathOrEmpty.takeIf { it.isNotEmpty() } ?: rhs.buildScriptClasspathOrEmpty
+    override fun fetchClasspathWithSources() = lhs.fetchClasspathWithSources().takeIf { it.isNotEmpty() } ?: rhs.fetchClasspathWithSources()
 }

shared/src/main/kotlin/org/javacs/kt/classpath/GradleClassPathResolver.kt

@@ -1,8 +1,6 @@
 package org.javacs.kt.classpath
 
 import org.javacs.kt.LOG
-import org.javacs.kt.util.firstNonNull
-import org.javacs.kt.util.tryResolving
 import org.javacs.kt.util.execAndReadStdoutAndStderr
 import org.javacs.kt.util.KotlinLSException
 import org.javacs.kt.util.isOSWindows
@@ -15,12 +13,13 @@ import java.nio.file.Paths
 internal class GradleClassPathResolver(private val path: Path, private val includeKotlinDSL: Boolean): ClassPathResolver {
     override val resolverType: String = "Gradle"
     private val projectDirectory: Path get() = path.getParent()
-    override val classpath: Set<Path> get() {
+    override val classpath: Set<ClassPathEntry> get() {
         val scripts = listOf("projectClassPathFinder.gradle")
         val tasks = listOf("kotlinLSPProjectDeps")
 
         return readDependenciesViaGradleCLI(projectDirectory, scripts, tasks)
             .apply { if (isNotEmpty()) LOG.info("Successfully resolved dependencies for '${projectDirectory.fileName}' using Gradle") }
+            .map { ClassPathEntry(it, null) }.toSet()
     }
     override val buildScriptClasspath: Set<Path> get() {
         return if (includeKotlinDSL) {