Merge pull request continuedev#3494 from lkk214/refactor/completion-inlay-rendering

refactor: Adjust inlay rendering and improve completion logic

Author: sestinj

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/activities/ContinuePluginStartupActivity.kt

@@ -10,7 +10,7 @@ import com.github.continuedev.continueintellijextension.listeners.ContinuePlugin
 import com.github.continuedev.continueintellijextension.services.ContinueExtensionSettings
 import com.github.continuedev.continueintellijextension.services.ContinuePluginService
 import com.github.continuedev.continueintellijextension.services.SettingsListener
-import com.intellij.openapi.Disposable
+import com.github.continuedev.continueintellijextension.utils.toUriOrNull
 import com.intellij.openapi.actionSystem.KeyboardShortcut
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.ApplicationNamesInfo
@@ -162,7 +162,7 @@ class ContinuePluginStartupActivity : StartupActivity, DumbAware {
                 override fun after(events: List<VFileEvent>) {
                     // Collect all relevant URIs for deletions
                     val deletedURIs = events.filterIsInstance<VFileDeleteEvent>()
-                        .map { event -> event.file.url }
+                        .mapNotNull { event -> event.file.toUriOrNull() }
 
                     // Send "files/deleted" message if there are any deletions
                     if (deletedURIs.isNotEmpty()) {
@@ -172,7 +172,7 @@ class ContinuePluginStartupActivity : StartupActivity, DumbAware {
 
                     // Collect all relevant URIs for content changes
                     val changedURIs = events.filterIsInstance<VFileContentChangeEvent>()
-                        .map { event -> event.file.url }
+                        .mapNotNull { event -> event.file.toUriOrNull() }
 
                     // Send "files/changed" message if there are any content changes
                     if (changedURIs.isNotEmpty()) {
@@ -220,7 +220,7 @@ class ContinuePluginStartupActivity : StartupActivity, DumbAware {
             // Reload the WebView
             continuePluginService?.let { pluginService ->
                 val allModulePaths = ModuleManager.getInstance(project).modules
-                    .flatMap { module -> ModuleRootManager.getInstance(module).contentRoots.map { it.url } }
+                    .flatMap { module -> ModuleRootManager.getInstance(module).contentRoots.mapNotNull { it.toUriOrNull() } }
 
                 val topLevelModulePaths = allModulePaths
                     .filter { modulePath -> allModulePaths.none { it != modulePath && modulePath.startsWith(it) } }

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/autocomplete/AcceptAutocompleteAction.kt

@@ -22,4 +22,4 @@ class AcceptAutocompleteAction : EditorAction(object : EditorActionHandler() {
                 && autocompleteService.pendingCompletion?.text != null
         return enabled
     }
-}) {}
+})

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/autocomplete/AutocompleteService.kt

@@ -2,6 +2,7 @@ package com.github.continuedev.continueintellijextension.autocomplete
 
 import com.github.continuedev.continueintellijextension.services.ContinueExtensionSettings
 import com.github.continuedev.continueintellijextension.services.ContinuePluginService
+import com.github.continuedev.continueintellijextension.utils.toUriOrNull
 import com.github.continuedev.continueintellijextension.utils.uuid
 import com.intellij.injected.editor.VirtualFileWindow
 import com.intellij.openapi.application.*
@@ -10,8 +11,10 @@ import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.InlayProperties
+import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.wm.WindowManager
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
@@ -32,9 +35,24 @@ fun PsiElement.isInjectedText(): Boolean {
     return false
 }
 
+fun Editor.addInlayElement(
+    lines: List<String>,
+    offset: Int,
+    properties: InlayProperties
+) {
+    if (this is EditorImpl) {
+        if (lines[0].isNotEmpty()) {
+            inlayModel.addInlineElement(offset, properties, ContinueInlayRenderer(listOf(lines[0])))
+        }
+        if (lines.size > 1) {
+            inlayModel.addBlockElement(offset, properties, ContinueInlayRenderer(lines.drop(1)))
+        }
+    }
+}
+
 @Service(Service.Level.PROJECT)
 class AutocompleteService(private val project: Project) {
-    var pendingCompletion: PendingCompletion? = null;
+    var pendingCompletion: PendingCompletion? = null
     private val autocompleteLookupListener = project.service<AutocompleteLookupListener>()
     private var widget: AutocompleteSpinnerWidget? = null
 
@@ -66,23 +84,23 @@ class AutocompleteService(private val project: Project) {
 
         // Request a completion from the core
         val virtualFile = FileDocumentManager.getInstance().getFile(editor.document)
+
+        val uri = virtualFile?.toUriOrNull() ?: return
+
+        val line = editor.caretModel.primaryCaret.logicalPosition.line
         val column = editor.caretModel.primaryCaret.logicalPosition.column
         val input = mapOf(
             "completionId" to completionId,
-            "filepath" to virtualFile?.url,
+            "filepath" to uri,
             "pos" to mapOf(
-                "line" to editor.caretModel.primaryCaret.logicalPosition.line,
+                "line" to line,
                 "character" to column
             ),
             "recentlyEditedFiles" to emptyList<String>(),
             "recentlyEditedRanges" to emptyList<String>(),
             "clipboardText" to ""
         )
 
-        val lineStart = editor.document.getLineStartOffset(editor.caretModel.primaryCaret.logicalPosition.line)
-        val lineEnd = editor.document.getLineEndOffset(editor.caretModel.primaryCaret.logicalPosition.line)
-        val lineLength = lineEnd - lineStart
-
         project.service<ContinuePluginService>().coreMessenger?.request(
             "autocomplete/complete",
             input,
@@ -95,9 +113,8 @@ class AutocompleteService(private val project: Project) {
                     val completion = completions[0].toString()
                     val finalTextToInsert = deduplicateCompletion(editor, offset, completion)
 
-                    if (shouldRenderCompletion(finalTextToInsert, column, lineLength, editor)) {
+                    if (shouldRenderCompletion(finalTextToInsert, offset, line, editor)) {
                         renderCompletion(editor, offset, finalTextToInsert)
-                        pendingCompletion = PendingCompletion(editor, offset, completionId, finalTextToInsert)
                         // Hide auto-popup
 //                    AutoPopupController.getInstance(project).cancelAllRequests()
                     }
@@ -106,13 +123,19 @@ class AutocompleteService(private val project: Project) {
         )
     }
 
-    private fun shouldRenderCompletion(completion: String, column: Int, lineLength: Int, editor: Editor): Boolean {
-        if (completion.isEmpty()) {
+    private fun shouldRenderCompletion(completion: String, offset: Int, line: Int, editor: Editor): Boolean {
+        if (completion.isEmpty() || runReadAction { offset != editor.caretModel.offset }) {
             return false
         }
 
+        if (completion.lines().size == 1) {
+            return true
+        }
+
+        val endOffset = editor.document.getLineEndOffset(line)
+
         // Do not render if completion is multi-line and caret is in middle of line
-        return !(completion.lines().size > 1 && column < lineLength)
+        return offset <= endOffset && editor.document.getText(TextRange(offset, endOffset)).isBlank()
     }
 
     private fun deduplicateCompletion(editor: Editor, offset: Int, completion: String): String {
@@ -126,9 +149,9 @@ class AutocompleteService(private val project: Project) {
 
             val N = 10
             var textAfterCursor = if (caretOffset + N <= document.textLength) {
-                document.getText(com.intellij.openapi.util.TextRange(caretOffset, caretOffset + N))
+                document.getText(TextRange(caretOffset, caretOffset + N))
             } else {
-                document.getText(com.intellij.openapi.util.TextRange(caretOffset, document.textLength))
+                document.getText(TextRange(caretOffset, document.textLength))
             }
 
             // Avoid truncating the completion text when the text after the cursor is blank
@@ -171,19 +194,9 @@ class AutocompleteService(private val project: Project) {
                 properties.relatesToPrecedingText(true)
                 properties.disableSoftWrapping(true)
 
-                if (completion.lines().size > 1) {
-                    editor.inlayModel.addBlockElement(
-                        offset,
-                        properties,
-                        ContinueMultilineCustomElementRenderer(editor, completion)
-                    )
-                } else {
-                    editor.inlayModel.addInlineElement(
-                        offset,
-                        properties,
-                        ContinueCustomElementRenderer(editor, completion)
-                    )
-                }
+                val lines = completion.lines()
+                pendingCompletion = pendingCompletion?.copy(text = lines.joinToString("\n"))
+                editor.addInlayElement(lines, offset, properties)
 
 //                val attributes = TextAttributes().apply {
 //                    backgroundColor = JBColor.GREEN
@@ -211,7 +224,7 @@ class AutocompleteService(private val project: Project) {
             ({})
         )
         invokeLater {
-            clearCompletions(editor)
+            clearCompletions(editor, completion)
         }
     }
 
@@ -274,23 +287,14 @@ class AutocompleteService(private val project: Project) {
         project.service<ContinuePluginService>().coreMessenger?.request("autocomplete/cancel", null, null, ({}))
     }
 
-    fun clearCompletions(editor: Editor) {
+    fun clearCompletions(editor: Editor, completion: PendingCompletion? = pendingCompletion) {
         if (isInjectedFile(editor)) return
 
-        if (pendingCompletion != null) {
-            cancelCompletion(pendingCompletion!!)
-            pendingCompletion = null
-        }
-        editor.inlayModel.getInlineElementsInRange(0, editor.document.textLength).forEach {
-            if (it.renderer is ContinueCustomElementRenderer) {
-                it.dispose()
-            }
-        }
-        editor.inlayModel.getBlockElementsInRange(0, editor.document.textLength).forEach {
-            if (it.renderer is ContinueMultilineCustomElementRenderer) {
-                it.dispose()
-            }
+        if (completion != null) {
+            cancelCompletion(completion)
+            if (completion.completionId == pendingCompletion?.completionId) pendingCompletion = null
         }
+        disposeInlayRenderer(editor)
     }
 
     private fun isInjectedFile(editor: Editor): Boolean {
@@ -306,13 +310,17 @@ class AutocompleteService(private val project: Project) {
     fun hideCompletions(editor: Editor) {
         if (isInjectedFile(editor)) return
 
+        disposeInlayRenderer(editor)
+    }
+
+    private fun disposeInlayRenderer(editor: Editor) {
         editor.inlayModel.getInlineElementsInRange(0, editor.document.textLength).forEach {
-            if (it.renderer is ContinueCustomElementRenderer) {
+            if (it.renderer is ContinueInlayRenderer) {
                 it.dispose()
             }
         }
         editor.inlayModel.getBlockElementsInRange(0, editor.document.textLength).forEach {
-            if (it.renderer is ContinueMultilineCustomElementRenderer) {
+            if (it.renderer is ContinueInlayRenderer) {
                 it.dispose()
             }
         }

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/autocomplete/ContinueCustomElementRenderer.kt

@@ -0,0 +1,58 @@
+package com.github.continuedev.continueintellijextension.autocomplete
+
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.EditorCustomElementRenderer
+import com.intellij.openapi.editor.Inlay
+import com.intellij.openapi.editor.colors.EditorFontType
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.openapi.editor.markup.TextAttributes
+import com.intellij.ui.JBColor
+import com.intellij.util.ui.UIUtil
+import java.awt.Font
+import java.awt.Graphics
+import java.awt.Rectangle
+
+/**
+ * The `ContinueInlayRenderer` class is responsible for rendering custom inlay elements within an editor.
+ * It implements the [EditorCustomElementRenderer] interface to provide custom rendering logic for inlays.
+ *
+ * This renderer is designed to display a list of text lines (`lines`) within the editor, calculating the
+ * necessary width and height based on the content and rendering each line with appropriate font and color.
+ *
+ * @author lk
+ */
+class ContinueInlayRenderer(val lines: List<String>) : EditorCustomElementRenderer {
+    override fun calcWidthInPixels(inlay: Inlay<*>): Int {
+        var maxLen = 0;
+        for (line in lines) {
+            val len = (inlay.editor as EditorImpl).getFontMetrics(Font.PLAIN).stringWidth(line)
+            if (len > maxLen) {
+                maxLen = len
+            }
+        }
+        return maxLen
+    }
+
+    override fun calcHeightInPixels(inlay: Inlay<*>): Int {
+        return (inlay.editor as EditorImpl).lineHeight * lines.size
+    }
+
+    private fun font(editor: Editor): Font {
+        val editorFont = editor.colorsScheme.getFont(EditorFontType.PLAIN)
+        return UIUtil.getFontWithFallbackIfNeeded(editorFont, lines.joinToString("\n"))
+            .deriveFont(editor.colorsScheme.editorFontSize)
+    }
+
+    override fun paint(inlay: Inlay<*>, g: Graphics, targetRegion: Rectangle, textAttributes: TextAttributes) {
+        val editor = inlay.editor
+        g.color = JBColor.GRAY
+        g.font = font(editor)
+        var additionalYOffset = 0
+        val ascent = editor.ascent
+        val lineHeight = editor.lineHeight
+        for (line in lines) {
+            g.drawString(line, targetRegion.x, targetRegion.y + ascent + additionalYOffset)
+            additionalYOffset += lineHeight
+        }
+    }
+}