From 3affb154bb8c30161012d6e018f99c64a013c917 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A1lint=20Hegyi?= <bhegyi@gradle.com>
Date: Fri, 26 Jul 2024 16:06:47 +0100
Subject: [PATCH] Improve how possible mutations are calculated, add DOM
 removal when file closed

---
 .../lsp/DeclarativeTextDocumentService.kt     | 52 ++++++++++---------
 .../lsp/storage/VersionedDocumentStore.kt     |  4 ++
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/lsp/src/main/kotlin/org/gradle/declarative/lsp/DeclarativeTextDocumentService.kt b/lsp/src/main/kotlin/org/gradle/declarative/lsp/DeclarativeTextDocumentService.kt
index 85c0b63..3b6a721 100644
--- a/lsp/src/main/kotlin/org/gradle/declarative/lsp/DeclarativeTextDocumentService.kt
+++ b/lsp/src/main/kotlin/org/gradle/declarative/lsp/DeclarativeTextDocumentService.kt
@@ -36,11 +36,11 @@ class DeclarativeTextDocumentService : TextDocumentService, LanguageClientAware
     private val availableMutations = listOf(
         AddDependency()
     )
-    private val applicableMutations: List<MutationApplicability> = mutableListOf()
 
     private lateinit var client: LanguageClient
     private lateinit var resources: ResolvedDeclarativeResourcesModel
     private lateinit var schemaEvaluator: SimpleAnalysisEvaluator
+    private var applicableMutations: List<MutationApplicability> = emptyList()
 
     override fun connect(client: LanguageClient?) {
         this.client = client!!
@@ -56,34 +56,36 @@ class DeclarativeTextDocumentService : TextDocumentService, LanguageClientAware
     }
 
     override fun didOpen(params: DidOpenTextDocumentParams?) {
-        params?.textDocument?.uri?.let { uri ->
-            URI(uri)
-        }?.let { uri ->
-            LOGGER.trace("Opened document: {}", uri)
+        params?.let {
+            LOGGER.trace("Opened document: {}", it.textDocument.uri)
+
+            val uri = URI(it.textDocument.uri)
             val text = File(uri).readText()
             documentStore.storeInitial(uri, parse(uri, text))
-            LOGGER.trace("Stored declarative model for document: {}", uri)
-
             processDocument(uri)
         }
     }
 
     override fun didChange(params: DidChangeTextDocumentParams?) {
-        params?.let { nonNullParams ->
-            val uri = URI(nonNullParams.textDocument.uri)
-            nonNullParams.contentChanges.forEach { change ->
+        params?.let {
+            LOGGER.trace("Changed document: ${params.textDocument.uri}")
+            val uri = URI(it.textDocument.uri)
+            it.contentChanges.forEach { change ->
                 documentStore.storeVersioned(
                     uri,
-                    nonNullParams.textDocument.version,
+                    it.textDocument.version,
                     parse(uri, change.text)
                 )
                 processDocument(uri)
             }
-            LOGGER.trace("Changed document: ${params.textDocument?.uri}")
         }
     }
 
     override fun didClose(params: DidCloseTextDocumentParams?) {
+        params?.let {
+            val uri = URI(it.textDocument.uri)
+            documentStore.remove(uri)
+        }
         LOGGER.trace("Closed document: ${params?.textDocument?.uri}")
     }
 
@@ -92,10 +94,10 @@ class DeclarativeTextDocumentService : TextDocumentService, LanguageClientAware
     }
 
     override fun hover(params: HoverParams?): CompletableFuture<Hover> {
-        val hover = params?.let { nonNullParams ->
-            val uri = URI(nonNullParams.textDocument.uri)
+        val hover = params?.let {
+            val uri = URI(it.textDocument.uri)
             withDom(uri) { dom ->
-                val position = nonNullParams.position
+                val position = it.position
 
                 // LSPs are supplying 0-based line and column numbers, while the DSL model is 1-based
                 val visitor = LocationMatchingVisitor(position.line + 1, position.character + 1)
@@ -131,15 +133,8 @@ class DeclarativeTextDocumentService : TextDocumentService, LanguageClientAware
     }
 
     override fun codeAction(params: CodeActionParams?): CompletableFuture<MutableList<Either<Command, CodeAction>>> {
-        params?.let { nonNullParams ->
-            val uri = URI(nonNullParams.textDocument.uri)
-            withDom(uri) { dom ->
-                val mac = MutationApplicabilityChecker(resources.analysisSchema, dom.result)
-                availableMutations.forEach { mutation ->
-                    val applicability = mac.checkApplicability(mutation)
-                    LOGGER.info("Applicability: {}", applicability)
-                }
-            }
+        params?.let {
+            // TODO: finish this
         }
 
         return CompletableFuture.completedFuture(null)
@@ -147,8 +142,15 @@ class DeclarativeTextDocumentService : TextDocumentService, LanguageClientAware
 
     // Common processing -----------------------------------------------------------------------------------------------
 
-    private fun processDocument(uri: URI) = withDom(uri) {
+    private fun processDocument(uri: URI) = withDom(uri) { dom ->
+        applicableMutations = processMutationApplications(dom)
+    }
 
+    private fun processMutationApplications(dom: DocumentOverlayResult): List<MutationApplicability> {
+        val mac = MutationApplicabilityChecker(resources.analysisSchema, dom.result)
+        return availableMutations.flatMap { mutation ->
+            mac.checkApplicability(mutation)
+        }
     }
 
     // Utility and other member functions ------------------------------------------------------------------------------
diff --git a/lsp/src/main/kotlin/org/gradle/declarative/lsp/storage/VersionedDocumentStore.kt b/lsp/src/main/kotlin/org/gradle/declarative/lsp/storage/VersionedDocumentStore.kt
index afbfe11..0b005cb 100644
--- a/lsp/src/main/kotlin/org/gradle/declarative/lsp/storage/VersionedDocumentStore.kt
+++ b/lsp/src/main/kotlin/org/gradle/declarative/lsp/storage/VersionedDocumentStore.kt
@@ -48,6 +48,10 @@ class VersionedDocumentStore {
         // If we reach this point, the document was not stored
         return false
     }
+
+    fun remove(uri: URI) {
+        store.remove(uri)
+    }
 }
 
 sealed class DocumentStoreEntry {