Extraction Engine: Fix extraction of labeled and non-local returns

compiler/frontend/src/org/jetbrains/kotlin/cfg/pseudocode/pseudocodeUtil.kt

@@ -16,21 +16,25 @@
 
 package org.jetbrains.kotlin.cfg.pseudocode
 
-import org.jetbrains.kotlin.cfg.pseudocode.instructions.*
-import org.jetbrains.kotlin.cfg.pseudocode.instructions.eval.*
-import org.jetbrains.kotlin.cfg.pseudocode.instructions.jumps.*
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.resolve.bindingContextUtil.*
-import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.resolve.BindingContext
-import java.util.*
-import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
-import org.jetbrains.kotlin.resolve.OverridingUtil
-import org.jetbrains.kotlin.types.TypeUtils
-import org.jetbrains.kotlin.types.JetType
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.Instruction
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.eval.*
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.jumps.ConditionalJumpInstruction
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.jumps.ReturnValueInstruction
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.jumps.ThrowExceptionInstruction
+import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.OverridingUtil
+import org.jetbrains.kotlin.resolve.bindingContextUtil.getTargetFunctionDescriptor
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.VarargValueArgument
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
+import org.jetbrains.kotlin.types.JetType
+import org.jetbrains.kotlin.types.TypeUtils
+import java.util.ArrayList
+import java.util.LinkedHashSet
 
 fun getReceiverTypePredicate(resolvedCall: ResolvedCall<*>, receiverValue: ReceiverValue): TypePredicate? {
     val callableDescriptor = resolvedCall.getResultingDescriptor()
@@ -176,3 +180,10 @@ fun Pseudocode.getElementValuesRecursively(element: JetElement): List<PseudoValu
     collectValues()
     return results
 }
+
+public fun Pseudocode.getPseudocodeByElement(element: JetElement): Pseudocode? {
+    if (getCorrespondingElement() == element) return this
+
+    getLocalDeclarations().forEach { decl -> decl.body.getPseudocodeByElement(element)?.let { return it } }
+    return null
+}

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/ExtractionData.kt

@@ -100,7 +100,8 @@ data class ExtractionData(
         getCodeFragmentTextRange()?.let { originalFile.getText()?.substring(it.getStartOffset(), it.getEndOffset()) } ?: ""
     }
 
-    val originalStartOffset = originalElements.firstOrNull()?.let { e -> e.getTextRange()!!.getStartOffset() }
+    val originalStartOffset: Int?
+        get() = originalElements.firstOrNull()?.let { e -> e.getTextRange()!!.getStartOffset() }
 
     private val itFakeDeclaration by Delegates.lazy { JetPsiFactory(originalFile).createParameter("it: Any?") }
 
@@ -111,6 +112,8 @@ data class ExtractionData(
             return function == null || !function.isInsideOf(originalElements)
         }
 
