ADFA-3320: introduce KtFileManager to manage parsed KtFile instances

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

@@ -0,0 +1,12 @@
+package com.itsaky.androidide.lsp.kotlin
+
+import java.nio.file.Path
+
+interface FileEventConsumer {
+
+	fun onFileOpened(path: Path, content: String)
+	fun onFileClosed(path: Path)
+
+	fun onFileContentChanged(path: Path, content: String)
+	fun onFileSaved(path: Path)
+}

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

@@ -22,13 +22,14 @@ import com.itsaky.androidide.app.configuration.IJdkDistributionProvider
 import com.itsaky.androidide.eventbus.events.editor.DocumentChangeEvent
 import com.itsaky.androidide.eventbus.events.editor.DocumentCloseEvent
 import com.itsaky.androidide.eventbus.events.editor.DocumentOpenEvent
+import com.itsaky.androidide.eventbus.events.editor.DocumentSaveEvent
 import com.itsaky.androidide.eventbus.events.editor.DocumentSelectedEvent
 import com.itsaky.androidide.lsp.api.ILanguageClient
 import com.itsaky.androidide.lsp.api.ILanguageServer
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.kotlin.compiler.Compiler
 import com.itsaky.androidide.lsp.kotlin.compiler.KotlinProjectModel
-import com.itsaky.androidide.lsp.kotlin.diagnostic.KotlinDiagnosticProvider
+import com.itsaky.androidide.lsp.kotlin.diagnostic.collectDiagnosticsFor
 import com.itsaky.androidide.lsp.models.CompletionParams
 import com.itsaky.androidide.lsp.models.CompletionResult
 import com.itsaky.androidide.lsp.models.DefinitionParams
@@ -76,7 +77,6 @@ class KotlinLanguageServer : ILanguageServer {
 		CoroutineScope(SupervisorJob() + CoroutineName(KotlinLanguageServer::class.simpleName!!))
 	private var projectModel: KotlinProjectModel? = null
 	private var compiler: Compiler? = null
-	private var diagnosticProvider: KotlinDiagnosticProvider? = null
 	private var analyzeJob: Job? = null
 
 	override val serverId: String = SERVER_ID
@@ -106,7 +106,6 @@ class KotlinLanguageServer : ILanguageServer {
 	override fun shutdown() {
 		EventBus.getDefault().unregister(this)
 		scope.cancel("LSP is being shut down")
-		diagnosticProvider?.close()
 		compiler?.close()
 		initialized = false
 	}
@@ -150,7 +149,6 @@ class KotlinLanguageServer : ILanguageServer {
 			)
 
 			this.compiler = compiler
-			this.diagnosticProvider = KotlinDiagnosticProvider(compiler, scope)
 		} else {
 			logger.info("Updating project model")
 
@@ -221,7 +219,7 @@ class KotlinLanguageServer : ILanguageServer {
 			return DiagnosticResult.NO_UPDATE
 		}
 
-		return diagnosticProvider?.analyze(file)
+		return compiler?.compilationEnvironmentFor(file)?.collectDiagnosticsFor(file)
 			?: DiagnosticResult.NO_UPDATE
 	}
 
@@ -232,6 +230,11 @@ class KotlinLanguageServer : ILanguageServer {
 			return
 		}
 
+		compiler?.compilationEnvironmentFor(event.openedFile)?.apply {
+			val content = FileManager.getDocumentContents(event.openedFile)
+			fileManager.onFileOpened(event.openedFile, content)
+		}
+
 		selectedFile = event.openedFile
 		debouncingAnalyze()
 	}
@@ -262,6 +265,13 @@ class KotlinLanguageServer : ILanguageServer {
 		if (!DocumentUtils.isKotlinFile(event.changedFile)) {
 			return
 		}
+
+		compiler?.compilationEnvironmentFor(event.changedFile)?.apply {
+			val content = FileManager.getDocumentContents(event.changedFile)
+			logger.info("Notifying KtFileManager for file {} with contents {}", event.changedFile, content)
+			fileManager.onFileContentChanged(event.changedFile, content)
+		}
+
 		debouncingAnalyze()
 	}
 
@@ -272,13 +282,29 @@ class KotlinLanguageServer : ILanguageServer {
 			return
 		}
 
