ADFA-6858: Kotlin declaration lookups without whole-package parses

lsp/jvm-symbol-index/src/main/kotlin/org/appdevforall/codeonthego/indexing/jvm/KtFileMetadataIndex.kt

@@ -8,6 +8,34 @@ import org.appdevforall.codeonthego.indexing.api.indexQuery
 import org.appdevforall.codeonthego.indexing.jvm.KtFileMetadataDescriptor.KEY_IS_INDEXED
 import org.appdevforall.codeonthego.indexing.jvm.KtFileMetadataDescriptor.KEY_PACKAGE
 import java.io.Closeable
+import java.util.concurrent.CopyOnWriteArrayList
+
+/**
+ * A completed write to a [KtFileMetadataIndex].
+ */
+sealed interface KtFileMetadataWrite {
+	/** The record for [filePath] was inserted or replaced; its package is [packageFqName]. */
+	data class Upserted(
+		val filePath: String,
+		val packageFqName: String,
+	) : KtFileMetadataWrite
+
+	/** The records for [filePaths] were removed. */
+	data class Removed(
+		val filePaths: Collection<String>,
+	) : KtFileMetadataWrite
+
+	/** Every record was removed. */
+	data object Cleared : KtFileMetadataWrite
+}
+
+/**
+ * Receives each [KtFileMetadataWrite] after it has been applied to the index.
+ */
+fun interface KtFileMetadataWriteListener {
+	/** Called synchronously on the writing thread once [write] is visible to readers. */
+	fun onWrite(write: KtFileMetadataWrite)
+}
 
 /**
  * An index of [KtFileMetadata] entries, one per Kotlin source file.
@@ -40,20 +68,42 @@ class KtFileMetadataIndex(
 			)
 	}
 
+	private val listeners = CopyOnWriteArrayList<KtFileMetadataWriteListener>()
+
+	/**
+	 * Registers [listener] to be told about every later write.
+	 *
+	 * Closing the returned handle unregisters it.
+	 */
+	fun addWriteListener(listener: KtFileMetadataWriteListener): AutoCloseable {
+		listeners.add(listener)
+		return AutoCloseable { listeners.remove(listener) }
+	}
+
+	private fun notifyListeners(write: KtFileMetadataWrite) {
+		listeners.forEach { it.onWrite(write) }
+	}
+
 	/**
 	 * Insert or replace the metadata record for a single file.
 	 *
 	 * Because [KtFileMetadata.key] and [KtFileMetadata.sourceId] are both the
 	 * file path, a file has at most one row and an insert replaces it.
 	 */
-	suspend fun upsert(metadata: KtFileMetadata) = backing.insert(metadata)
+	suspend fun upsert(metadata: KtFileMetadata) {
+		backing.insert(metadata)
+		notifyListeners(KtFileMetadataWrite.Upserted(metadata.filePath, metadata.packageFqName))
+	}
 
 	/**
 	 * Remove the metadata record for [filePath].
 	 *
 	 * No-op if the file is not in the index.
 	 */
-	suspend fun remove(filePath: String) = backing.removeBySource(filePath)
+	suspend fun remove(filePath: String) {
+		backing.removeBySource(filePath)
+		notifyListeners(KtFileMetadataWrite.Removed(listOf(filePath)))
+	}
 
 	/**
 	 * Remove the metadata records for every path in [filePaths] as a single
@@ -63,7 +113,10 @@ class KtFileMetadataIndex(
 	 * transaction instead of N - see [Index.removeBySources]. Paths not present in
 	 * the index are ignored.
 	 */
-	suspend fun removeAll(filePaths: Collection<String>) = backing.removeBySources(filePaths)
+	suspend fun removeAll(filePaths: Collection<String>) {
+		backing.removeBySources(filePaths)
+		notifyListeners(KtFileMetadataWrite.Removed(filePaths))
+	}
 
 	/**
 	 * Return the [KtFileMetadata] for [filePath], or `null` if the file is
@@ -154,7 +207,10 @@ class KtFileMetadataIndex(
 			).map { it.filePath }
 
 	/** Remove all records from the index. */
