Minor: Extract lambda return expression handling

This commit is contained in:
Yan Zhulanow
2020-10-07 00:51:09 +09:00
parent f76e98868c
commit 55038e19ec
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.quickfix
import com.intellij.codeInsight.intention.IntentionAction import com.intellij.codeInsight.intention.IntentionAction
import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project import com.intellij.openapi.project.Project
import com.intellij.psi.PsiElement
import org.jetbrains.kotlin.diagnostics.Diagnostic import org.jetbrains.kotlin.diagnostics.Diagnostic
import org.jetbrains.kotlin.diagnostics.Errors import org.jetbrains.kotlin.diagnostics.Errors
import org.jetbrains.kotlin.idea.KotlinBundle import org.jetbrains.kotlin.idea.KotlinBundle
@@ -62,23 +63,28 @@ class ChangeToLabeledReturnFix(
override fun doCreateActions(diagnostic: Diagnostic): List<IntentionAction> { override fun doCreateActions(diagnostic: Diagnostic): List<IntentionAction> {
val element = diagnostic.psiElement as? KtElement ?: return emptyList() val element = diagnostic.psiElement as? KtElement ?: return emptyList()
val context by lazy { element.analyze() } val context by lazy { element.analyze() }
val returnExpression = when (diagnostic.factory) { val returnExpression = when (diagnostic.factory) {
Errors.RETURN_NOT_ALLOWED -> { Errors.RETURN_NOT_ALLOWED ->
diagnostic.psiElement as? KtReturnExpression diagnostic.psiElement as? KtReturnExpression
} Errors.TYPE_MISMATCH,
Errors.TYPE_MISMATCH, Errors.CONSTANT_EXPECTED_TYPE_MISMATCH, Errors.NULL_FOR_NONNULL_TYPE -> { Errors.CONSTANT_EXPECTED_TYPE_MISMATCH,
val returnExpression = diagnostic.psiElement.getStrictParentOfType<KtReturnExpression>() ?: return emptyList() Errors.NULL_FOR_NONNULL_TYPE ->
val lambda = returnExpression.getStrictParentOfType<KtLambdaExpression>() ?: return emptyList() getLambdaReturnExpression(diagnostic.psiElement, context)
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
}
else -> null else -> null
} ?: return emptyList() } ?: 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
} }
} }
} }