Allow return value hint for lambda expressions (KT-26689)

#KT-26689 Fixed

idea/src/org/jetbrains/kotlin/idea/parameterInfo/KotlinInlayParameterHintsProvider.kt

@@ -81,30 +81,30 @@ enum class HintType(desc: String, enabled: Boolean) {
         override fun isApplicable(elem: PsiElement): Boolean = elem is KtValueArgumentList
     },
 
-
     LAMBDA_RETURN_EXPRESSION("Show lambda return expression hints", true) {
         override fun isApplicable(elem: PsiElement) =
-            elem is KtExpression && elem !is KtLambdaExpression && elem !is KtFunctionLiteral &&
-                    !elem.isNameReferenceInCall()
+            elem is KtExpression && elem !is KtFunctionLiteral && !elem.isNameReferenceInCall()
 
         override fun provideHints(elem: PsiElement): List<InlayInfo> {
             if (elem is KtExpression) {
                 return provideLambdaReturnValueHints(elem)
             }
+
             return emptyList()
         }
     },
 
     LAMBDA_IMPLICIT_PARAMETER_RECEIVER("Show hints for implicit receivers and parameters of lambdas", true) {
-        override fun isApplicable(elem: PsiElement) = elem is KtLambdaExpression
+        override fun isApplicable(elem: PsiElement) = elem is KtFunctionLiteral
 
         override fun provideHints(elem: PsiElement): List<InlayInfo> {
-            (elem as? KtLambdaExpression)?.let {
-                return provideLambdaImplicitHints(elem)
+            ((elem as? KtFunctionLiteral)?.parent as? KtLambdaExpression)?.let {
+                return provideLambdaImplicitHints(it)
             }
             return emptyList()
         }
     },
+
     SUSPENDING_CALL("Show hints for suspending calls", false) {
         override fun isApplicable(elem: PsiElement) = elem.isNameReferenceInCall() && ApplicationManager.getApplication().isInternal
 
@@ -115,8 +115,11 @@ enum class HintType(desc: String, enabled: Boolean) {
     };
 
     companion object {
+        fun resolve(elem: PsiElement): HintType? {
+            val applicableTypes = HintType.values().filter { it.isApplicable(elem) }
+            return applicableTypes.firstOrNull()
+        }
 
-        fun resolve(elem: PsiElement): HintType? = HintType.values().find { it.isApplicable(elem) }
         fun resolveToEnabled(elem: PsiElement?): HintType? {
 
             val resolved = elem?.let { resolve(it) } ?: return null

idea/tests/org/jetbrains/kotlin/idea/parameterInfo/InlayHintTypeTestUtil.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2018 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
+
+import com.intellij.codeInsight.hints.HintInfo
+import com.intellij.codeInsight.hints.InlayParameterHintsExtension
+import com.intellij.codeInsight.hints.isOwnsInlayInEditor
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.testFramework.fixtures.JavaCodeInsightTestFixture
+import org.junit.Assert
+
+internal fun JavaCodeInsightTestFixture.checkHintType(text: String, hintType: HintType) {
+    configureByText("A.kt", text.trimIndent())
+    doHighlighting()
+
+    val hintInfo = getHintInfoFromProvider(caretOffset, file, editor)
+    Assert.assertNotNull("No hint available at caret", hintInfo)
+    Assert.assertEquals(hintType.option.name, (hintInfo as HintInfo.OptionInfo).optionName)
+}
+
+// It's crucial for this method to be conformable with IDEA internals.
+// Originally copied from com.intellij.codeInsight.hints.getHintInfoFromProvider()
+private fun getHintInfoFromProvider(offset: Int, file: PsiFile, editor: Editor): HintInfo? {
+    val element = file.findElementAt(offset) ?: return null
+    val provider = InlayParameterHintsExtension.forLanguage(file.language) ?: return null
+
+    val isHintOwnedByElement: (PsiElement) -> Boolean = { e -> provider.getHintInfo(e) != null && e.isOwnsInlayInEditor(editor) }
+    val method = PsiTreeUtil.findFirstParent(element, isHintOwnedByElement) ?: return null
+
+    return provider.getHintInfo(method)
+}

idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaImplicitHintsTest.kt

@@ -5,9 +5,11 @@
 
 package org.jetbrains.kotlin.idea.parameterInfo
 
+import com.intellij.codeInsight.hints.HintInfo
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
+import org.junit.Assert
 
 class LambdaImplicitHintsTest : KotlinLightCodeInsightFixtureTestCase() {
     override fun getProjectDescriptor(): KotlinLightProjectDescriptor = KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE
@@ -17,6 +19,16 @@ class LambdaImplicitHintsTest : KotlinLightCodeInsightFixtureTestCase() {
         myFixture.testInlays()
     }
 
+    fun testHintType() {
+        myFixture.checkHintType(
+            """
+            val x = listOf("").filter { <caret>
+            }
+            """,
+            HintType.LAMBDA_IMPLICIT_PARAMETER_RECEIVER
+        )
+    }
+
     fun testSimpleIt() {
         check(
             """

idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt

@@ -17,6 +17,18 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
         myFixture.testInlays()
     }
 
+    fun testHintType() {
+        myFixture.checkHintType(
+            """
+            val x = run {
+                println("foo")
+                <caret>1
+            }
+            """,
+            HintType.LAMBDA_RETURN_EXPRESSION
+        )
+    }
+
     fun testSimple() {
         check(
             """
@@ -177,4 +189,15 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
         )
     }
+
+    fun testReturnFunctionType() {
+        check(
+            """
+            fun test() = run {
+                val a = 1
+                <hint text="^run"/>{ a }
+            }
+            """
+        )
+    }
 }