-	suspend fun clear() = backing.clear()
+	suspend fun clear() {
+		backing.clear()
+		notifyListeners(KtFileMetadataWrite.Cleared)
+	}
 
 	/** Whether the backing index still serves reads and writes. */
 	val isOpen: Boolean

lsp/jvm-symbol-index/src/test/kotlin/org/appdevforall/codeonthego/indexing/jvm/KtFileMetadataIndexWriteListenerTest.kt

@@ -0,0 +1,78 @@
+package org.appdevforall.codeonthego.indexing.jvm
+
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.appdevforall.codeonthego.indexing.InMemoryIndex
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import java.time.Instant
+
+@RunWith(JUnit4::class)
+class KtFileMetadataIndexWriteListenerTest {
+	private val index = KtFileMetadataIndex(InMemoryIndex(KtFileMetadataDescriptor))
+	private val writes = mutableListOf<KtFileMetadataWrite>()
+
+	private fun metadata(
+		path: String,
+		pkg: String,
+	) = KtFileMetadata(path, pkg, Instant.EPOCH, 1L)
+
+	@Test
+	fun upsertReportsPathAndPackage() =
+		runTest {
+			index.addWriteListener { writes += it }
+
+			index.upsert(metadata("/a/A.kt", "com.a"))
+
+			assertThat(writes).containsExactly(KtFileMetadataWrite.Upserted("/a/A.kt", "com.a"))
+		}
+
+	@Test
+	fun removeAndRemoveAllReportPaths() =
+		runTest {
+			index.addWriteListener { writes += it }
+
+			index.remove("/a/A.kt")
+			index.removeAll(listOf("/a/B.kt", "/a/C.kt"))
+
+			assertThat(writes)
+				.containsExactly(
+					KtFileMetadataWrite.Removed(listOf("/a/A.kt")),
+					KtFileMetadataWrite.Removed(listOf("/a/B.kt", "/a/C.kt")),
+				).inOrder()
+		}
+
+	@Test
+	fun clearReportsCleared() =
+		runTest {
+			index.addWriteListener { writes += it }
+
+			index.clear()
+
+			assertThat(writes).containsExactly(KtFileMetadataWrite.Cleared)
+		}
+
+	@Test
+	fun listenerSeesTheWriteAlreadyApplied() =
+		runTest {
+			var pathsSeenInListener: List<String>? = null
+			index.addWriteListener {
+				pathsSeenInListener = index.getFilePathsForPackage("com.a").toList()
+			}
+
+			index.upsert(metadata("/a/A.kt", "com.a"))
+
+			assertThat(pathsSeenInListener).containsExactly("/a/A.kt")
+		}
+
+	@Test
+	fun closedRegistrationIsNotCalled() =
+		runTest {
+			index.addWriteListener { writes += it }.close()
+
+			index.upsert(metadata("/a/A.kt", "com.a"))
+
+			assertThat(writes).isEmpty()
+		}
+}

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/KotlinLanguageServer.kt