+        val originalStartOffset = originalStartOffset
+
         if (originalStartOffset != null) {
             val resultMap = HashMap<Int, ResolveResult>()
 

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractableAnalysisUtil.kt

@@ -16,64 +16,62 @@
 
 package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 
+import com.intellij.openapi.diagnostic.Logger
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.types.*
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.idea.refactoring.JetRefactoringBundle
-import org.jetbrains.kotlin.psi.psiUtil.isInsideOf
-import java.util.*
-import org.jetbrains.kotlin.idea.refactoring.createTempCopy
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.idea.refactoring.JetNameSuggester
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import com.intellij.psi.util.PsiTreeUtil
-import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
-import org.jetbrains.kotlin.resolve.scopes.receivers.ThisReceiver
-import org.jetbrains.kotlin.cfg.pseudocode.*
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.idea.refactoring.JetNameValidatorImpl
-import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
-import org.jetbrains.kotlin.idea.imports.canBeReferencedViaImport
 import com.intellij.psi.PsiNamedElement
-import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
-import org.jetbrains.kotlin.utils.DFS
-import org.jetbrains.kotlin.utils.DFS.*
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.refactoring.util.RefactoringUIUtil
 import com.intellij.util.containers.MultiMap
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.AnalysisResult.Status
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.AnalysisResult.ErrorMessage
-import org.jetbrains.kotlin.cfg.pseudocode.instructions.special.*
-import org.jetbrains.kotlin.cfg.pseudocode.instructions.*
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.cfg.pseudocode.*
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.Instruction
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.InstructionVisitorWithResult
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.InstructionWithNext
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.JetElementInstruction
 import org.jetbrains.kotlin.cfg.pseudocode.instructions.eval.*
 import org.jetbrains.kotlin.cfg.pseudocode.instructions.jumps.*
-import kotlin.properties.Delegates
-import org.jetbrains.kotlin.cfg.pseudocodeTraverser.traverse
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.special.LocalFunctionDeclarationInstruction
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.special.MarkInstruction
 import org.jetbrains.kotlin.cfg.pseudocodeTraverser.TraversalOrder
-import org.jetbrains.kotlin.resolve.bindingContextUtil.getTargetFunctionDescriptor
-import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsStatement
-import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
-import org.jetbrains.kotlin.idea.imports.importableFqNameSafe
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.Initializer
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.ParameterUpdate
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.ExpressionValue
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.Jump
+import org.jetbrains.kotlin.cfg.pseudocodeTraverser.traverse
 import org.jetbrains.kotlin.cfg.pseudocodeTraverser.traverseFollowingInstructions
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValueBoxer.AsList
-import org.jetbrains.kotlin.idea.refactoring.getContextForContainingDeclarationBody
-import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
-import org.jetbrains.kotlin.idea.caches.resolve.findModuleDescriptor
-import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
-import org.jetbrains.kotlin.idea.refactoring.comparePossiblyOverridingDescriptors
-import org.jetbrains.kotlin.idea.util.makeNullable
-import org.jetbrains.kotlin.resolve.calls.CallTransformer
-import org.jetbrains.kotlin.resolve.calls.callUtil.*
-import org.jetbrains.kotlin.diagnostics.*
-import java.util.logging.*
-import com.intellij.openapi.diagnostic.Logger
+import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
+import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.caches.resolve.analyzeFully
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
+import org.jetbrains.kotlin.idea.caches.resolve.findModuleDescriptor
+import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
+import org.jetbrains.kotlin.idea.imports.canBeReferencedViaImport
+import org.jetbrains.kotlin.idea.imports.importableFqNameSafe
+import org.jetbrains.kotlin.idea.refactoring.*
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.AnalysisResult.ErrorMessage
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.AnalysisResult.Status
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.ExpressionValue
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.Initializer
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.Jump
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.ParameterUpdate
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValueBoxer.AsList
+import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
+import org.jetbrains.kotlin.idea.util.makeNullable
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.isInsideOf
+import org.jetbrains.kotlin.psi.psiUtil.parents
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsStatement
+import org.jetbrains.kotlin.resolve.calls.CallTransformer
+import org.jetbrains.kotlin.resolve.calls.callUtil.getCalleeExpressionIfAny
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
+import org.jetbrains.kotlin.resolve.scopes.receivers.ThisReceiver
+import org.jetbrains.kotlin.types.*
+import org.jetbrains.kotlin.utils.DFS
+import org.jetbrains.kotlin.utils.DFS.CollectingNodeHandler
+import org.jetbrains.kotlin.utils.DFS.Neighbors
+import org.jetbrains.kotlin.utils.DFS.VisitedWithSet
+import java.util.*
+import kotlin.properties.Delegates
 
 private val DEFAULT_RETURN_TYPE = KotlinBuiltIns.getInstance().getUnitType()
 private val DEFAULT_PARAMETER_TYPE = KotlinBuiltIns.getInstance().getNullableAnyType()
@@ -221,12 +219,6 @@ private fun ExtractionData.analyzeControlFlow(
         options: ExtractionOptions,
         parameters: Set<Parameter>
 ): Pair<ControlFlow, ErrorMessage?> {
-    fun isCurrentFunctionReturn(expression: JetReturnExpression): Boolean {
-        val functionDescriptor = expression.getTargetFunctionDescriptor(bindingContext)
-        val currentDescriptor = bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, pseudocode.getCorrespondingElement()]
-        return currentDescriptor == functionDescriptor
-    }
-
     val exitPoints = localInstructions.getExitPoints()
 
     val valuedReturnExits = ArrayList<ReturnValueInstruction>()
@@ -246,21 +238,19 @@ private fun ExtractionData.analyzeControlFlow(
 
         when (insn) {
             is ReturnValueInstruction -> {
-                val returnExpression = insn.returnExpressionIfAny
-                if (returnExpression == null) {
-                    val containingDeclaration = insn.returnedValue.element?.getNonStrictParentOfType<JetDeclarationWithBody>()
-                    if (containingDeclaration == pseudocode.getCorrespondingElement()) {
+                if (insn.owner == pseudocode) {
+                    if (insn.returnExpressionIfAny == null) {
                         defaultExits.add(insn)
                     }
-                }
-                else if (isCurrentFunctionReturn(returnExpression)) {
-                    valuedReturnExits.add(insn)
+                    else {
+                        valuedReturnExits.add(insn)
+                    }
                 }
             }
 
             is AbstractJumpInstruction -> {
                 val element = insn.element
-                if ((element is JetReturnExpression && isCurrentFunctionReturn(element))
+                if ((element is JetReturnExpression && insn.owner == pseudocode)
                         || element is JetBreakExpression
                         || element is JetContinueExpression) {
                     jumpExits.add(insn)
@@ -775,11 +765,23 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
     val bindingContext = commonParent.getContextForContainingDeclarationBody()
     if (bindingContext == null) return noContainerError
 
-    val pseudocodeDeclaration = PsiTreeUtil.getParentOfType(
-            commonParent, javaClass<JetDeclarationWithBody>(), javaClass<JetClassOrObject>()
-    ) ?: commonParent.getNonStrictParentOfType<JetProperty>()
-    ?: return noContainerError
-    val pseudocode = PseudocodeUtil.generatePseudocode(pseudocodeDeclaration, bindingContext)
+    val pseudocodeDeclaration =
+            PsiTreeUtil.getParentOfType(commonParent, javaClass<JetDeclarationWithBody>(), javaClass<JetClassOrObject>())
+            ?: commonParent.getNonStrictParentOfType<JetProperty>()
+            ?: return noContainerError
+
+    val enclosingPseudocodeDeclaration = if (pseudocodeDeclaration is JetFunctionLiteral) {
+        commonParent.parents(withItself = false)
+                .firstOrNull { it is JetDeclaration && it !is JetFunctionLiteral } as? JetDeclaration
+        ?: pseudocodeDeclaration
+    }
+    else {
+        pseudocodeDeclaration
+    }
+
+    val enclosingPseudocode = PseudocodeUtil.generatePseudocode(enclosingPseudocodeDeclaration, bindingContext)
+    val pseudocode = enclosingPseudocode.getPseudocodeByElement(pseudocodeDeclaration)
+                     ?: throw AssertionError("Can't find nested pseudocode for element: ${JetPsiUtil.getElementTextWithContext(pseudocodeDeclaration)}")
     val localInstructions = getLocalInstructions(pseudocode)
 
     val modifiedVarDescriptorsWithExpressions = localInstructions.getModifiedVarDescriptors(bindingContext)

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractorUtil.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.idea.intentions.ConvertToExpressionBodyAction
 import org.jetbrains.kotlin.idea.refactoring.JetNameSuggester
@@ -286,10 +287,17 @@ private fun makeCall(
 
     fun wrapCall(outputValue: OutputValue, callText: String): List<PsiElement> {
         return when (outputValue) {
-            is OutputValue.ExpressionValue ->
-                Collections.singletonList(
-                        if (outputValue.callSiteReturn) psiFactory.createReturn(callText) else psiFactory.createExpression(callText)
-                )
+            is OutputValue.ExpressionValue -> {
+                val exprText = if (outputValue.callSiteReturn) {
+                    val firstReturn = outputValue.originalExpressions.filterIsInstance<JetReturnExpression>().firstOrNull()
+                    val label = firstReturn?.getTargetLabel()?.getText() ?: ""
+                    "return$label $callText"
+                }
+                else {
+                    callText
+                }
+                Collections.singletonList(psiFactory.createExpression(exprText))
+            }
 
             is ParameterUpdate ->
                 Collections.singletonList(
@@ -403,6 +411,17 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         nameByOffset.entrySet().forEach { e -> counterpartMap[e.getValue()]?.let { e.setValue(it) } }
     }
 
+    fun getCounterparts<T : JetExpression>(originalExpressions: Collection<T>,
+                                           body: JetExpression,
+                                           bodyOffset: Int,
+                                           file: PsiFile): List<T> {
+        return originalExpressions.map { originalExpression ->
+            val offsetInBody = originalExpression.getTextRange()!!.getStartOffset() - descriptor.extractionData.originalStartOffset!!
+            file.findElementAt(bodyOffset + offsetInBody)?.getNonStrictParentOfType(originalExpression.javaClass)
+            ?: throw AssertionError("Couldn't find expression at $offsetInBody in '${body.getText()}'")
+        }
+    }
+
     fun adjustDeclarationBody(declaration: JetNamedDeclaration) {
         val body = declaration.getGeneratedBody()
 
@@ -432,16 +451,14 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         val replacingReturn: JetExpression?
         val expressionsToReplaceWithReturn: List<JetElement>
 
+        val returnsForLabelRemoval = descriptor.controlFlow.outputValues
+                .flatMapTo(ArrayList<JetReturnExpression>()) { it.originalExpressions.filterIsInstance<JetReturnExpression>() }
+
         val jumpValue = descriptor.controlFlow.jumpOutputValue
         if (jumpValue != null) {
             replacingReturn = psiFactory.createExpression(if (jumpValue.conditional) "return true" else "return")
-            expressionsToReplaceWithReturn = jumpValue.elementsToReplace.map { jumpElement ->
-                val offsetInBody = jumpElement.getTextRange()!!.getStartOffset() - descriptor.extractionData.originalStartOffset!!
-                val expr = file.findElementAt(bodyOffset + offsetInBody)?.getNonStrictParentOfType(jumpElement.javaClass)
-                assert(expr != null, "Couldn't find expression at $offsetInBody in '${body.getText()}'")
-
-                expr!!
-            }
+            returnsForLabelRemoval.removeAll(jumpValue.elementsToReplace)
+            expressionsToReplaceWithReturn = getCounterparts(jumpValue.elementsToReplace, body, bodyOffset, file)
         }
         else {
             replacingReturn = null
@@ -454,6 +471,8 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
             }
         }
 
+        getCounterparts(returnsForLabelRemoval, body, bodyOffset, file).forEach { it.getTargetLabel()?.delete() }
+
         for ((expr, originalOffset) in originalOffsetByExpr) {
             if (expr.isValid()) {
                 nameByOffset.put(originalOffset, exprReplacementMap[expr]?.invoke(expr) ?: expr)

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/labeledAndUnlabeledReturn1.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(a: Int): Int {
+    a.let {
+        <selection>if (it > 0) return it else return@foo -it</selection>
+    }
+    return 0
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/labeledAndUnlabeledReturn1.kt.after

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(a: Int): Int {
+    a.let {
+        return i(it)
+    }
+    return 0
+}
+
+private fun i(it: Int): Int {
+    if (it > 0) return it else return -it
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/labeledAndUnlabeledReturn2.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(a: Int): Int {
+    a.let {
+        <selection>if (it > 0) return@foo it else return -it</selection>
+    }
+    return 0
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/labeledAndUnlabeledReturn2.kt.after

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(a: Int): Int {
+    a.let {
+        return@foo i(it)
+    }
+    return 0
+}
+
+private fun i(it: Int): Int {
+    if (it > 0) return it else return -it
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/nonLocalReturn.kt

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(): Int {
+    1.let { <selection>return it + 1</selection> }
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/nonLocalReturn.kt.after

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(): Int {
+    1.let { return i(it) }
+}
+
+private fun i(it: Int): Int {
+    return it + 1
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/nonLocalReturnWithLabel.kt

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(): Int {
+    1.let { <selection>return@foo it + 1</selection> }
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/nonLocalReturnWithLabel.kt.after

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// PARAM_DESCRIPTOR: value-parameter val it: kotlin.Int defined in foo.<anonymous>
+// PARAM_TYPES: kotlin.Int
+fun foo(): Int {
+    1.let { return@foo i(it) }
+}
+
+private fun i(it: Int): Int {
+    return it + 1
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/returnWithLabel.kt

@@ -0,0 +1,3 @@
+fun foo(): Int {
+    <selection>return@foo 1</selection>
+}

idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/returnWithLabel.kt.after

@@ -0,0 +1,7 @@
+fun foo(): Int {
+    return@foo i()
+}
+
+private fun i(): Int {
+    return 1
+}

idea/testData/refactoring/extractFunction/controlFlow/unextractable/unmatchedReturnTargets.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo(a: Int): Int {
+    a.let {
+        <selection>if (it > 0) return@let it else return@foo -it</selection>
+    }
+    return 0
+}

idea/testData/refactoring/extractFunction/controlFlow/unextractable/unmatchedReturnTargets.kt.conflicts

@@ -0,0 +1 @@
+Selected code fragment has multiple exit points

idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/JetExtractionTestGenerated.java

@@ -668,6 +668,36 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                     doExtractFunctionTest(fileName);
                 }
 
+                @TestMetadata("labeledAndUnlabeledReturn1.kt")
+                public void testLabeledAndUnlabeledReturn1() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/labeledAndUnlabeledReturn1.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("labeledAndUnlabeledReturn2.kt")
+                public void testLabeledAndUnlabeledReturn2() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/labeledAndUnlabeledReturn2.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("nonLocalReturn.kt")
+                public void testNonLocalReturn() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/nonLocalReturn.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("nonLocalReturnWithLabel.kt")
+                public void testNonLocalReturnWithLabel() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/nonLocalReturnWithLabel.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("returnWithLabel.kt")
+                public void testReturnWithLabel() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/returnWithLabel.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
                 @TestMetadata("singleDefiniteReturn.kt")
                 public void testSingleDefiniteReturn() throws Exception {
                     String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/definiteReturns/singleDefiniteReturn.kt");
@@ -1118,6 +1148,12 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                     doExtractFunctionTest(fileName);
                 }
 
+                @TestMetadata("unmatchedReturnTargets.kt")
+                public void testUnmatchedReturnTargets() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/unextractable/unmatchedReturnTargets.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
                 @TestMetadata("variablesOutOfScope.kt")
                 public void testVariablesOutOfScope() throws Exception {
                     String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/unextractable/variablesOutOfScope.kt");