perf: add declarations in bulk

Author: amgdev9

app/src/main/kotlin/org/kotlinlsp/index/db/Database.kt

@@ -9,7 +9,7 @@ import org.kotlinlsp.index.db.adapters.put
 import java.io.File
 import kotlin.io.path.absolutePathString
 
-const val CURRENT_SCHEMA_VERSION = 4    // Increment on schema changes
+const val CURRENT_SCHEMA_VERSION = 5    // Increment on schema changes
 const val VERSION_KEY = "__version"
 
 class Database(rootFolder: String) {

app/src/main/kotlin/org/kotlinlsp/index/db/Declaration.kt

@@ -4,6 +4,7 @@ import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import org.kotlinlsp.common.info
 import org.kotlinlsp.index.db.adapters.put
+import org.kotlinlsp.index.db.adapters.putBulk
 
 @Serializable
 sealed class Declaration() {
@@ -76,7 +77,8 @@ sealed class Declaration() {
     ) : Declaration()
 }
 
-fun Database.putDeclaration(decl: Declaration) {
-    val key = "${decl.name}:${decl.file}:${decl.startOffset}:${decl.endOffset}"
-    declarationsDb.put(key, decl)
+fun Declaration.id() = "${name}:${file}:${startOffset}:${endOffset}"
+
+fun Database.putDeclarations(declarations: Iterable<Declaration>) {
+    declarationsDb.putBulk(declarations.map { Pair(it.id(), it) })
 }

app/src/main/kotlin/org/kotlinlsp/index/db/File.kt

@@ -14,6 +14,7 @@ data class File(
     val lastModified: Instant,
     val modificationStamp: Long,
     val indexed: Boolean,
+    val declarationKeys: MutableList<String> = mutableListOf()
 ) {
     companion object {
         fun fromKtFile(ktFile: KtFile, project: Project, indexed: Boolean): File = project.read {
@@ -43,6 +44,7 @@ fun File.toDto(): FileDto = FileDto(
     lastModified = lastModified.toEpochMilli(),
     modificationStamp = modificationStamp,
     indexed = indexed,
+    declarationKeys = declarationKeys
 )
 
 @Serializable
@@ -51,6 +53,7 @@ data class FileDto(
     val lastModified: Long,
     val modificationStamp: Long,
     val indexed: Boolean,
+    val declarationKeys: List<String>
 )
 
 fun Database.file(path: String): File? {

app/src/main/kotlin/org/kotlinlsp/index/db/adapters/DatabaseAdapter.kt

@@ -7,6 +7,7 @@ import kotlinx.serialization.protobuf.ProtoBuf
 
 interface DatabaseAdapter {
     fun putByteArray(key: String, value: ByteArray)
+    fun putByteArrayBulk(values: Iterable<Pair<String, ByteArray>>)
     fun getByteArray(key: String): ByteArray?
     fun prefixSearch(prefix: String): Sequence<Pair<String, ByteArray>>
     fun remove(key: String)
@@ -19,6 +20,11 @@ inline fun <reified T: Any> DatabaseAdapter.put(key: String, value: T) {
     putByteArray(key, ProtoBuf.encodeToByteArray(value))
 }
 
+@OptIn(ExperimentalSerializationApi::class)
+inline fun <reified T: Any> DatabaseAdapter.putBulk(values: Iterable<Pair<String, T>>) {
+    putByteArrayBulk(values.map { Pair(it.first, ProtoBuf.encodeToByteArray(it.second)) })
+}
+
 @OptIn(ExperimentalSerializationApi::class)
 inline fun <reified T> DatabaseAdapter.get(key: String): T? {
     val data = getByteArray(key) ?: return null

app/src/main/kotlin/org/kotlinlsp/index/db/adapters/RocksDBAdapter.kt

@@ -1,17 +1,10 @@
 package org.kotlinlsp.index.db.adapters
 
-import kotlinx.serialization.*
-import kotlinx.serialization.protobuf.ProtoBuf
-import org.rocksdb.InfoLogLevel
-import org.rocksdb.Options
-import org.rocksdb.ReadOptions
-import org.rocksdb.RocksDB
+import org.rocksdb.*
 import java.io.File
-import java.io.Serial
 import java.nio.charset.Charset
 import java.nio.file.Path
 import kotlin.io.path.absolutePathString
-import kotlin.reflect.KClass
 
 class RocksDBAdapter(private val path: Path): DatabaseAdapter {
     companion object {
@@ -32,6 +25,15 @@ class RocksDBAdapter(private val path: Path): DatabaseAdapter {
         db.put(key.toByteArray(), value)
     }
 
+    override fun putByteArrayBulk(values: Iterable<Pair<String, ByteArray>>) {
+        val batch = WriteBatch()
+        values.forEach { (key, value) ->
+            batch.put(key.toByteArray(), value)
+        }
+        db.write(WriteOptions(), batch)
+        batch.close()
+    }
+
     override fun getByteArray(key: String): ByteArray? {
         val data = db.get(key.toByteArray()) ?: return null
         return data

app/src/main/kotlin/org/kotlinlsp/index/worker/IndexKtFile.kt

@@ -13,34 +13,38 @@ import org.kotlinlsp.index.db.*
 fun indexKtFile(project: Project, ktFile: KtFile, db: Database) {
     val newFile = File.fromKtFile(ktFile, project, indexed = true)
 
+    // Check if skipping the indexing of that file
     val existingFile = db.file(newFile.path)
     if (
         File.shouldBeSkipped(existingFile = existingFile, newFile = newFile) &&
         existingFile?.indexed == true  // Already indexed
     ) return
 
-    // TODO Remove declarations for this file first
+    // Remove declarations for this file first
+    existingFile?.declarationKeys?.forEach {
+        db.declarationsDb.remove(it)
+    }
+
+    // Get declarations metadata
     val declarations = project.read {
-        val list = mutableListOf<KtDeclaration>()
+        val list = mutableListOf<Declaration>()
         ktFile.accept(object : KtTreeVisitorVoid() {
             override fun visitDeclaration(dcl: KtDeclaration) {
-                list.add(dcl)
+                val record = analyze(dcl) {
+                    analyzeDeclaration(newFile.path, dcl)
+                } ?: return
+                list.add(record)
                 super.visitDeclaration(dcl)
             }
         })
         return@read list
     }
 
-    declarations.forEach {
-        val record = analyze(it) {
-            analyzeDeclaration(newFile.path, it)
-        } ?: return@forEach
-
-        // TODO Put all declarations in bulk to improve performance
-        db.putDeclaration(record)
-    }
+    // Save declarations
+    db.putDeclarations(declarations)
 
-    // Update the file timestamp and package
+    // Update the file timestamp, package and declaration names
+    newFile.declarationKeys.addAll(declarations.map { it.id() })
     db.setFile(newFile)
 }