@@ -475,12 +475,13 @@ class KotlinLanguageServer : ILanguageServer {
 	@Suppress("unused")
 	fun onFileDeleted(event: FileDeletionEvent) {
 		val path = event.file.toPath()
-		if (!DocumentUtils.isKotlinFile(path)) {
+		// The file is already gone, so only its extension can say it was Kotlin.
+		if (!DocumentUtils.hasKotlinExtension(path)) {
 			return
 		}
 
 		scope.launch {
-			runCatching { compiler?.compilationEnvironmentFor(path) }
+			runCatching { compiler?.compilationEnvironmentForKotlinPath(path) }
 				.getOrNull()
 				?.onFileRemoved(path)
 		}
@@ -495,7 +496,8 @@ class KotlinLanguageServer : ILanguageServer {
 		scope.launch {
 			// One read: a rebuild can swap the compiler between two reads of the property.
 			val compiler = compiler ?: return@launch
-			val oldIsKotlinFile = DocumentUtils.isKotlinFile(fromPath)
+			// The old path no longer exists, so only its extension can say it was Kotlin.
+			val oldIsKotlinFile = DocumentUtils.hasKotlinExtension(fromPath)
 			val newIsKotlinFile = DocumentUtils.isKotlinFile(toPath)
 
 			if (!oldIsKotlinFile && newIsKotlinFile) {
@@ -511,17 +513,16 @@ class KotlinLanguageServer : ILanguageServer {
 				// only the old file was a Kotlin file
 				// so just remove it from the index
 				compiler
-					.compilationEnvironmentFor(fromPath)
+					.compilationEnvironmentForKotlinPath(fromPath)
 					?.onFileRemoved(fromPath)
 				return@launch
 			}
 
-			val fromKind = runCatching { compiler.compilationKindFor(fromPath) }.getOrNull()
-			val toKind = runCatching { compiler.compilationKindFor(toPath) }.getOrNull()
-			val fromEnv = fromKind?.let { compiler.compilationEnvironmentFor(it) }
-			val toEnv = toKind?.let { compiler.compilationEnvironmentFor(it) }
+			// Path-based, so a .kts side yields no environment instead of the kind-based overload's throw.
+			val fromEnv = runCatching { compiler.compilationEnvironmentForKotlinPath(fromPath) }.getOrNull()
+			val toEnv = runCatching { compiler.compilationEnvironmentForKotlinPath(toPath) }.getOrNull()
 
-			if (fromKind != null && fromEnv == toEnv && toEnv != null) {
+			if (fromEnv != null && fromEnv == toEnv) {
 				// file was renamed within the same compilation environment
 				toEnv.onFileMoved(fromPath, toPath)
 				return@launch

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/compiler/AbstractCompilationEnvironment.kt

@@ -1,6 +1,7 @@
 package com.itsaky.androidide.lsp.kotlin.compiler
 
 import com.itsaky.androidide.lsp.kotlin.compiler.index.KtSymbolIndex
+import com.itsaky.androidide.lsp.kotlin.compiler.modules.AbstractKtModule
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.AnalysisPriority
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.AnalysisScheduler
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.KtModule
@@ -9,18 +10,28 @@ import com.itsaky.androidide.lsp.kotlin.compiler.modules.isSourceModule
 import com.itsaky.androidide.lsp.kotlin.compiler.services.JavaModuleAccessibilityChecker
 import com.itsaky.androidide.lsp.kotlin.compiler.services.JavaModuleAnnotationsProvider
 import com.itsaky.androidide.lsp.kotlin.compiler.services.KtLspService
+import com.itsaky.androidide.lsp.kotlin.compiler.services.ProjectStructureProvider
+import com.itsaky.androidide.lsp.kotlin.compiler.services.ResolutionScopeProvider
 import com.itsaky.androidide.lsp.kotlin.compiler.services.WriteAccessGuard
 import com.itsaky.androidide.lsp.kotlin.compiler.services.latestLanguageVersionSettings
 import com.itsaky.androidide.lsp.kotlin.compiler.util.SLF4JLogger
+import com.itsaky.androidide.lsp.kotlin.utils.toVirtualFileOrNull
 import com.itsaky.androidide.progress.ICancelChecker
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withTimeoutOrNull
 import org.jetbrains.kotlin.K1Deprecation
+import org.jetbrains.kotlin.analysis.api.KaImplementationDetail
+import org.jetbrains.kotlin.analysis.api.platform.analysisMessageBus
 import org.jetbrains.kotlin.analysis.api.platform.declarations.KotlinAnnotationsResolverFactory
 import org.jetbrains.kotlin.analysis.api.platform.declarations.KotlinDeclarationProviderFactory
 import org.jetbrains.kotlin.analysis.api.platform.declarations.KotlinDirectInheritorsProvider
 import org.jetbrains.kotlin.analysis.api.platform.java.KotlinJavaModuleAccessibilityChecker
 import org.jetbrains.kotlin.analysis.api.platform.java.KotlinJavaModuleAnnotationsProvider
+import org.jetbrains.kotlin.analysis.api.platform.modification.KaElementModificationType
+import org.jetbrains.kotlin.analysis.api.platform.modification.KaSourceModificationService
+import org.jetbrains.kotlin.analysis.api.platform.modification.KotlinModuleStateModificationEvent
+import org.jetbrains.kotlin.analysis.api.platform.modification.KotlinModuleStateModificationKind
+import org.jetbrains.kotlin.analysis.api.platform.modification.publishModificationEvent
 import org.jetbrains.kotlin.analysis.api.platform.packages.KotlinPackagePartProviderFactory
 import org.jetbrains.kotlin.analysis.api.platform.packages.KotlinPackageProviderFactory
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinModuleDependentsProvider
@@ -31,6 +42,7 @@ import org.jetbrains.kotlin.analysis.api.standalone.base.projectStructure.Standa
 import org.jetbrains.kotlin.analysis.api.standalone.base.projectStructure.registerProjectExtensionPoints
 import org.jetbrains.kotlin.analysis.api.standalone.base.projectStructure.registerProjectModelServices
 import org.jetbrains.kotlin.analysis.api.standalone.base.projectStructure.registerProjectServices
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidationTopics
 import org.jetbrains.kotlin.cli.common.intellijPluginRoot
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSourceLocation
@@ -80,6 +92,7 @@ import org.jetbrains.kotlin.config.useFir
 import org.jetbrains.kotlin.load.kotlin.MetadataFinderFactory
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.slf4j.LoggerFactory
 import java.nio.file.Path
@@ -339,6 +352,72 @@ internal abstract class AbstractCompilationEnvironment(
 		postInit(libraryRoots)
 	}
 
+	suspend fun onFileCreated(path: Path) {
+		if (project.isDisposed) return
+		notifyElementModifiedForPath(path) { KaElementModificationType.ElementAdded }
+		ktSymbolIndex.submitForIndexing(path)
+	}
+
+	suspend fun onFileRemoved(path: Path) {
+		if (project.isDisposed) return
+		notifyElementModifiedForPath(path) { ktFile ->
+			KaElementModificationType.ElementRemoved(ktFile)
+		}
+		ProjectStructureProvider.getInstance(project).unregisterInMemoryFile(path.pathString)
+		ktSymbolIndex.removeFromIndex(path)
+	}
+
+	@OptIn(KaImplementationDetail::class)
+	protected inline fun notifyElementModifiedForPath(
+		path: Path,
+		crossinline typeProvider: (KtFile) -> KaElementModificationType,
+	) {
+		// Resolve PSI/module structure under the read lock; driving psiManager.findFile /
+		// structureProvider concurrently with an `analyze` read section otherwise races.
+		val (ktFile, module) =
+			project.readUnlessDisposed {
+				val structureProvider = ProjectStructureProvider.getInstance(project)
+				val ktFile =
+					path.toVirtualFileOrNull()?.let {
+						psiManager.findFile(it) as? KtFile
+					}
+
+				val module =
+					(
+						ktFile?.let { structureProvider.getModule(it, null) }
+							?: structureProvider.findModuleForSourceId(path.pathString)
+					) as? AbstractKtModule
+
+				ktFile to module
+			} ?: return
+
+		project.writeUnlessDisposed {
+			// Must run under the write lock so the session mutation can't race a concurrent
+			// `analyze` (which only holds the read lock); see KtSymbolIndex.refreshToCurrent.
+			if (ktFile != null) {
+				KaSourceModificationService
+					.getInstance(project)
+					.handleElementModification(ktFile, typeProvider(ktFile))
+			}
+
+			// A path in no module is in no session's scope; invalidating globally would only force
+			// the platform and library sessions to rebuild.
+			if (module != null) {
+				module.invalidateSearchScope()
+				project.publishModificationEvent(
+					KotlinModuleStateModificationEvent(
+						module,
+						KotlinModuleStateModificationKind.UPDATE,
+					),
+				)
+				project.analysisMessageBus
+					.syncPublisher(LLFirSessionInvalidationTopics.SESSION_INVALIDATION)
+					.afterInvalidation(setOf(module))
+				ResolutionScopeProvider.getInstance(project).invalidate(module)
+			}
+		}
+	}
+
 	private fun createCompilerConfiguration(): CompilerConfiguration =
 		CompilerConfiguration().apply {
 			this.moduleName = JvmProtoBufUtil.DEFAULT_MODULE_NAME