From cd54c7922d764c4e9dfc9b79dce29317ebe6d5f2 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Thu, 7 Jun 2018 18:54:13 +0300
Subject: [PATCH] Fix lambda return hint for labeled expressions (KT-24828)

 #KT-24828 Fixed
---
 .../parameterInfo/LambdaReturnValueHints.kt   | 25 ++++++++++++++-----
 .../LambdaReturnValueHintsTest.kt             | 19 ++++++++++++++
 2 files changed, 38 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 2be89662300..6ffe3639249 100644
--- a/idea/src/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHints.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHints.kt
@@ -55,16 +55,29 @@ private fun getNameOfFunctionThatTakesLambda(expression: KtExpression): String?
 }
 
 private fun allowLabelOnExpressionPart(expression: KtExpression): Boolean {
-    val parent = expression.parent
-    return parent is KtAnnotatedExpression && parent.baseExpression == expression && parent.lineSeparatorBeforeBase()
+    val parent = expression.parent as? KtExpression ?: return false
+    return expression == expressionStatementPart(parent)
 }
 
 private fun forceLabelOnExpressionPart(expression: KtExpression): Boolean {
-    return expression is KtAnnotatedExpression && expression.lineSeparatorBeforeBase()
+    return expressionStatementPart(expression) != null
 }
 
-private fun KtAnnotatedExpression.lineSeparatorBeforeBase(): Boolean {
-    val base = baseExpression ?: return false
-    val whiteSpace = base.node.treePrev?.psi as? PsiWhiteSpace ?: return false
+private fun expressionStatementPart(expression: KtExpression): KtExpression? {
+    val splitPart: KtExpression = when (expression) {
+        is KtAnnotatedExpression -> expression.baseExpression
+        is KtLabeledExpression -> expression.baseExpression
+        else -> null
+    } ?: return null
+
+    if (!isNewLineBeforeExpression(splitPart)) {
+        return null
+    }
+
+    return splitPart
+}
+
+private fun isNewLineBeforeExpression(expression: KtExpression): Boolean {
+    val whiteSpace = expression.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 7cd9b96cc7e..5d1daa643b0 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/parameterInfo/LambdaReturnValueHintsTest.kt
@@ -158,4 +158,23 @@ class LambdaReturnValueHintsTest : KotlinLightCodeInsightFixtureTestCase() {
             """
         )
     }
+
+    fun testLabeledStatement() {
+        check(
+            """
+            fun test() {
+                run {
+                    val files: Any? = null
+                    run@
+                    <hint text="^run"/>12
+                }
+
+                run {
+                    val files: Any? = null
+                    <hint text="^run"/>run@12
+                }
+            }
+            """
+        )
+    }
 }