From 78592873372b37b93fe6f0b93d8549a3536e43d1 Mon Sep 17 00:00:00 2001
From: Andrei Klunnyi <andrei.klunnyi@jetbrains.com>
Date: Tue, 3 Nov 2020 19:08:14 +0100
Subject: [PATCH] [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.
---
 .../hints/KotlinAbstractHintsProvider.kt      | 13 +++++++++-
 .../hints/KotlinLambdasHintsProvider.kt       | 25 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinAbstractHintsProvider.kt b/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinAbstractHintsProvider.kt
index 5a104f7a9f2..1ef12fff9be 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinAbstractHintsProvider.kt
+++ b/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
             }
 
diff --git a/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinLambdasHintsProvider.kt b/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinLambdasHintsProvider.kt
index 3dbe6d3b84f..780377bda13 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinLambdasHintsProvider.kt
+++ b/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)
     }