Fix lambda return hint for labeled expressions (KT-24828)

#KT-24828 Fixed

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")
 }

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
+                }
+            }
+            """
+        )
+    }
 }