ADFA-4174: prioritize code completion requests over over analyze requests

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/actions/AddImportAction.kt

@@ -18,7 +18,6 @@ import com.itsaky.androidide.lsp.models.DocumentChange
 import com.itsaky.androidide.lsp.models.TextEdit
 import com.itsaky.androidide.resources.R
 import org.appdevforall.codeonthego.indexing.jvm.JvmSymbol
-import org.jetbrains.kotlin.analysis.api.fir.diagnostics.KaFirDiagnostic
 import org.slf4j.LoggerFactory
 
 class AddImportAction : BaseKotlinCodeAction() {
@@ -46,14 +45,13 @@ class AddImportAction : BaseKotlinCodeAction() {
 			return
 		}
 
-		val diagnostic = extra.diagnostic as? KaFirDiagnostic.UnresolvedReference?
-		if (diagnostic == null) {
+		val reference = extra.unresolvedReference
+		if (reference == null) {
 			markInvisible()
 			return
 		}
 
 		val env = extra.compilationEnv
-		val reference = diagnostic.reference
 		val hasImportableSymbols = env.ktSymbolIndex
 			.findSymbolBySimpleName(reference, limit = 0)
 			.any { it.kind.isClassifier }
@@ -65,10 +63,10 @@ class AddImportAction : BaseKotlinCodeAction() {
 	}
 
 	override suspend fun execAction(data: ActionData): Map<JvmSymbol, List<TextEdit>> {
-		val (diagnostic, env) = data.require<DiagnosticItem>().extra as? KotlinDiagnosticExtra
+		val (reference, env) = data.require<DiagnosticItem>().extra as? KotlinDiagnosticExtra
 			?: return emptyMap()
 
-		diagnostic as KaFirDiagnostic.UnresolvedReference
+		if (reference == null) return emptyMap()
 
 		val file = data.requireFile()
 		val nioPath = file.toPath()
@@ -77,7 +75,7 @@ class AddImportAction : BaseKotlinCodeAction() {
 			?: return emptyMap()
 
 		return env.ktSymbolIndex
-			.findSymbolBySimpleName(diagnostic.reference, limit = 0)
+			.findSymbolBySimpleName(reference, limit = 0)
 			.filter { it.kind.isClassifier }
 			.associateWith { symbol -> insertImport(ktFile, symbol.fqName) }
 	}

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

@@ -3,6 +3,7 @@ package com.itsaky.androidide.lsp.kotlin.compiler
 import com.itsaky.androidide.lsp.api.ILanguageClient
 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.AnalysisPreemptedException
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.KtModule
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.asFlatSequence
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.backingFilePath
@@ -166,9 +167,16 @@ internal class CompilationEnvironment(
 			scope = coroutineScope,
 			debounceDuration = DEFAULT_FILE_MOD_EVENT_DEBOUNCE_DURATION,
 		) { path, cancelChecker ->
-			val result = collectDiagnosticsFor(path, cancelChecker)
-			withContext(Dispatchers.Main.immediate) {
-				languageClient?.publishDiagnostics(result)
+			try {
+				val result = collectDiagnosticsFor(path, cancelChecker)
+				withContext(Dispatchers.Main.immediate) {
+					languageClient?.publishDiagnostics(result)
+				}
+			} catch (e: AnalysisPreemptedException) {
+				// A higher-priority analysis (completion) preempted this diagnostics run.
+				// Re-schedule so diagnostics still run once the higher-priority work finishes.
+				logger.debug("diagnostics for {} preempted; rescheduling", path)
+				fileAnalyzer.schedule(path)
 			}
 		}
 	}
@@ -208,22 +216,24 @@ internal class CompilationEnvironment(
 	@OptIn(KaImplementationDetail::class)
 	private inline fun notifyElementModifiedForPath(
 		path: Path,
-		typeProvider: (KtFile) -> KaElementModificationType,
+		crossinline typeProvider: (KtFile) -> KaElementModificationType,
 	) {
 		val structureProvider = ProjectStructureProvider.getInstance(project)
 		val ktFile = path.toVirtualFileOrNull()?.let {
 			psiManager.findFile(it) as? KtFile
 		}
 
-		if (ktFile != null) {
-			KaSourceModificationService.getInstance(project)
-				.handleElementModification(ktFile, typeProvider(ktFile))
-		}
-
 		val module = (ktFile?.let { structureProvider.getModule(it, null) }
 			?: structureProvider.findModuleForSourceId(path.pathString)) as? AbstractKtModule
 
 		project.write {
+			// Must run under the write lock so the session mutation can't race a concurrent
+			// `analyze` (which only holds the read lock); see onFileContentChanged.
+			if (ktFile != null) {
+				KaSourceModificationService.getInstance(project)
+					.handleElementModification(ktFile, typeProvider(ktFile))
+			}
+
 			if (module != null) {
 				module.invalidateSearchScope()
 				project.publishModificationEvent(

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

@@ -1,13 +1,15 @@
 package com.itsaky.androidide.lsp.kotlin.compiler.index
 
 import com.itsaky.androidide.lsp.kotlin.compiler.CompilationEnvironment
+import com.itsaky.androidide.lsp.kotlin.compiler.modules.AnalysisPreemptedException
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.backingFilePath
 import com.itsaky.androidide.lsp.kotlin.compiler.read
 import com.itsaky.androidide.progress.ICancelChecker
 import com.itsaky.androidide.utils.KeyedDebouncingAction
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.isActive
+import kotlinx.coroutines.launch
 import org.appdevforall.codeonthego.indexing.jvm.JvmSymbolIndex
 import org.appdevforall.codeonthego.indexing.jvm.KtFileMetadata
 import org.appdevforall.codeonthego.indexing.jvm.KtFileMetadataIndex
@@ -54,8 +56,14 @@ internal class IndexWorker(
 			debounceDuration = CompilationEnvironment.DEFAULT_FILE_MOD_EVENT_DEBOUNCE_DURATION
 		) { (path, ktFile), cancelChecker ->
 			logger.debug("Indexing modified file: {}", path)
-			indexSourceFile(project, ktFile, fileIndex, sourceIndex, cancelChecker)
-			sourceIndexCount++
+			try {
+				indexSourceFile(project, ktFile, fileIndex, sourceIndex, cancelChecker)
+				sourceIndexCount++
+			} catch (e: AnalysisPreemptedException) {
+				// Preempted by higher-priority analysis; re-queue so the edit still gets indexed.
+				logger.debug("Indexing of modified file {} preempted; re-queueing", path)
+				scope.launch { submitCommand(IndexCommand.IndexModifiedFile(ktFile)) }
+			}
 		}
 
 		while (isActive) {
@@ -82,15 +90,23 @@ internal class IndexWorker(
 						continue
 					}
 
-					indexSourceFile(
-						project = project,
-						ktFile = ktFile,
-						fileIndex = fileIndex,
-						symbolsIndex = sourceIndex,
-						cancelChecker = ICancelChecker.NOOP
-					)
+					try {
+						indexSourceFile(
+							project = project,
+							ktFile = ktFile,
+							fileIndex = fileIndex,
+							symbolsIndex = sourceIndex,
+							// A real (cancellable) checker so the scheduler can preempt this pass
+							// in favour of completion/diagnostics.
+							cancelChecker = ICancelChecker.Default()
+						)
 
-					sourceIndexCount++
+						sourceIndexCount++
+					} catch (e: AnalysisPreemptedException) {
+						// Preempted by higher-priority analysis; re-queue so the file still gets indexed.
+						logger.debug("Indexing of {} preempted; re-queueing", cmd.vf.path)
+						scope.launch { submitCommand(cmd) }
+					}
 				}
 
 				is IndexCommand.IndexModifiedFile -> {

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

@@ -1,5 +1,7 @@
 package com.itsaky.androidide.lsp.kotlin.compiler.index
 
+import com.itsaky.androidide.lsp.kotlin.compiler.modules.AnalysisPriority
+import com.itsaky.androidide.lsp.kotlin.compiler.modules.ScheduledCancelChecker
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.analyzeMaybeDangling
 import com.itsaky.androidide.lsp.kotlin.compiler.modules.backingFilePath
 import com.itsaky.androidide.lsp.kotlin.compiler.read
@@ -74,9 +76,15 @@ internal suspend fun indexSourceFile(
 	symbolsIndex: JvmSymbolIndex,
 	cancelChecker: ICancelChecker,
 ) {
+	// Indexing runs at the lowest (INDEXING) priority: it yields to both completion and diagnostics.
+	// Wrapping the checker lets the scheduler preempt an in-progress index pass; the preemption
+	// surfaces as AnalysisPreemptedException at the abortIfCancelled() checkpoints below, which the
+	// IndexWorker catches to re-queue the file.
+	val checker = cancelChecker as? ScheduledCancelChecker ?: ScheduledCancelChecker(cancelChecker)
+
 	val newFile = ktFile.toMetadata(project, isIndexed = true)
 	val existingFile = fileIndex.get(newFile.filePath)
-	cancelChecker.abortIfCancelled()
+	checker.abortIfCancelled()
 
 	if (KtFileMetadata.shouldBeSkipped(existingFile, newFile) && existingFile?.isIndexed == true) {
 		return
@@ -85,18 +93,18 @@ internal suspend fun indexSourceFile(
 	// Remove stale symbols written during the previous indexing pass.
 	if (existingFile?.isIndexed == true) {
 		symbolsIndex.removeBySource(newFile.filePath)
-		cancelChecker.abortIfCancelled()
+		checker.abortIfCancelled()
 	}
 
 	val symbols = project.read {
 		val list = mutableListOf<JvmSymbol>()
-		analyzeMaybeDangling(ktFile) {
+		analyzeMaybeDangling(ktFile, AnalysisPriority.INDEXING, checker) {
 			val session = this
 			ktFile.accept(object : KtTreeVisitorVoid() {
 				override fun visitDeclaration(dcl: KtDeclaration) {
-					cancelChecker.abortIfCancelled()
+					checker.abortIfCancelled()
 					val symbol = with(session) { analyzeDeclaration(newFile.filePath, dcl) }
-					cancelChecker.abortIfCancelled()
+					checker.abortIfCancelled()
 					symbol?.let { list.add(it) }
 					super.visitDeclaration(dcl)
 				}

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

@@ -0,0 +1,159 @@
+package com.itsaky.androidide.lsp.kotlin.compiler.modules
+
+import com.itsaky.androidide.progress.ICancelChecker
+import java.util.concurrent.CancellationException
+import java.util.concurrent.locks.ReentrantLock
+import kotlin.concurrent.withLock
+
+/**
+ * Priority of an Analysis API request. Higher [ordinal] wins: a request can preempt any strictly
+ * lower-priority analysis that is currently running, and is served before any lower-priority request
+ * that is merely waiting.
+ *
+ * Order: [INDEXING] < [DIAGNOSTICS] < [COMPLETION] — interactive completion beats background
+ * diagnostics, which beats bulk indexing.
+ */
+internal enum class AnalysisPriority {
+	INDEXING,
+	DIAGNOSTICS,
+	COMPLETION,
+}
+
+/**
+ * Thrown at an `abortIfCancelled()` checkpoint when the running analysis has been preempted by a
+ * higher-priority request (see [AnalysisScheduler]). It is a [CancellationException] so it unwinds
+ * cleanly through the existing cancellation-aware `catch` blocks; callers that want the preempted work
+ * to run later catch this specific type and re-schedule it.
+ */
+internal class AnalysisPreemptedException :
+	CancellationException("analysis preempted by a higher-priority request")
+
+/**
+ * An [ICancelChecker] that adds a cooperative *preemption* signal on top of an existing [delegate]
+ * checker. The Analysis API cannot be interrupted mid-`analyze` (it runs with a no-op
+ * `ProgressManager`), so [AnalysisScheduler] flags preemption here and the running analysis notices it
+ * at its next [abortIfCancelled] checkpoint.
+ *
+ * Preemption is distinct from ordinary cancellation: [abortIfCancelled] throws
+ * [AnalysisPreemptedException] (so the source can re-schedule the work) while still honouring the
+ * delegate's own cancellation (e.g. a superseding edit or a closed file).
+ */
+internal class ScheduledCancelChecker(
+	private val delegate: ICancelChecker,
+) : ICancelChecker {
+
+	@Volatile
+	private var preempted = false
+
+	/** Marks this analysis as preempted; the next [abortIfCancelled] will throw. */
+	fun preempt() {
+		preempted = true
+	}
+
+	override fun cancel() {
+		delegate.cancel()
+	}
+
+	override fun isCancelled(): Boolean = preempted || delegate.isCancelled()
+
+	override fun abortIfCancelled() {
+		if (preempted) {
+			throw AnalysisPreemptedException()
+		}
+		delegate.abortIfCancelled()
+	}
+}
+
+/**
+ * A process-global, priority-aware, preemptive lock that serializes all Kotlin Analysis API access.
+ *
+ * It replaces the plain FIFO lock that previously guarded `analyze` / `analyzeCopy`. Semantics:
+ * - only one analysis runs at a time (the Analysis API is not safe to drive concurrently);
+ * - a higher-priority requester **preempts** a strictly lower-priority holder by invoking its
+ *   `onPreempt` callback once (cooperative — the holder bails at its next `abortIfCancelled()`);
+ * - when the lock frees, the highest-priority waiter acquires it next;
+ * - it is **reentrant**: a nested analysis on the same thread re-enters without deadlocking.
+ *
+ * Access it through [withAnalysisLock] / [analyzeMaybeDangling] rather than directly.
+ */
+internal object AnalysisScheduler {
+
+	private val mutex = ReentrantLock()
+	private val available = mutex.newCondition()
+
+	private var holderThread: Thread? = null
+	private var holderPriority: AnalysisPriority? = null
+	private var holderReentry = 0
+	private var holderPreempted = false
+	private var holderPreempt: (() -> Unit)? = null
+
+	/** Number of threads currently waiting to acquire, per priority. */
+	private val waiting = IntArray(AnalysisPriority.entries.size)
+
+	/**
+	 * Acquire the analysis lock at the given [priority]. Blocks until the current thread may run. If a
+	 * strictly lower-priority analysis is in progress, [onPreempt] of *that* holder is invoked so it
+	 * yields; [onPreempt] passed here is stored and used if this acquisition is later preempted.
+	 */
+	fun acquire(priority: AnalysisPriority, onPreempt: () -> Unit) {
+		mutex.withLock {
+			val me = Thread.currentThread()
+			if (holderThread === me) {
+				// Reentrant: nested analysis on the same thread shares the outer hold.
+				holderReentry++
+				return
+			}
+
+			waiting[priority.ordinal]++
+			try {
+				while (true) {
+					val hp = holderPriority
+					if (holderThread != null && hp != null &&
+						hp.ordinal < priority.ordinal && !holderPreempted
+					) {
+						// Signal the lower-priority holder to bail (once).
+						holderPreempted = true
+						holderPreempt?.invoke()
+					}
+
+					if (holderThread == null && !higherPriorityWaiting(priority)) {
+						break
+					}
+					available.await()
+				}
+			} finally {
+				waiting[priority.ordinal]--
+			}
+
+			holderThread = me
+			holderPriority = priority
+			holderPreempt = onPreempt
+			holderPreempted = false
+			holderReentry = 1
+		}
+	}
+
+	/** Release a hold acquired via [acquire]. Wakes waiters when the outermost hold is released. */
+	fun release() {
+		mutex.withLock {
+			if (holderThread !== Thread.currentThread()) {
+				return
+			}
+			if (--holderReentry > 0) {
+				return
+			}
+			holderThread = null
+			holderPriority = null
+			holderPreempt = null
+			holderPreempted = false
+			available.signalAll()
+		}
+	}
+
+	private fun higherPriorityWaiting(priority: AnalysisPriority): Boolean {
+		for (i in priority.ordinal + 1 until waiting.size) {
+			if (waiting[i] > 0) return true
+		}
+		return false
+	}
+}