Setup incremental indexing on file changes

daplf · daplf · commit e2e3dc02d5d7 · 2021-12-05T00:55:50.000Z
diff --git a/server/src/main/kotlin/org/javacs/kt/SourcePath.kt b/server/src/main/kotlin/org/javacs/kt/SourcePath.kt
@@ -7,17 +7,14 @@ import org.javacs.kt.util.filePath
 import org.javacs.kt.util.describeURI
 import org.javacs.kt.index.SymbolIndex
 import org.javacs.kt.progress.Progress
-import org.javacs.kt.IndexingConfiguration
 import com.intellij.lang.Language
-import com.intellij.psi.PsiFile
-import com.intellij.openapi.fileTypes.FileType
-import com.intellij.openapi.fileTypes.LanguageFileType
 import org.jetbrains.kotlin.container.ComponentProvider
 import org.jetbrains.kotlin.container.getService
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.CompositeBindingContext
+import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import kotlin.concurrent.withLock
 import java.nio.file.Path
 import java.nio.file.Paths
@@ -99,14 +96,16 @@ class SourcePath(
         private fun doCompile() {
             LOG.debug("Compiling {}", path?.fileName)
 
+            val oldFile = clone()
+
             val (context, container) = cp.compiler.compileKtFile(parsed!!, allIncludingThis(), kind)
             parseDataWriteLock.withLock {
                 compiledContext = context
                 compiledContainer = container
                 compiledFile = parsed
             }
 
-            initializeIndexAsyncIfNeeded(container)
+            refreshIndexes(container, listOfNotNull(oldFile), listOfNotNull(this))
         }
 
         private fun doCompileIfChanged() {
@@ -125,6 +124,9 @@ class SourcePath(
             if (isTemporary) (all().asSequence() + sequenceOf(parsed!!)).toList()
             else all()
         }
+
+        // Creates a shallow copy
+        fun clone(): SourceFile = SourceFile(uri, content, path, parsed, compiledFile, compiledContext, compiledContainer, language, isTemporary)
     }
 
     private fun sourceFile(uri: URI): SourceFile {
@@ -161,6 +163,8 @@ class SourcePath(
         }
 
     fun delete(uri: URI) {
+        files[uri]?.let { refreshIndexes(files[uri]?.compiledContainer!!, listOf(it), listOf()) }
+
         files.remove(uri)
     }
 
@@ -195,7 +199,20 @@ class SourcePath(
         // Compile changed files
         fun compileAndUpdate(changed: List<SourceFile>, kind: CompilationKind): BindingContext? {
             if (changed.isEmpty()) return null
+
+            // Get clones of the old files, so we can remove the old declarations from the index
+            val oldFiles = changed.mapNotNull {
+                if (it.compiledFile?.text != it.content || it.parsed?.text != it.content) {
+                    it.clone()
+                } else {
+                    null
+                }
+            }
+
+            // Parse the files that have changed
             val parse = changed.associateWith { it.apply { parseIfChanged() }.parsed!! }
+
+            // Get all the files. This will parse them if they changed
             val allFiles = all()
             beforeCompileCallback.invoke()
             val (context, container) = cp.compiler.compileKtFiles(parse.values, allFiles, kind)
@@ -214,7 +231,7 @@ class SourcePath(
 
             // Only index normal files, not build files
             if (kind == CompilationKind.DEFAULT) {
-                initializeIndexAsyncIfNeeded(container)
+                refreshIndexes(container, oldFiles, parse.keys.toList())
             }
 
             return context
@@ -231,17 +248,41 @@ class SourcePath(
     }
 
     /**
-     * Initialized the symbol index asynchronously, if not
-     * already done.
+     * Refreshes the indexes. If already done, refreshes only the declarations in the files that were changed.
      */
-    private fun initializeIndexAsyncIfNeeded(container: ComponentProvider) = indexAsync.execute {
-        if (indexEnabled && !indexInitialized) {
-            indexInitialized = true
+    private fun refreshIndexes(container: ComponentProvider, oldFiles: List<SourceFile>, newFiles: List<SourceFile>) = indexAsync.execute {
+        if (indexEnabled) {
             val module = container.getService(ModuleDescriptor::class.java)
-            index.refresh(module)
+            val oldDeclarations = getDeclarationDescriptors(oldFiles)
+            val newDeclarations = getDeclarationDescriptors(newFiles)
+
+            // Index all the declarations except any new declarations that were just compiled
+            // TODO: Move this logic to a different place when re-indexing is triggered on configuration changes
+            if (!indexInitialized) {
+                indexInitialized = true
+                index.refresh(module, newDeclarations)
+            }
+
+            // Index the new declarations in the Kotlin source files that were just compiled, removing the old ones
+            index.updateIndexes(oldDeclarations, newDeclarations)
         }
     }
 
+    // Gets all the declaration descriptors for the collection of files
+    private fun getDeclarationDescriptors(files: Collection<SourceFile>) =
+        files.flatMap { file ->
+            val compiledFile = file.compiledFile ?: file.parsed
+            val compiledContainer = file.compiledContainer
+            if (compiledFile != null && compiledContainer != null) {
+                val module = compiledContainer.getService(ModuleDescriptor::class.java)
+                module.getPackage(compiledFile.packageFqName).memberScope.getContributedDescriptors(
+                    DescriptorKindFilter.ALL
+                ) { name -> compiledFile.declarations.map { it.name }.contains(name.toString()) }
+            } else {
+                listOf()
+            }
+        }.asSequence()
+
     /**
      * Recompiles all source files that are initialized.
      */
diff --git a/server/src/main/kotlin/org/javacs/kt/index/SymbolIndex.kt b/server/src/main/kotlin/org/javacs/kt/index/SymbolIndex.kt
@@ -1,45 +1,79 @@
 package org.javacs.kt.index
 
-import org.jetbrains.exposed.sql.and
-import org.jetbrains.exposed.sql.count
-import org.jetbrains.exposed.sql.deleteAll
-import org.jetbrains.exposed.sql.innerJoin
-import org.jetbrains.exposed.sql.replace
-import org.jetbrains.exposed.sql.select
-import org.jetbrains.exposed.sql.selectAll
-import org.jetbrains.exposed.sql.Table
 import org.jetbrains.exposed.sql.transactions.transaction
-import org.jetbrains.exposed.sql.SchemaUtils
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.psi2ir.intermediate.extensionReceiverType
 import org.javacs.kt.LOG
 import org.javacs.kt.progress.Progress
-import org.jetbrains.exposed.sql.Database
-import org.jetbrains.exposed.sql.SqlExpressionBuilder.like
-import org.jetbrains.exposed.sql.insert
-
-private val MAX_FQNAME_LENGTH = 255
-private val MAX_SHORT_NAME_LENGTH = 80
-
-private object Symbols : Table() {
-    val fqName = varchar("fqname", length = MAX_FQNAME_LENGTH) references FqNames.fqName
+import org.jetbrains.exposed.dao.IntEntity
+import org.jetbrains.exposed.dao.IntEntityClass
+import org.jetbrains.exposed.dao.id.EntityID
+import org.jetbrains.exposed.dao.id.IntIdTable
+import org.jetbrains.exposed.sql.*
+import java.util.concurrent.CompletableFuture
+import kotlin.sequences.Sequence
+
+private const val MAX_FQNAME_LENGTH = 255
+private const val MAX_SHORT_NAME_LENGTH = 80
+private const val MAX_URI_LENGTH = 511
+
+private object Symbols : IntIdTable() {
+    val fqName = varchar("fqname", length = MAX_FQNAME_LENGTH).index()
+    val shortName = varchar("shortname", length = MAX_SHORT_NAME_LENGTH)
     val kind = integer("kind")
     val visibility = integer("visibility")
     val extensionReceiverType = varchar("extensionreceivertype", length = MAX_FQNAME_LENGTH).nullable()
+    val location = optReference("location", Locations)
+}
 
-    override val primaryKey = PrimaryKey(fqName)
+private object Locations : IntIdTable() {
+    val uri = varchar("uri", length = MAX_URI_LENGTH)
+    val range = reference("range", Ranges)
 }
 
-private object FqNames : Table() {
-    val fqName = varchar("fqname", length = MAX_FQNAME_LENGTH)
-    val shortName = varchar("shortname", length = MAX_SHORT_NAME_LENGTH)
+private object Ranges : IntIdTable() {
+    val start = reference("start", Positions)
+    val end = reference("end", Positions)
+}
+
+private object Positions : IntIdTable() {
+    val line = integer("line")
+    val character = integer("character")
+}
+
+class SymbolEntity(id: EntityID<Int>) : IntEntity(id) {
+    companion object : IntEntityClass<SymbolEntity>(Symbols)
+
+    var fqName by Symbols.fqName
+    var shortName by Symbols.shortName
+    var kind by Symbols.kind
+    var visibility by Symbols.visibility
+    var extensionReceiverType by Symbols.extensionReceiverType
+    var location by LocationEntity optionalReferencedOn Symbols.location
+}
+
+class LocationEntity(id: EntityID<Int>) : IntEntity(id) {
+    companion object : IntEntityClass<LocationEntity>(Locations)
 
-    override val primaryKey = PrimaryKey(fqName)
+    var uri by Locations.uri
+    var range by RangeEntity referencedOn Locations.range
+}
+
+class RangeEntity(id: EntityID<Int>) : IntEntity(id) {
+    companion object : IntEntityClass<RangeEntity>(Ranges)
+
+    var start by PositionEntity referencedOn Ranges.start
+    var end by PositionEntity referencedOn Ranges.end
+}
+
+class PositionEntity(id: EntityID<Int>) : IntEntity(id) {
+    companion object : IntEntityClass<PositionEntity>(Positions)
+
+    var line by Positions.line
+    var character by Positions.character
 }
 
 /**
@@ -48,47 +82,49 @@ private object FqNames : Table() {
 class SymbolIndex {
     private val db = Database.connect("jdbc:h2:mem:symbolindex;DB_CLOSE_DELAY=-1", "org.h2.Driver")
 
-    var progressFactory: Progress.Factory = Progress.Factory.None
+    var progressFactory: Progress.Factory = object: Progress.Factory {
+        override fun create(label: String): CompletableFuture<Progress> = CompletableFuture.supplyAsync { Progress.None }
+    }
 
     init {
         transaction(db) {
-            SchemaUtils.create(Symbols, FqNames)
+            SchemaUtils.create(Symbols, Locations, Ranges, Positions)
         }
     }
 
     /** Rebuilds the entire index. May take a while. */
-    fun refresh(module: ModuleDescriptor) {
+    fun refresh(module: ModuleDescriptor, exclusions: Sequence<DeclarationDescriptor>) {
+        val started = System.currentTimeMillis()
+        LOG.info("Updating full symbol index...")
+
+        progressFactory.create("Indexing").thenApplyAsync { progress ->
+            try {
+                transaction(db) {
+                    addDeclarations(allDescriptors(module, exclusions))
+
+                    val finished = System.currentTimeMillis()
+                    val count = Symbols.slice(Symbols.fqName.count()).selectAll().first()[Symbols.fqName.count()]
+                    LOG.info("Updated full symbol index in ${finished - started} ms! (${count} symbol(s))")
+                }
+            } catch (e: Exception) {
+                LOG.error("Error while updating symbol index")
+                LOG.printStackTrace(e)
+            }
+
+            progress.close()
+        }
+    }
+
+    // Removes a list of indexes and adds another list. Everything is done in the same transaction.
+    fun updateIndexes(remove: Sequence<DeclarationDescriptor>, add: Sequence<DeclarationDescriptor>) {
         val started = System.currentTimeMillis()
         LOG.info("Updating symbol index...")
 
-        progressFactory.create("Indexing").thenApply { progress ->
+        progressFactory.create("Indexing").thenApplyAsync { progress ->
             try {
-                // TODO: Incremental updates
                 transaction(db) {
-                    Symbols.deleteAll()
-
-                    for (descriptor in allDescriptors(module)) {
-                        val descriptorFqn = descriptor.fqNameSafe
-                        val extensionReceiverFqn = descriptor.accept(ExtractSymbolExtensionReceiverType, Unit)?.takeIf { !it.isRoot }
-
-                        if (canStoreFqName(descriptorFqn) && (extensionReceiverFqn?.let { canStoreFqName(it) } ?: true)) {
-                            for (fqn in listOf(descriptorFqn, extensionReceiverFqn).filterNotNull()) {
-                                FqNames.replace {
-                                    it[fqName] = fqn.toString()
-                                    it[shortName] = fqn.shortName().toString()
-                                }
-                            }
-
-                            Symbols.replace {
-                                it[fqName] = descriptorFqn.toString()
-                                it[kind] = descriptor.accept(ExtractSymbolKind, Unit).rawValue
-                                it[visibility] = descriptor.accept(ExtractSymbolVisibility, Unit).rawValue
-                                it[extensionReceiverType] = extensionReceiverFqn?.toString()
-                            }
-                        } else {
-                            LOG.warn("Excluding symbol {} from index since its name is too long", descriptorFqn.toString())
-                        }
-                    }
+                    removeDeclarations(remove)
+                    addDeclarations(add)
 
                     val finished = System.currentTimeMillis()
                     val count = Symbols.slice(Symbols.fqName.count()).selectAll().first()[Symbols.fqName.count()]
@@ -103,29 +139,68 @@ class SymbolIndex {
         }
     }
 
-    private fun canStoreFqName(fqName: FqName) =
-           fqName.toString().length <= MAX_FQNAME_LENGTH
-        && fqName.shortName().toString().length <= MAX_SHORT_NAME_LENGTH
+    private fun Transaction.removeDeclarations(declarations: Sequence<DeclarationDescriptor>) =
+        declarations.forEach { declaration ->
+            val (descriptorFqn, extensionReceiverFqn) = getFqNames(declaration)
+
+            if (validFqName(descriptorFqn) && (extensionReceiverFqn?.let { validFqName(it) } != false)) {
+                Symbols.deleteWhere {
+                    (Symbols.fqName eq descriptorFqn.toString()) and (Symbols.extensionReceiverType eq extensionReceiverFqn?.toString())
+                }
+            } else {
+                LOG.warn("Excluding symbol {} from index since its name is too long", descriptorFqn.toString())
+            }
+        }
+
+    private fun Transaction.addDeclarations(declarations: Sequence<DeclarationDescriptor>) =
+        declarations.forEach { declaration ->
+            val (descriptorFqn, extensionReceiverFqn) = getFqNames(declaration)
+
+            if (validFqName(descriptorFqn) && (extensionReceiverFqn?.let { validFqName(it) } != false)) {
+                SymbolEntity.new {
+                    fqName = descriptorFqn.toString()
+                    shortName = descriptorFqn.shortName().toString()
+                    kind = declaration.accept(ExtractSymbolKind, Unit).rawValue
+                    visibility = declaration.accept(ExtractSymbolVisibility, Unit).rawValue
+                    extensionReceiverType = extensionReceiverFqn?.toString()
+                }
+            } else {
+                LOG.warn("Excluding symbol {} from index since its name is too long", descriptorFqn.toString())
+            }
+        }
+
+    private fun getFqNames(declaration: DeclarationDescriptor): Pair<FqName, FqName?> {
+        val descriptorFqn = declaration.fqNameSafe
+        val extensionReceiverFqn = declaration.accept(ExtractSymbolExtensionReceiverType, Unit)?.takeIf { !it.isRoot }
+
+        return Pair(descriptorFqn, extensionReceiverFqn)
+    }
+
+    private fun validFqName(fqName: FqName) =
+        fqName.toString().length <= MAX_FQNAME_LENGTH
+            && fqName.shortName().toString().length <= MAX_SHORT_NAME_LENGTH
 
     fun query(prefix: String, receiverType: FqName? = null, limit: Int = 20): List<Symbol> = transaction(db) {
         // TODO: Extension completion currently only works if the receiver matches exactly,
         //       ideally this should work with subtypes as well
-        (Symbols innerJoin FqNames)
-            .select { FqNames.shortName.like("$prefix%") and (Symbols.extensionReceiverType eq receiverType?.toString()) }
-            .limit(limit)
+        SymbolEntity.find {
+            (Symbols.shortName like "$prefix%") and (Symbols.extensionReceiverType eq receiverType?.toString())
+        }.limit(limit)
             .map { Symbol(
-                fqName = FqName(it[Symbols.fqName]),
-                kind = Symbol.Kind.fromRaw(it[Symbols.kind]),
-                visibility = Symbol.Visibility.fromRaw(it[Symbols.visibility]),
-                extensionReceiverType = it[Symbols.extensionReceiverType]?.let(::FqName)
+                fqName = FqName(it.fqName),
+                kind = Symbol.Kind.fromRaw(it.kind),
+                visibility = Symbol.Visibility.fromRaw(it.visibility),
+                extensionReceiverType = it.extensionReceiverType?.let(::FqName)
             ) }
     }
 
-    private fun allDescriptors(module: ModuleDescriptor): Sequence<DeclarationDescriptor> = allPackages(module)
+    private fun allDescriptors(module: ModuleDescriptor, exclusions: Sequence<DeclarationDescriptor>): Sequence<DeclarationDescriptor> = allPackages(module)
         .map(module::getPackage)
         .flatMap {
             try {
-                it.memberScope.getContributedDescriptors(DescriptorKindFilter.ALL, MemberScope.ALL_NAME_FILTER)
+                it.memberScope.getContributedDescriptors(
+                    DescriptorKindFilter.ALL
+                ) { name -> !exclusions.any { declaration -> declaration.name == name } }
             } catch (e: IllegalStateException) {
                 LOG.warn("Could not query descriptors in package $it")
                 emptyList()
