From f19f38a4db4779393ad51a56b8ef3b7e0da5f929 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Wed, 9 Jan 2019 12:40:25 +0300
Subject: [PATCH] Rework "lambda return expression" hint (KT-28870)

Move the hint to the end of the line, to avoid breaking indentation.

Also use `EditorLinePainter` API for painting to prevent problems with
typing at the line end, when caret is placed after the hint
(See IDEA-204702 for implementing such hints in platform).

 #KT-28870 Fixed
---
 idea/resources/META-INF/plugin-common.xml     |  11 ++
 .../KotlinInlayParameterHintsProvider.kt      |  21 ++-
 .../DisableReturnLambdaHintOptionAction.kt    |  50 +++++
 .../custom/KotlinCodeHintsModel.kt            | 120 ++++++++++++
 .../custom/KotlinCodeHintsPass.kt             | 103 +++++++++++
 .../custom/ReturnHintLinePainter.kt           |  39 ++++
 .../org/jetbrains/kotlin/idea/util/editors.kt |  27 +++
 .../LambdaReturnValueHintsTest.kt             | 174 +++++++++++++++---
 8 files changed, 514 insertions(+), 31 deletions(-)
 create mode 100644 idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/DisableReturnLambdaHintOptionAction.kt
 create mode 100644 idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsModel.kt
 create mode 100644 idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsPass.kt
 create mode 100644 idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/ReturnHintLinePainter.kt
 create mode 100644 idea/src/org/jetbrains/kotlin/idea/util/editors.kt

diff --git a/idea/resources/META-INF/plugin-common.xml b/idea/resources/META-INF/plugin-common.xml
index 12690a1199e..508810dcd94 100644
--- a/idea/resources/META-INF/plugin-common.xml
+++ b/idea/resources/META-INF/plugin-common.xml
@@ -3,6 +3,9 @@
     <component>
       <implementation-class>org.jetbrains.kotlin.idea.highlighter.KotlinBeforeResolveHighlightingPass$Factory</implementation-class>
     </component>
+    <component>
+      <implementation-class>org.jetbrains.kotlin.idea.parameterInfo.custom.KotlinCodeHintsPass$Companion$Factory</implementation-class>
+    </component>
     <component>
       <implementation-class>org.jetbrains.kotlin.idea.refactoring.cutPaste.MoveDeclarationsPassFactory</implementation-class>
       <skipForDefaultProject/>
@@ -26,6 +29,9 @@
     <component>
       <implementation-class>org.jetbrains.kotlin.idea.configuration.KotlinMigrationProjectComponent</implementation-class>
     </component>
+    <component>
+      <implementation-class>org.jetbrains.kotlin.idea.parameterInfo.custom.KotlinCodeHintsModel</implementation-class>
+    </component>
   </project-components>
 
   <application-components>
@@ -378,6 +384,7 @@
     <lang.psiStructureViewFactory language="kotlin" implementationClass="org.jetbrains.kotlin.idea.structureView.KotlinStructureViewFactory"/>
     <structureViewBuilder order="first" key="CLASS" factoryClass="org.jetbrains.kotlin.idea.structureView.KtClsStructureViewBuilderProvider"/>
     <editorNotificationProvider implementation="org.jetbrains.kotlin.idea.script.configutation.MultipleScriptDefinitionsChecker"/>
+    <editor.linePainter order="first" implementation="org.jetbrains.kotlin.idea.parameterInfo.custom.ReturnHintLinePainter"/>
 
     <lang.foldingBuilder language="kotlin" implementationClass="org.jetbrains.kotlin.idea.KotlinFoldingBuilder"/>
     <lang.formatter language="kotlin" implementationClass="org.jetbrains.kotlin.idea.formatter.KotlinFormattingModelBuilder"/>
@@ -1694,6 +1701,10 @@
       <category>Kotlin</category>
     </intentionAction>
 
+    <intentionAction>
+      <className>org.jetbrains.kotlin.idea.parameterInfo.custom.DisableReturnLambdaHintOptionAction</className>
+    </intentionAction>
+
     <intentionAction>
       <className>org.jetbrains.kotlin.idea.intentions.AddUnderscoresToNumericLiteralIntention</className>
       <category>Kotlin</category>