-		diagnosticProvider?.clearTimestamp(event.closedFile)
+		compiler?.compilationEnvironmentFor(event.closedFile)?.apply {
+			fileManager.onFileClosed(event.closedFile)
+			fileManager.clearAnalyzeTimestampOf(event.closedFile)
+		}
+
 		if (FileManager.getActiveDocumentCount() == 0) {
 			selectedFile = null
 			analyzeJob?.cancel("No active files")
 		}
 	}
 
+	@Subscribe(threadMode = ThreadMode.ASYNC)
+	@Suppress("unused")
+	fun onDocumentSaved(event: DocumentSaveEvent) {
+		if (!DocumentUtils.isKotlinFile(event.savedFile)) {
+			return
+		}
+
+		compiler?.compilationEnvironmentFor(event.savedFile)?.apply {
+			fileManager.onFileSaved(event.savedFile)
+		}
+	}
+
 	@Subscribe(threadMode = ThreadMode.ASYNC)
 	@Suppress("unused")
 	fun onDocumentSelected(event: DocumentSelectedEvent) {

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

@@ -0,0 +1,193 @@
+package com.itsaky.androidide.lsp.kotlin
+
+import org.jetbrains.kotlin.analysis.api.KaSession
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.analyzeCopy
+import org.jetbrains.kotlin.analysis.api.projectStructure.KaDanglingFileResolutionMode
+import org.jetbrains.kotlin.com.intellij.openapi.editor.Document
+import org.jetbrains.kotlin.com.intellij.openapi.vfs.StandardFileSystems
+import org.jetbrains.kotlin.com.intellij.openapi.vfs.VirtualFileManager
+import org.jetbrains.kotlin.com.intellij.psi.PsiDocumentManager
+import org.jetbrains.kotlin.com.intellij.psi.PsiManager
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.slf4j.LoggerFactory
+import java.nio.file.Path
+import java.util.concurrent.ConcurrentHashMap
+import kotlin.io.path.name
+import kotlin.io.path.pathString
+import kotlin.time.Clock
+import kotlin.time.Instant
+
+/**
+ * Manages [KtFile] instances for all open files.
+ */
+class KtFileManager(
+	private val psiFactory: KtPsiFactory,
+	private val psiManager: PsiManager,
+	private val psiDocumentManager: PsiDocumentManager,
+) : FileEventConsumer, AutoCloseable {
+
+	companion object {
+		private val logger = LoggerFactory.getLogger(KtFileManager::class.java)
+	}
+
+	private val entries = ConcurrentHashMap<Path, ManagedFile>()
+
+	@ConsistentCopyVisibility
+	data class ManagedFile @Deprecated("Use ManagedFile.create instead") internal constructor(
+		val file: Path,
+		val diskKtFile: KtFile,
+		@Volatile var inMemoryKtFile: KtFile,
+		val document: Document,
+		@Volatile var lastModified: Instant,
+		@Volatile var isDirty: Boolean,
+		@Volatile var analyzeTimestamp: Instant,
+	) {
+
+		/**
+		 * Analyze this [ManagedFile] contents.
+		 *
+		 * @param action The analysis action.
+		 */
+		fun <R> analyze(action: KaSession.(file: KtFile) -> R): R {
+			if (diskKtFile === inMemoryKtFile) {
+				return analyze(useSiteElement = inMemoryKtFile) { action(inMemoryKtFile) }
+			}
+
+			return analyzeCopy(
+				useSiteElement = inMemoryKtFile,
+				resolutionMode = KaDanglingFileResolutionMode.PREFER_SELF
+			) {
+				action(inMemoryKtFile)
+			}
+		}
+
+		fun createInMemoryFileWithContent(psiFactory: KtPsiFactory, content: String): KtFile {
+			val inMemoryFile = psiFactory.createFile(file.name, content)
+			inMemoryFile.originalFile = diskKtFile
+			return inMemoryFile
+		}
+
+		companion object {
+			@Suppress("DEPRECATION")
+			fun create(
+				file: Path,
+				ktFile: KtFile,
+				document: Document,
+				inMemoryKtFile: KtFile = ktFile,
+				lastModified: Instant = Clock.System.now(),
+				isDirty: Boolean = false,
+				analyzeTimestamp: Instant = Instant.DISTANT_PAST,
+			) =
+				ManagedFile(
+					file = file,
+					diskKtFile = ktFile,
+					inMemoryKtFile = inMemoryKtFile,
+					document = document,
+					lastModified = lastModified,
+					isDirty = isDirty,
+					analyzeTimestamp = analyzeTimestamp,
+				)
+		}
+	}
+
+	override fun onFileOpened(path: Path, content: String) {
+		logger.debug("onFileOpened: {}", path)
+
+		entries[path]?.let { existing ->
+			logger.info("File is already opened, updating content")
+			updateDocumentContent(existing, content)
+			return
+		}
+
+		val ktFile = resolveKtFile(path)
+
+		if (ktFile == null) {
+			logger.warn("Cannot resolve KtFile for: {}", path)
+			return
+		}
+
+		val document = getOrCreateDocument(ktFile)
+		if (document == null) {
+			logger.warn("Cannot obtain Document for: {}", path)
+			return
+		}
+
+		logger.info("Creating managed file entry")
+		val entry = ManagedFile.create(
+			file = path,
+			ktFile = ktFile,
+			document = document,
+		)
+
+		entries[path] = entry
+
+		updateDocumentContent(entry, content)
+		logger.debug("File opened and managed: {}", path)
+		return
+	}
+
+	override fun onFileContentChanged(path: Path, content: String) {
+		logger.debug("onFileContentChanged: {}", path)
+		val entry = entries[path] ?: run {
+			logger.debug("Content changed for unmanaged file: {}. Ignoring.", path)
+			return
+		}
+
+		updateDocumentContent(entry, content)
+	}
+
+	override fun onFileSaved(path: Path) {
+		val entry = entries[path] ?: return
+		entry.isDirty = false
+
+		logger.debug("File saved: {}", path)
+	}
+
+	override fun onFileClosed(path: Path) {
+		entries.remove(path) ?: return
+		logger.debug("File closed: {}", path)
+	}
+
+	fun getOpenFile(path: Path): ManagedFile? = entries[path]
+
+	fun allOpenFiles(): Collection<ManagedFile> =
+		entries.values.toList()
+
+	fun clearAnalyzeTimestampOf(file: Path) {
+		val managed = getOpenFile(file) ?: return
+		managed.analyzeTimestamp = Instant.DISTANT_PAST
+	}
+
+	private fun resolveKtFile(path: Path): KtFile? {
+		val vfs = VirtualFileManager.getInstance()
+			.getFileSystem(StandardFileSystems.FILE_PROTOCOL)
+
+		val virtualFile = vfs.refreshAndFindFileByPath(path.pathString)
+			?: return null
+
+		val psiFile = psiManager.findFile(virtualFile)
+
+		return psiFile as? KtFile
+	}
+
+	private fun getOrCreateDocument(ktFile: KtFile): Document? {
+		return psiDocumentManager.getDocument(ktFile)
+	}
+
+	private fun updateDocumentContent(entry: ManagedFile, content: String) {
+		logger.info("Updating doc content for {}", entry.file)
+
+		val normalized = content.replace("\r", "")
+		if (entry.inMemoryKtFile.text == normalized) return
+
+		entry.inMemoryKtFile = entry.createInMemoryFileWithContent(psiFactory, content)
+		entry.lastModified = Clock.System.now()
+		entry.isDirty = true
+	}
+
+	override fun close() {
+		entries.clear()
+	}
+}

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

@@ -1,5 +1,7 @@
 package com.itsaky.androidide.lsp.kotlin.compiler
 
+import com.itsaky.androidide.lsp.kotlin.FileEventConsumer
+import com.itsaky.androidide.lsp.kotlin.KtFileManager
 import org.jetbrains.kotlin.analysis.api.KaExperimentalApi
 import org.jetbrains.kotlin.analysis.api.platform.declarations.KotlinAnnotationsResolverFactory
 import org.jetbrains.kotlin.analysis.api.platform.declarations.KotlinDeclarationProviderFactory
@@ -61,9 +63,13 @@ class CompilationEnvironment(
 
 	var session: StandaloneAnalysisAPISession
 		private set
+
 	var parser: KtPsiFactory
 		private set
 
+	var fileManager: KtFileManager
+		private set
+
 	val psiManager: PsiManager
 		get() = PsiManager.getInstance(session.project)
 
@@ -101,6 +107,7 @@ class CompilationEnvironment(
 	init {
 		session = buildSession()
 		parser = KtPsiFactory(session.project, eventSystemEnabled = enableParserEventSystem)
+		fileManager = KtFileManager(parser, psiManager, psiDocumentManager)
 
 		projectModel.addListener(this)
 	}
@@ -225,6 +232,7 @@ class CompilationEnvironment(
 	}
 
 	override fun close() {
+		fileManager.close()
 		projectModel.removeListener(this)
 		disposable.dispose()
 	}