From 55038e19ec696df573e2c5a30ae85b6ba5e53767 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Wed, 7 Oct 2020 00:51:09 +0900
Subject: [PATCH] Minor: Extract lambda return expression handling

---
 .../idea/quickfix/ChangeToLabeledReturnFix.kt | 32 +++++++++++--------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/ChangeToLabeledReturnFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/ChangeToLabeledReturnFix.kt
index 93ca545289f..d5b25c3cbf0 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/ChangeToLabeledReturnFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/ChangeToLabeledReturnFix.kt
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.quickfix
 import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.KotlinBundle
@@ -62,23 +63,28 @@ class ChangeToLabeledReturnFix(
         override fun doCreateActions(diagnostic: Diagnostic): List<IntentionAction> {
             val element = diagnostic.psiElement as? KtElement ?: return emptyList()
             val context by lazy { element.analyze() }
+
             val returnExpression = when (diagnostic.factory) {
-                Errors.RETURN_NOT_ALLOWED -> {
+                Errors.RETURN_NOT_ALLOWED ->
                     diagnostic.psiElement as? KtReturnExpression
-                }
-                Errors.TYPE_MISMATCH, Errors.CONSTANT_EXPECTED_TYPE_MISMATCH, Errors.NULL_FOR_NONNULL_TYPE -> {
-                    val returnExpression = diagnostic.psiElement.getStrictParentOfType<KtReturnExpression>() ?: return emptyList()
-                    val lambda = returnExpression.getStrictParentOfType<KtLambdaExpression>() ?: return emptyList()
-                    val lambdaReturnType = context[BindingContext.FUNCTION, lambda.functionLiteral]?.returnType ?: return emptyList()
-                    val returnType = returnExpression.returnedExpression?.getType(context) ?: return emptyList()
-                    if (!returnType.isSubtypeOf(lambdaReturnType)) return emptyList()
-                    returnExpression
-                }
+                Errors.TYPE_MISMATCH,
+                Errors.CONSTANT_EXPECTED_TYPE_MISMATCH,
+                Errors.NULL_FOR_NONNULL_TYPE ->
+                    getLambdaReturnExpression(diagnostic.psiElement, context)
                 else -> null
             } ?: return emptyList()
-            return findAccessibleLabels(context, returnExpression).map {
-                ChangeToLabeledReturnFix(returnExpression, labeledReturn = "return@${it.render()}")
-            }
+
+            val candidates = findAccessibleLabels(context, returnExpression)
+            return candidates.map { ChangeToLabeledReturnFix(returnExpression, labeledReturn = "return@${it.render()}") }
+        }
+
+        private fun getLambdaReturnExpression(element: PsiElement, bindingContext: BindingContext): KtReturnExpression? {
+            val returnExpression = element.getStrictParentOfType<KtReturnExpression>() ?: return null
+            val lambda = returnExpression.getStrictParentOfType<KtLambdaExpression>() ?: return null
+            val lambdaReturnType = bindingContext[BindingContext.FUNCTION, lambda.functionLiteral]?.returnType ?: return null
+            val returnType = returnExpression.returnedExpression?.getType(bindingContext) ?: return null
+            if (!returnType.isSubtypeOf(lambdaReturnType)) return null
+            return returnExpression
         }
     }
 }
\ No newline at end of file