diff --git a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/KotlinInlayParameterHintsProvider.kt b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/KotlinInlayParameterHintsProvider.kt
index 328745c47d9..3a7540d4674 100644
--- a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/KotlinInlayParameterHintsProvider.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/KotlinInlayParameterHintsProvider.kt
@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameOrNull
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 
-enum class HintType(desc: String, enabled: Boolean) {
+enum class HintType(val desc: String, defaultEnabled: Boolean) {
 
     PROPERTY_HINT("Show property type hints", false) {
         override fun provideHints(elem: PsiElement): List<InlayInfo> {
@@ -86,9 +86,11 @@ enum class HintType(desc: String, enabled: Boolean) {
             elem is KtExpression && elem !is KtFunctionLiteral && !elem.isNameReferenceInCall()
 
         override fun provideHints(elem: PsiElement): List<InlayInfo> {
-            if (elem is KtExpression) {
-                return provideLambdaReturnValueHints(elem)
-            }
+            // Will be painted with ReturnHintLinePainter
+
+            // Enable/Disable setting will be present in the list with other hints.
+            // Enable action will be provided by the platform.
+            // Disable action need to be reimplemented as hints are not actually added, see DisableReturnLambdaHintOptionAction.
 
             return emptyList()
         }
@@ -133,11 +135,12 @@ enum class HintType(desc: String, enabled: Boolean) {
 
     abstract fun isApplicable(elem: PsiElement): Boolean
     abstract fun provideHints(elem: PsiElement): List<InlayInfo>
-    val option = Option("SHOW_${this.name}", desc, enabled)
+    val option = Option("SHOW_${this.name}", desc, defaultEnabled)
     val enabled
         get() = option.get()
 }
 
+@Suppress("UnstableApiUsage")
 class KotlinInlayParameterHintsProvider : InlayParameterHintsProvider {
 
     override fun getSupportedOptions(): List<Option> = HintType.values().map { it.option }
@@ -160,8 +163,10 @@ class KotlinInlayParameterHintsProvider : InlayParameterHintsProvider {
         }
     }
 
-    override fun getParameterHints(element: PsiElement?): List<InlayInfo> =
-        HintType.resolveToEnabled(element)?.provideHints(element!!) ?: emptyList()
+    override fun getParameterHints(element: PsiElement?): List<InlayInfo> {
+        val resolveToEnabled = HintType.resolveToEnabled(element) ?: return emptyList()
+        return resolveToEnabled.provideHints(element!!)
+    }
 
     override fun getBlackListDependencyLanguage(): Language = JavaLanguage.INSTANCE
 
@@ -188,5 +193,5 @@ class KotlinInlayParameterHintsProvider : InlayParameterHintsProvider {
     }
 }
 
-private fun PsiElement.isNameReferenceInCall() =
+fun PsiElement.isNameReferenceInCall() =
     this is KtNameReferenceExpression && parent is KtCallExpression
diff --git a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/DisableReturnLambdaHintOptionAction.kt b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/DisableReturnLambdaHintOptionAction.kt
new file mode 100644
index 00000000000..2194b070055
--- /dev/null
+++ b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/DisableReturnLambdaHintOptionAction.kt
@@ -0,0 +1,50 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.parameterInfo.custom
+
+import com.intellij.codeInsight.CodeInsightBundle
+import com.intellij.codeInsight.hints.InlayParameterHintsExtension
+import com.intellij.codeInsight.intention.IntentionAction
+import com.intellij.codeInsight.intention.LowPriorityAction
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.ex.EditorSettingsExternalizable
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.idea.parameterInfo.HintType
+import org.jetbrains.kotlin.idea.util.refreshAllOpenEditors
+
+class DisableReturnLambdaHintOptionAction : IntentionAction, LowPriorityAction {
+    override fun getText(): String {
+        val optionName = HintType.LAMBDA_RETURN_EXPRESSION.option.name
+        return if (optionName.startsWith("show", ignoreCase = true)) {
+            "Do not ${optionName.toLowerCase()}"
+        } else {
+            CodeInsightBundle.message("inlay.hints.disable.custom.option", optionName)
+        }
+    }
+
+    override fun getFamilyName(): String = text
+
+    override fun isAvailable(project: Project, editor: Editor, file: PsiFile): Boolean {
+        if (file.language != KotlinLanguage.INSTANCE) return false
+
+        InlayParameterHintsExtension.forLanguage(file.language) ?: return false
+
+        if (!EditorSettingsExternalizable.getInstance().isShowParameterNameHints || !HintType.LAMBDA_RETURN_EXPRESSION.enabled) {
+            return false
+        }
+
+        return KotlinCodeHintsModel.getInstance(project).getExtensionInfoAtOffset(editor) != null
+    }
+
+    override fun invoke(project: Project, editor: Editor, file: PsiFile) {
+        HintType.LAMBDA_RETURN_EXPRESSION.option.set(false)
+        refreshAllOpenEditors()
+    }
+
+    override fun startInWriteAction(): Boolean = false
+}
diff --git a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsModel.kt b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsModel.kt
new file mode 100644
index 00000000000..3987f3de725
--- /dev/null
+++ b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsModel.kt
@@ -0,0 +1,120 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license 
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.parameterInfo.custom
+
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.EditorFactory
+import com.intellij.openapi.editor.RangeMarker
+import com.intellij.openapi.editor.event.EditorFactoryEvent
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.util.containers.ContainerUtil
+import org.jetbrains.kotlin.idea.conversion.copy.range
+import org.jetbrains.kotlin.idea.util.application.runReadAction
+import org.jetbrains.kotlin.idea.util.compat.EditorFactoryListenerWrapper
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import java.util.concurrent.ConcurrentHashMap
+
+class KotlinCodeHintsModel(val project: Project) : EditorFactoryListenerWrapper {
+    companion object {
+        fun getInstance(project: Project): KotlinCodeHintsModel =
+            project.getComponent(KotlinCodeHintsModel::class.java) ?: error("Component `KotlinCodeHintsModel` is expected to be registered")
+    }
+
+    private class DocumentExtensionInfoModel(val document: Document) {
+        private val lineEndMarkers = ConcurrentHashMap<RangeMarker, String>()
+
+        fun markEndOfLine(lineEndOffset: Int, hint: String) {
+            val endLineMarker = document.createRangeMarker(lineEndOffset, lineEndOffset)
+            endLineMarker.isGreedyToRight = true
+
+            lineEndMarkers[endLineMarker] = hint
+        }
+
+        fun getExtensionAtOffset(offset: Int): String? {
+            return runReadAction {
+                // Protect operations working with the document offsets
+
+                lineEndMarkers
+                    .entries
+                    .firstOrNull { (marker, _) ->
+                        val textRange = marker.range
+                        if (textRange == null || !(textRange.startOffset <= offset && offset <= textRange.endOffset)) {
+                            return@firstOrNull false
+                        }
+
+                        val document = marker.document
+                        val hasNewLine = document.getText(textRange).contains('\n')
+                        if (!hasNewLine) {
+                            textRange.endOffset == offset
+                        } else {
+                            // New line may appear after session of fast typing with one or several enter hitting.
+                            // We can't believe startOffset too because typing session may had started with
+                            // typing adding several chars at the original line.
+                            val originalLineNumber = document.getLineNumber(textRange.startOffset)
+                            val currentOriginalLineEnd = document.getLineEndOffset(originalLineNumber)
+
+                            currentOriginalLineEnd == offset
+                        }
+                    }
+                    ?.value
+            }
+        }
+
+        fun dispose() {
+            for (marker in lineEndMarkers.keys()) {
+                marker.dispose()
+            }
+        }
+    }
+
+    private val documentModels =
+        ContainerUtil.createConcurrentSoftMap<Document, DocumentExtensionInfoModel>()
+
+    init {
+        val editorFactory = EditorFactory.getInstance()
+        editorFactory.addEditorFactoryListener(this, project)
+    }
+
+    override fun editorReleased(event: EditorFactoryEvent) {
+        // Pass for other editor with the same document if present should re-add needed hints
+        removeAll(event.editor.document)
+    }
+
+    fun getExtensionInfoAtOffset(editor: Editor): String? {
+        return getExtensionInfo(editor.document, editor.caretModel.offset)
+    }
+
+    fun getExtensionInfo(document: Document, offset: Int): String? {
+        return documentModels[document]?.getExtensionAtOffset(offset)
+    }
+
+    fun removeAll(document: Document) {
+        documentModels.remove(document)?.dispose()
+    }
+
+    fun update(document: Document, actualHints: Map<PsiElement, String>) {
+        if (actualHints.isEmpty()) {
+            removeAll(document)
+            return
+        }
+
+        val updatedModel = runReadAction {
+            val model = DocumentExtensionInfoModel(document)
+            for ((element, hint) in actualHints) {
+                val lineNumber = document.getLineNumber(element.endOffset)
+                val lineEndOffset = document.getLineEndOffset(lineNumber)
+
+                model.markEndOfLine(lineEndOffset, hint)
+            }
+
+            model
+        }
+
+        documentModels.put(document, updatedModel)?.dispose()
+    }
+}
\ No newline at end of file
diff --git a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsPass.kt b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsPass.kt
new file mode 100644
index 00000000000..f7ade739abf
--- /dev/null
+++ b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/KotlinCodeHintsPass.kt
@@ -0,0 +1,103 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.parameterInfo.custom
+
+import com.intellij.codeHighlighting.EditorBoundHighlightingPass
+import com.intellij.codeHighlighting.TextEditorHighlightingPass
+import com.intellij.codeHighlighting.TextEditorHighlightingPassFactory
+import com.intellij.codeHighlighting.TextEditorHighlightingPassRegistrar
+import com.intellij.codeInsight.hints.InlayParameterHintsExtension
+import com.intellij.diff.util.DiffUtil
+import com.intellij.openapi.components.ProjectComponent
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.ex.EditorSettingsExternalizable
+import com.intellij.openapi.progress.ProgressIndicator
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.SyntaxTraverser
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.idea.parameterInfo.HintType
+import org.jetbrains.kotlin.idea.parameterInfo.TYPE_INFO_PREFIX
+import org.jetbrains.kotlin.idea.parameterInfo.provideLambdaReturnValueHints
+import org.jetbrains.kotlin.psi.KtExpression
+import java.util.*
+
+class KotlinCodeHintsPass(private val myRootElement: PsiElement, editor: Editor) :
+    EditorBoundHighlightingPass(editor, myRootElement.containingFile, true) {
+
+    private val myTraverser: SyntaxTraverser<PsiElement> = SyntaxTraverser.psiTraverser(myRootElement)
+
+    override fun doCollectInformation(progress: ProgressIndicator) {
+        if (myFile.language != KotlinLanguage.INSTANCE) return
+        if (myDocument == null) return
+
+        val kotlinCodeHintsModel = KotlinCodeHintsModel.getInstance(myRootElement.project)
+
+        val provider = InlayParameterHintsExtension.forLanguage(KotlinLanguage.INSTANCE)
+        if (provider == null || !provider.canShowHintsWhenDisabled() && !isEnabled || DiffUtil.isDiffEditor(myEditor)) {
+            kotlinCodeHintsModel.removeAll(myDocument)
+            return
+        }
+
+        if (HintType.LAMBDA_RETURN_EXPRESSION.enabled) {
+            val actualHints = HashMap<PsiElement, String>()
+            myTraverser.forEach { element -> processLambdaReturnHints(element, actualHints) }
+
+            kotlinCodeHintsModel.update(myDocument, actualHints)
+        } else {
+            kotlinCodeHintsModel.removeAll(myDocument)
+        }
+    }
+
+    private fun processLambdaReturnHints(element: PsiElement, actualElements: MutableMap<PsiElement, String>) {
+        if (element !is KtExpression) return
+
+        for (returnHint in provideLambdaReturnValueHints(element)) {
+            val offset = returnHint.offset
+
+            if (!canShowHintsAtOffset(offset)) continue
+
+            actualElements[element] = returnHint.text.substringAfter(TYPE_INFO_PREFIX)
+        }
+    }
+
+    override fun doApplyInformationToEditor() {
+        // Information will be painted with org.jetbrains.kotlin.idea.parameterInfo.custom.ReturnHintLinePainter
+    }
+
+    /**
+     * Adding hints on the borders of root element (at startOffset or endOffset)
+     * is allowed only in the case when root element is a document
+     *
+     * @return true iff a given offset can be used for hint rendering
+     */
+    private fun canShowHintsAtOffset(offset: Int): Boolean {
+        val rootRange = myRootElement.textRange
+
+        if (rootRange.startOffset < offset && offset < rootRange.endOffset) {
+            return true
+        }
+
+        return myDocument != null && myDocument.textLength == rootRange.length
+    }
+
+    companion object {
+        class Factory(registrar: TextEditorHighlightingPassRegistrar) : ProjectComponent, TextEditorHighlightingPassFactory {
+            init {
+                registrar.registerTextEditorHighlightingPass(this, null, null, false, -1)
+            }
+
+            override fun createHighlightingPass(file: PsiFile, editor: Editor): TextEditorHighlightingPass? {
+                if (file.language != KotlinLanguage.INSTANCE) return null
+                return KotlinCodeHintsPass(file, editor)
+            }
+        }
+
+        private val isEnabled: Boolean
+            get() =
+                EditorSettingsExternalizable.getInstance().isShowParameterNameHints
+    }
+}
\ No newline at end of file
diff --git a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/ReturnHintLinePainter.kt b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/ReturnHintLinePainter.kt
new file mode 100644
index 00000000000..f03329b298c
--- /dev/null
+++ b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/custom/ReturnHintLinePainter.kt
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.parameterInfo.custom
+
+import com.intellij.openapi.editor.*
+import com.intellij.openapi.editor.markup.TextAttributes
+import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiManager
+import org.jetbrains.kotlin.idea.KotlinLanguage
+
+class ReturnHintLinePainter : EditorLinePainter() {
+    override fun getLineExtensions(project: Project, file: VirtualFile, lineNumber: Int): List<LineExtensionInfo>? {
+        val psiFile = PsiManager.getInstance(project).findFile(file) ?: return null
+        if (psiFile.language != KotlinLanguage.INSTANCE) {
+            return null
+        }
+
+        val hint = getLineHint(project, file, lineNumber)
+            ?: return null
+
+        val hintLineInfo = LineExtensionInfo(" $hint", TextAttributes())
+        return listOf(hintLineInfo)
+    }
+
+    private fun getLineHint(project: Project, file: VirtualFile, lineNumber: Int): String? {
+        val doc = FileDocumentManager.getInstance().getDocument(file) ?: return null
+        if (lineNumber >= doc.lineCount) {
+            return null
+        }
+        val lineEndOffset = doc.getLineEndOffset(lineNumber)
+
+        return KotlinCodeHintsModel.getInstance(project).getExtensionInfo(doc, lineEndOffset)
+    }
+}
\ No newline at end of file
diff --git a/idea/src/org/jetbrains/kotlin/idea/util/editors.kt b/idea/src/org/jetbrains/kotlin/idea/util/editors.kt
new file mode 100644
index 00000000000..3d8760549be
--- /dev/null
+++ b/idea/src/org/jetbrains/kotlin/idea/util/editors.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.util
+
+import com.intellij.codeInsight.daemon.DaemonCodeAnalyzer
+import com.intellij.codeInsight.hints.ParameterHintsPassFactory
+import com.intellij.openapi.fileEditor.FileEditorManager
+import com.intellij.openapi.project.ProjectManager
+import com.intellij.psi.PsiManager
+
+fun refreshAllOpenEditors() {
+    ParameterHintsPassFactory.forceHintsUpdateOnNextPass();
+    ProjectManager.getInstance().openProjects.forEach { project ->
+        val psiManager = PsiManager.getInstance(project)
+        val daemonCodeAnalyzer = DaemonCodeAnalyzer.getInstance(project)
+        val fileEditorManager = FileEditorManager.getInstance(project)
+
+        DaemonCodeAnalyzer.getInstance(project).restart()
+
+        fileEditorManager.selectedFiles.forEach {
+            psiManager.findFile(it)?.let { daemonCodeAnalyzer.restart(it) }
+        }
+    }
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt b/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
index e6b2a9d5478..12bbaad1b07 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
@@ -5,20 +5,148 @@
 
 package org.jetbrains.kotlin.idea.parameterInfo
 
+import com.intellij.codeInsight.intention.IntentionAction
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.LineExtensionInfo
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.testFramework.utils.inlays.InlayHintsChecker
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
+import org.jetbrains.kotlin.test.TagsTestDataUtil
+import org.junit.Assert
 
 class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
+    companion object {
+        const val DISABLE_ACTION_TEXT = "Do not show lambda return expression hints"
+        const val ENABLE_ACTION_TEXT = "Show lambda return expression hints"
+    }
+
     override fun getProjectDescriptor(): KotlinLightProjectDescriptor = KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE
 
+    private class LineExtensionInfoTag(offset: Int, data: LineExtensionInfo) :
+        TagsTestDataUtil.TagInfo<LineExtensionInfo>(offset, true, true, false, data) {
+
+        override fun getName() = "hint"
+        override fun getAttributesString(): String = "text=\"${data.text}\""
+    }
+
+    private class CaretTag(editor: Editor) :
+        TagsTestDataUtil.TagInfo<Any>(
+            editor.caretModel.currentCaret.offset,
+            /*isStart = */true, /*isClosed = */false, /*isFixed = */true,
+            "caret"
+        )
+
+
+    private fun collectActualLineExtensionsTags(): List<LineExtensionInfoTag> {
+        val tags = ArrayList<LineExtensionInfoTag>()
+        val lineCount = myFixture.editor.document.lineCount
+        for (i in 0 until lineCount) {
+            val lineEndOffset = myFixture.editor.document.getLineEndOffset(i)
+
+            (myFixture.editor as EditorImpl).processLineExtensions(i) { lineExtensionInfo ->
+                tags.add(LineExtensionInfoTag(lineEndOffset, lineExtensionInfo))
+                true
+            }
+        }
+
+        return tags
+    }
+
     fun check(text: String) {
         myFixture.configureByText("A.kt", text.trimIndent())
-        myFixture.testInlays()
 
-        if (myFixture.editor.caretModel.offset != 0) {
-            myFixture.checkHintType(HintType.LAMBDA_RETURN_EXPRESSION)
+        val expectedText = run {
+            val tags = if (editor.caretModel.offset > 0) listOf(CaretTag(editor)) else emptyList()
+            TagsTestDataUtil.insertTagsInText(tags, editor.document.text)
         }
+
+        // Clean test file from the hints tags
+        InlayHintsChecker(myFixture).extractInlaysAndCaretInfo(editor.document)
+
+        myFixture.doHighlighting()
+
+        Assert.assertTrue(
+            "No other inlays should be present in the file",
+            editor.inlayModel.getInlineElementsInRange(0, editor.document.textLength).isEmpty()
+        )
+
+        if (editor.caretModel.offset > 0) {
+            val availableIntentions = myFixture.availableIntentions
+            Assert.assertTrue(
+                "Disable action with text `$DISABLE_ACTION_TEXT` is expected: \n${availableIntentions.joinToString(separator = "\n") { "  $it" }}",
+                availableIntentions.any { it.text == DISABLE_ACTION_TEXT }
+            )
+        }
+
+        val actualText = run {
+            val tags = ArrayList<TagsTestDataUtil.TagInfo<*>>()
+
+            tags.addAll(collectActualLineExtensionsTags())
+
+            if (editor.caretModel.offset > 0) {
+                tags.add(CaretTag(editor))
+            }
+
+            TagsTestDataUtil.insertTagsInText(tags, editor.document.text)
+        }
+
+        Assert.assertEquals(expectedText, actualText)
+    }
+
+    fun testDisableEnableActions() {
+        myFixture.configureByText(
+            "A.kt",
+            """
+            val x = run {
+                println("foo")
+                1<caret>
+            }
+            """.trimIndent()
+        )
+
+        Assert.assertTrue("Return expression hint should be enabled", HintType.LAMBDA_RETURN_EXPRESSION.enabled)
+        try {
+            val disableIntention = findDisableReturnHintsIntention()
+            disableIntention.invoke(project, editor, file)
+
+            Assert.assertFalse("Disable action doesn't work", HintType.LAMBDA_RETURN_EXPRESSION.option.get())
+
+            val availableIntentions = myFixture.availableIntentions
+            Assert.assertTrue(
+                intentionsPresenceErrorMessage(
+                    "Disable action shouldn't be present when option is already disabled", availableIntentions
+                ),
+                availableIntentions.find { it.text == DISABLE_ACTION_TEXT } == null
+            )
+
+            val enableAction = availableIntentions.find { it.text == ENABLE_ACTION_TEXT }
+            Assert.assertTrue(
+                intentionsPresenceErrorMessage("No enable action with text $ENABLE_ACTION_TEXT found", availableIntentions),
+                enableAction != null
+            )
+
+            enableAction!!.invoke(project, editor, file)
+            Assert.assertTrue("Enable action doesn't work", HintType.LAMBDA_RETURN_EXPRESSION.option.get())
+
+        } finally {
+            HintType.LAMBDA_RETURN_EXPRESSION.option.set(true)
+        }
+
+    }
+
+    private fun findDisableReturnHintsIntention(): IntentionAction {
+        val availableIntentions = myFixture.availableIntentions
+        return availableIntentions.find { it.text == DISABLE_ACTION_TEXT }
+            ?: throw AssertionError(
+                intentionsPresenceErrorMessage("Disable action with text `$DISABLE_ACTION_TEXT` is expected", availableIntentions)
+            )
+    }
+
+    private fun intentionsPresenceErrorMessage(message: String, intentions: List<IntentionAction>): String {
+        return "$message: \n" +
+                intentions.joinToString(separator = "\n") { "  $it" }
     }
 
     fun testSimple() {
@@ -26,7 +154,7 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
             val x = run {
                 println("foo")
-                <caret><hint text="^run" />1
+                1<caret><hint text=" ^run"/>
             }
             """
         )
@@ -37,7 +165,7 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
             val x = run {
                 var s = "abc"
-                <hint text="^run" />s.length
+                s.length<caret><hint text=" ^run"/>
             }
             """
         )
@@ -48,9 +176,9 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
             val x = run {
                 if (true) {
-                    <hint text="^run" />1
+                    1<hint text=" ^run"/>
                 } else {
-                    <hint text="^run" />0
+                    0<hint text=" ^run"/>
                 }
             }
             """
@@ -62,7 +190,7 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
             val x = run {
                 println(1)
-                <caret><hint text="^run"/>if (true) 1 else { 0 }
+                if (true) 1 else { 0 }<caret><hint text=" ^run"/>
             }
             """
         )
@@ -73,8 +201,8 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
             val x = run {
                 when (true) {
-                    true -> <hint text="^run" />1
-                    false -><hint text="^run" />0
+                    true -> 1<hint text=" ^run"/>
+                    false ->0<hint text=" ^run"/>
                 }
             }
             """
@@ -96,7 +224,7 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
             val x = run foo@{
                 println("foo")
-                <hint text="^foo" />1
+                1<hint text=" ^foo"/>
             }
             """
         )
@@ -109,12 +237,12 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
                 if (true) {
                 }
 
-                <hint text="^hello" />run { // Two hints here
+                run { // Two hints here
                     when (true) {
-                        true -> <hint text="^run" />1
-                        false -> <hint text="^run" />0
+                        true -> 1<hint text=" ^run"/>
+                        false -> 0<hint text=" ^run"/>
                     }
-                }
+                }<hint text=" ^hello"/>
             }
             """
         )
@@ -126,7 +254,7 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             fun foo() {
                 run {
                     val length: Int? = null
-                    <hint text="^run" />length ?: 0
+                    length ?: 0<hint text=" ^run"/>
                 }
             }
             """
@@ -140,12 +268,12 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
                 var test = 0
                 run {
                     test
-                    <hint text="^run"/>test++
+                    test++<hint text=" ^run"/>
                 }
 
                 run {
                     test
-                    <hint text="^run"/>++test
+                    ++test<hint text=" ^run"/>
                 }
             }
             """
@@ -162,12 +290,12 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
                 run {
                     val files: Any? = null
                     @Some
-                    <hint text="^run"/>12
+                    12<hint text=" ^run"/>
                 }
 
                 run {
                     val files: Any? = null
-                    <hint text="^run"/>@Some 12
+                    @Some 12<hint text=" ^run"/>
                 }
             }
             """
@@ -181,12 +309,12 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
                 run {
                     val files: Any? = null
                     run@
-                    <hint text="^run"/>12
+                    12<hint text=" ^run"/>
                 }
 
                 run {
                     val files: Any? = null
-                    <hint text="^run"/>run@12
+                    run@12<hint text=" ^run"/>
                 }
             }
             """
@@ -198,7 +326,7 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
             fun test() = run {
                 val a = 1
-                <hint text="^run"/>{ a }
+                { a }<hint text=" ^run"/>
             }
             """
         )