[inline-completion] IDEA-333789: add possibility to stream multi-line suggestions

GitOrigin-RevId: b09d83667c29306781a40f0388ed9130398d1770

platform/platform-impl/src/com/intellij/codeInsight/inline/completion/elements/InlineCompletionGrayTextElement.kt

@@ -1,19 +1,9 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.codeInsight.inline.completion.elements
 
-import com.intellij.codeInsight.inline.completion.render.InlineBlockElementRenderer
-import com.intellij.codeInsight.inline.completion.render.InlineSuffixRenderer
-import com.intellij.openapi.actionSystem.IdeActions
+import com.intellij.codeInsight.inline.completion.render.InlineCompletionGrayTextElementRenderer
 import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.editor.EditorCustomElementRenderer
-import com.intellij.openapi.editor.Inlay
-import com.intellij.openapi.editor.VisualPosition
-import com.intellij.openapi.editor.ex.util.EditorActionAvailabilityHint
-import com.intellij.openapi.editor.ex.util.addActionAvailabilityHint
-import com.intellij.openapi.editor.markup.TextAttributes
-import com.intellij.openapi.util.Disposer
 import org.jetbrains.annotations.ApiStatus
-import java.awt.Graphics
 import java.awt.Rectangle
 
 class InlineCompletionGrayTextElement(override val text: String) : InlineCompletionElement {
@@ -21,10 +11,12 @@ class InlineCompletionGrayTextElement(override val text: String) : InlineComplet
   override fun toPresentable(): InlineCompletionElement.Presentable = Presentable(this)
 
   open class Presentable(override val element: InlineCompletionElement) : InlineCompletionElement.Presentable {
-    private var suffixInlay: Inlay<*>? = null
-    private var blockInlay: Inlay<*>? = null
+    private var offset: Int? = null
+    private var rectangle: Rectangle? = null
 
-    override fun isVisible(): Boolean = suffixInlay != null || blockInlay != null
+    override fun isVisible(): Boolean = offset != null
+    override fun startOffset(): Int? = offset
+    override fun endOffset(): Int? = offset
 
     /**
      * Temporal workaround for an internal plugin. **Should not be used.**
@@ -33,75 +25,16 @@ class InlineCompletionGrayTextElement(override val text: String) : InlineComplet
     @ApiStatus.Experimental
     protected open fun getText(): String = element.text
 
+    // TODO It returns the rectangle for all the elements, but now it's not important
+    override fun getBounds(): Rectangle? = rectangle
+
     override fun render(editor: Editor, offset: Int) {
-      val text = getText()
-      if (text.isEmpty()) return
-      val lines = text.lines()
-      renderSuffix(editor, lines, offset)
-      if (lines.size > 1) {
-        renderBlock(lines.drop(1), editor, offset)
-      }
+      rectangle = InlineCompletionGrayTextElementRenderer.render(editor, getText(), offset, this)
+      this.offset = offset
     }
 
-    override fun getBounds(): Rectangle? {
-      val bounds = suffixInlay?.bounds?.let { Rectangle(it) }
-      blockInlay?.bounds?.let { bounds?.add(Rectangle(it)) }
-      return bounds
-    }
-
-    override fun startOffset(): Int? = suffixInlay?.offset
-    override fun endOffset(): Int? = suffixInlay?.offset
-
     override fun dispose() {
-      blockInlay?.also(Disposer::dispose)
-      blockInlay = null
-      suffixInlay?.also(Disposer::dispose)
-      suffixInlay = null
-    }
-
-    private fun renderSuffix(editor: Editor, lines: List<String>, offset: Int) {
-      // The following is a hacky solution to the effect described in ML-977
-      // ML-1781 Inline completion renders on the left to the caret after moving it
-      editor.forceLeanLeft()
-
-      val line = lines.first()
-      if (line.isEmpty()) {
-        suffixInlay = editor.inlayModel.addInlineElement(editor.caretModel.offset, object : EditorCustomElementRenderer {
-          override fun calcWidthInPixels(inlay: Inlay<*>) = 1
-          override fun calcHeightInPixels(inlay: Inlay<*>) = 1
-          override fun paint(inlay: Inlay<*>, g: Graphics, targetRegion: Rectangle, textAttributes: TextAttributes) {}
-        })
-        return
-      }
-      editor.inlayModel.execute(true) {
-        val element = editor.inlayModel.addInlineElement(offset, true, InlineSuffixRenderer(editor, line)) ?: return@execute
-        element.addActionAvailabilityHint(EditorActionAvailabilityHint(
-          IdeActions.ACTION_INSERT_INLINE_COMPLETION,
-          EditorActionAvailabilityHint.AvailabilityCondition.CaretOnStart,
-        ))
-        suffixInlay = element
-      }
-    }
-
-    private fun renderBlock(
-      lines: List<String>,
-      editor: Editor,
-      offset: Int
-    ) {
-      val element = editor.inlayModel.addBlockElement(
-        offset, true, false, 1,
-        InlineBlockElementRenderer(editor, lines)
-      ) ?: return
-
-      blockInlay = element
-    }
-
-    private fun Editor.forceLeanLeft() {
-      val visualPosition = caretModel.visualPosition
-      if (visualPosition.leansRight) {
-        val leftLeaningPosition = VisualPosition(visualPosition.line, visualPosition.column, false)
-        caretModel.moveToVisualPosition(leftLeaningPosition)
-      }
+      offset = null
     }
   }
 }

platform/platform-impl/src/com/intellij/codeInsight/inline/completion/render/InlineCompletionGrayTextElementRenderer.kt

@@ -0,0 +1,176 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight.inline.completion.render
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.actionSystem.IdeActions
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.Inlay
+import com.intellij.openapi.editor.VisualPosition
+import com.intellij.openapi.editor.ex.util.EditorActionAvailabilityHint
+import com.intellij.openapi.editor.ex.util.addActionAvailabilityHint
+import com.intellij.openapi.observable.util.whenDisposed
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.util.removeUserData
+import com.intellij.util.concurrency.ThreadingAssertions
+import com.intellij.util.concurrency.annotations.RequiresEdt
+import org.jetbrains.annotations.ApiStatus
+import java.awt.Rectangle
+
+// TODO docs
+@ApiStatus.Experimental
+class InlineCompletionGrayTextElementRenderer private constructor(
+  editor: Editor,
+  offset: Int,
+  private val onDispose: () -> Unit
+) {
+
+  private val renderer = Renderer(editor, offset)
+  private var elementsCounter = 0
+  private var isActive = true
+
+  private fun append(text: String, disposable: Disposable): Rectangle? {
+    check(isActive) { "Cannot render an element since the renderer is already disposed." }
+    disposable.whenDisposed {
+      // TODO exclude [text] from rendering
+      ThreadingAssertions.assertEventDispatchThread()
+      elementsCounter--
+      if (elementsCounter == 0) {
+        cleanUp()
+      }
+    }
+    elementsCounter++
+    return renderer.append(text)
+  }
+
+  private fun cleanUp() {
+    isActive = false
+    Disposer.dispose(renderer)
+    onDispose()
+  }
+
+  private class Renderer(private val editor: Editor, private val offset: Int) : Disposable {
+
+    private var suffixInlay: Inlay<*>? = null
+    private var blockInlay: Inlay<*>? = null
+    private val currentLines = mutableListOf<String>()
+
+    fun append(text: String): Rectangle? {
+      if (text.isEmpty()) {
+        return null
+      }
+      val lines = text.lines()
+      if (currentLines.isEmpty()) {
+        currentLines.add("")
+      }
+      currentLines[currentLines.size - 1] += lines.first()
+      if (lines.size > 1) {
+        currentLines += lines.subList(1, lines.size)
+      }
+      return render()
+    }
+
+    override fun dispose() {
+      reset()
+      currentLines.clear()
+    }
+
+    private fun render(): Rectangle? {
+      editor.forceLeanLeft()
+      editor.inlayModel.execute(true) {
+        reset()
+        renderSuffix()
+        renderBlock()
+      }
+      return getRectangle()
+    }
+
+    private fun renderSuffix() {
+      if (currentLines.isEmpty() || currentLines.first().isEmpty()) {
+        return
+      }
+      val suffix = currentLines.first()
+      val element = editor.inlayModel.addInlineElement(offset, true, InlineSuffixRenderer(editor, suffix))
+      if (element != null) {
+        element.addActionAvailabilityHint(
+          EditorActionAvailabilityHint(
+            IdeActions.ACTION_INSERT_INLINE_COMPLETION,
+            EditorActionAvailabilityHint.AvailabilityCondition.CaretOnStart,
+          )
+        )
+        suffixInlay = element
+      }
+    }
+
+    private fun renderBlock() {
+      if (currentLines.size <= 1) {
+        return
+      }
+      val lines = currentLines.subList(1, currentLines.size)
+      blockInlay = editor.inlayModel.addBlockElement(
+        offset,
+        true,
+        false,
+        1,
+        InlineBlockElementRenderer(editor, lines)
+      )
+    }
+
+    private fun getRectangle(): Rectangle? {
+      val bounds = suffixInlay?.bounds?.let { Rectangle(it) }
+      blockInlay?.bounds?.let { bounds?.add(Rectangle(it)) }
+      return bounds
+    }
+
+    private fun reset() {
+      suffixInlay?.let { Disposer.dispose(it) }
+      blockInlay?.let { Disposer.dispose(it) }
+      suffixInlay = null
+      blockInlay = null
+    }
+
+    private fun Editor.forceLeanLeft() {
+      val visualPosition = caretModel.visualPosition
+      if (visualPosition.leansRight) {
+        val leftLeaningPosition = VisualPosition(visualPosition.line, visualPosition.column, false)
+        caretModel.moveToVisualPosition(leftLeaningPosition)
+      }
+    }
+  }
+
+  companion object {
+
+    private val STORAGE_KEY = Key.create<Storage<Int, InlineCompletionGrayTextElementRenderer>>("inline.completion.gray.text.render")
+
+    @ApiStatus.Experimental
+    @RequiresEdt
+    fun render(editor: Editor, text: String, offset: Int, disposable: Disposable): Rectangle? {
+      ThreadingAssertions.assertEventDispatchThread()
+
+      val storage = editor.getUserData(STORAGE_KEY) ?: Storage()
+      editor.putUserData(STORAGE_KEY, storage)
+
+      val renderer = storage.getOrInitialize(offset) {
+        InlineCompletionGrayTextElementRenderer(editor, offset) {
+          storage.remove(offset)
+          if (storage.isEmpty()) {
+            editor.removeUserData(STORAGE_KEY)
+          }
+        }
+      }
+      return renderer.append(text, disposable)
+    }
+
+    private class Storage<K : Any, V : Any> {
+      private val map = mutableMapOf<K, V>()
+
+      fun getOrInitialize(key: K, init: () -> V): V = map.computeIfAbsent(key) { init() }
+
+      fun remove(key: K) {
+        checkNotNull(map.remove(key))
+      }
+
+      fun isEmpty(): Boolean = map.isEmpty()
+    }
+  }
+}