[KT-40688] Inlay Hints: lambda hints removal on editing

Due to the platform API limitation lambda related hints are treated
separately to all others. Hints removal on editing is not an exclusion
and is now supported as well.

idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinAbstractHintsProvider.kt

@@ -38,6 +38,13 @@ abstract class KotlinAbstractHintsProvider<T : Any> : InlayHintsProvider<T> {
         }
     }
 
+    /**
+     * Similar to [handlePresentations] we need to have a special handler to remove lambda related hints.
+     * Therefore [KotlinLambdasHintsProvider] provides its own "crutch" implementation.
+     */
+    protected open fun handleAfterLineEndHintsRemoval(editor: Editor, resolved: HintType, element: PsiElement) {
+    }
+
     override fun getCollectorFor(file: PsiFile, editor: Editor, settings: T, sink: InlayHintsSink): InlayHintsCollector? {
         return object : FactoryInlayHintsCollector(editor) {
             override fun collect(element: PsiElement, editor: Editor, sink: InlayHintsSink): Boolean {
@@ -45,8 +52,12 @@ abstract class KotlinAbstractHintsProvider<T : Any> : InlayHintsProvider<T> {
                 if (!isElementSupported(resolved, settings)) return true
 
                 val presentations = resolved.provideHints(element).mapNotNull { info -> convert(info, editor.project) }
-                if (presentations.isNotEmpty())
+                if (presentations.isNotEmpty()) {
                     handlePresentations(presentations, editor, sink)
+                } else {
+                    handleAfterLineEndHintsRemoval(editor, resolved, element)
+                }
+
                 return true
             }
 

idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinLambdasHintsProvider.kt

@@ -13,8 +13,13 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorCustomElementRenderer
 import com.intellij.openapi.editor.Inlay
 import com.intellij.openapi.editor.markup.TextAttributes
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.util.application.invokeLater
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtLambdaExpression
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import java.awt.Graphics
 import java.awt.Rectangle
 
@@ -64,6 +69,26 @@ class KotlinLambdasHintsProvider : KotlinAbstractHintsProvider<KotlinLambdasHint
         }
     }
 
+    override fun handleAfterLineEndHintsRemoval(editor: Editor, resolved: HintType, element: PsiElement) {
+        invokeLater {
+            val offset = when (resolved) {
+                HintType.LAMBDA_IMPLICIT_PARAMETER_RECEIVER -> {
+                    val lambdaExpression = (element as? KtFunctionLiteral)?.parent as? KtLambdaExpression
+                    lambdaExpression?.leftCurlyBrace?.textRange?.endOffset
+                }
+                HintType.LAMBDA_RETURN_EXPRESSION -> (element as? KtExpression)?.endOffset
+                else -> null
+            }
+
+            offset?.let {
+                val logicalLine = editor.offsetToLogicalPosition(offset).line
+                editor.inlayModel.getAfterLineEndElementsForLogicalLine(logicalLine)
+                    .filter { it.renderer is LambdaHintsRenderer }
+                    .forEach { it.dispose() }
+            }
+        }
+    }
+
     override fun createConfigurable(settings: Settings): ImmediateConfigurable {
         return createLambdaHintsImmediateConfigurable(settings)
     }