From 71023cd596c37b08248561a4ea319a342dca05f2 Mon Sep 17 00:00:00 2001
From: Andrei Klunnyi <andrei.klunnyi@jetbrains.com>
Date: Thu, 22 Oct 2020 15:42:22 +0200
Subject: [PATCH] [KT-40688] Inlay Hints: inlay hints are duplicated on code
 editing

Code editing changes inlay hints offset.
InlayModel.getAfterLineEndElementsInRange(int, int) call didn't return
previously existed hint by its "outdated" offset. From now on hints are
taken by logical line.
---
 .../hints/KotlinLambdasHintsProvider.kt       | 38 ++++++++++++++++++-
 1 file changed, 36 insertions(+), 2 deletions(-)

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 0db01dbca73..3dbe6d3b84f 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinLambdasHintsProvider.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/codeInsight/hints/KotlinLambdasHintsProvider.kt
@@ -7,10 +7,16 @@ package org.jetbrains.kotlin.idea.codeInsight.hints
 
 import com.intellij.codeInsight.hints.ImmediateConfigurable
 import com.intellij.codeInsight.hints.InlayHintsSink
+import com.intellij.codeInsight.hints.presentation.InlayPresentation
 import com.intellij.codeInsight.hints.presentation.PresentationRenderer
 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 org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.util.application.invokeLater
+import java.awt.Graphics
+import java.awt.Rectangle
 
 /**
  * Please note that the executable test class is currently not generated. The thing is that [KotlinLambdasHintsProvider] utilizes a hack
@@ -49,8 +55,11 @@ class KotlinLambdasHintsProvider : KotlinAbstractHintsProvider<KotlinLambdasHint
         // sink should remain empty for the outer infrastructure - we place hints ourselves
         invokeLater {
             presentations.forEach { p ->
-                editor.inlayModel.getAfterLineEndElementsInRange(p.offset, p.offset).singleOrNull()?.dispose()
-                editor.inlayModel.addAfterLineEndElement(p.offset, p.relatesToPrecedingText, PresentationRenderer(p.presentation))
+                val logicalLine = editor.offsetToLogicalPosition(p.offset).line
+                editor.inlayModel.getAfterLineEndElementsForLogicalLine(logicalLine)
+                    .filter { it.renderer is LambdaHintsRenderer }
+                    .forEach { it.dispose() }
+                editor.inlayModel.addAfterLineEndElement(p.offset, p.relatesToPrecedingText, LambdaHintsRenderer(p.presentation))
             }
         }
     }
@@ -73,4 +82,29 @@ class KotlinLambdasHintsProvider : KotlinAbstractHintsProvider<KotlinLambdasHint
             }
         }
     """.trimIndent()
+
+
+    /**
+     * This renderer is not more than just a filter criterion. [PresentationRenderer] is not extensible, instead delegation is used.
+     */
+    private class LambdaHintsRenderer(presentation: InlayPresentation) : EditorCustomElementRenderer {
+        private val delegate: PresentationRenderer = PresentationRenderer(presentation)
+
+        override fun paint(inlay: Inlay<*>, g: Graphics, targetRegion: Rectangle, textAttributes: TextAttributes) {
+            delegate.paint(inlay, g, targetRegion, textAttributes)
+        }
+
+        override fun calcWidthInPixels(inlay: Inlay<*>): Int {
+            return delegate.calcWidthInPixels(inlay)
+        }
+
+        // this should not be shown anywhere
+        override fun getContextMenuGroupId(inlay: Inlay<*>): String {
+            return "DummyActionGroup"
+        }
+
+        override fun toString(): String {
+            return delegate.toString()
+        }
+    }
 }
\ No newline at end of file