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()
 	}