From 076870ee05467ea0e4d80674e5780d33941c07b0 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Thu, 7 Jun 2018 20:23:57 +0300
Subject: [PATCH] Fix double return hint in lambda for annotated expression
 (KT-23238)

#KT-23238 Fixed
---
 .../parameterInfo/LambdaReturnValueHints.kt   | 27 ++++++++++++++-----
 .../LambdaReturnValueHintsTest.kt             | 22 +++++++++++++++
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHints.kt b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHints.kt
index ebb7fadc388..2be89662300 100644
--- a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHints.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHints.kt
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.parameterInfo
 
 import com.intellij.codeInsight.hints.InlayInfo
+import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
@@ -18,12 +19,11 @@ fun provideLambdaReturnValueHints(expression: KtExpression): List<InlayInfo> {
         return emptyList()
     }
 
-    val parent = expression.parent
-    if (parent is KtDotQualifiedExpression ||
-        parent is KtSafeQualifiedExpression ||
-        parent is KtBinaryExpression ||
-        parent is KtUnaryExpression
-    ) {
+    if (!KtPsiUtil.isStatement(expression)) {
+        if (!allowLabelOnExpressionPart(expression)) {
+            return emptyList()
+        }
+    } else if (forceLabelOnExpressionPart(expression)) {
         return emptyList()
     }
 
@@ -53,3 +53,18 @@ private fun getNameOfFunctionThatTakesLambda(expression: KtExpression): String?
     }
     return null
 }
+
+private fun allowLabelOnExpressionPart(expression: KtExpression): Boolean {
+    val parent = expression.parent
+    return parent is KtAnnotatedExpression && parent.baseExpression == expression && parent.lineSeparatorBeforeBase()
+}
+
+private fun forceLabelOnExpressionPart(expression: KtExpression): Boolean {
+    return expression is KtAnnotatedExpression && expression.lineSeparatorBeforeBase()
+}
+
+private fun KtAnnotatedExpression.lineSeparatorBeforeBase(): Boolean {
+    val base = baseExpression ?: return false
+    val whiteSpace = base.node.treePrev?.psi as? PsiWhiteSpace ?: return false
+    return whiteSpace.text.contains("\n")
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt b/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
index 35fa4e7484f..7cd9b96cc7e 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
@@ -136,4 +136,26 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
         )
     }
+
+    fun testAnnotatedStatement() {
+        check(
+            """
+            @Target(AnnotationTarget.EXPRESSION)
+            annotation class Some
+
+            fun test() {
+                run {
+                    val files: Any? = null
+                    @Some
+                    <hint text="^run"/>12
+                }
+
+                run {
+                    val files: Any? = null
+                    <hint text="^run"/>@Some 12
+                }
+            }
+            """
+        )
+    }
 }