Extract Function: Merge jumps with default exits if all exit points are equivalent

#KT-6598 Fixed

idea/src/org/jetbrains/kotlin/idea/refactoring/extractFunction/ExtractableCodeDescriptor.kt

@@ -132,7 +132,7 @@ trait OutputValue {
 
     class Jump(
             val elementsToReplace: List<JetExpression>,
-            val elementToInsertAfterCall: JetElement,
+            val elementToInsertAfterCall: JetElement?,
             val conditional: Boolean
     ): OutputValue {
         override val originalExpressions: List<JetExpression> get() = elementsToReplace

idea/src/org/jetbrains/kotlin/idea/refactoring/extractFunction/extractableAnalysisUtil.kt

@@ -32,7 +32,6 @@ 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.cfg.Label
 import org.jetbrains.kotlin.idea.refactoring.JetNameValidatorImpl
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToDeclarationUtil
 import org.jetbrains.kotlin.idea.imports.canBeReferencedViaImport
@@ -45,7 +44,7 @@ import com.intellij.util.containers.MultiMap
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.refactoring.extractFunction.AnalysisResult.Status
 import org.jetbrains.kotlin.idea.refactoring.extractFunction.AnalysisResult.ErrorMessage
-import org.jetbrains.kotlin.cfg.pseudocode.instructions.special.LocalFunctionDeclarationInstruction
+import org.jetbrains.kotlin.cfg.pseudocode.instructions.special.*
 import org.jetbrains.kotlin.cfg.pseudocode.instructions.*
 import org.jetbrains.kotlin.cfg.pseudocode.instructions.eval.*
 import org.jetbrains.kotlin.cfg.pseudocode.instructions.jumps.*
@@ -152,9 +151,37 @@ private fun List<Instruction>.getResultTypeAndExpressions(
     return resultType to expressions
 }
 
-private fun List<AbstractJumpInstruction>.checkEquivalence(checkPsi: Boolean): Boolean {
+private fun getCommonNonTrivialSuccessorIfAny(instructions: List<Instruction>): Instruction? {
-    if (mapTo(HashSet<Label?>()) { it.targetLabel }.size > 1) return false
+    val singleSuccessorCheckingVisitor = object: InstructionVisitorWithResult<Boolean>() {
-    return !checkPsi || mapTo(HashSet<String?>()) { it.element.getText() }.size <= 1
+        var target: Instruction? = null
+
+        override fun visitInstructionWithNext(instruction: InstructionWithNext): Boolean {
+            return when (instruction) {
+                is LoadUnitValueInstruction,
+                is MergeInstruction,
+                is MarkInstruction -> {
+                    instruction.next?.accept(this) ?: true
+                }
+                else -> visitInstruction(instruction)
+            }
+        }
+
+        override fun visitJump(instruction: AbstractJumpInstruction): Boolean {
+            return when (instruction) {
+                is ConditionalJumpInstruction -> visitInstruction(instruction)
+                else -> instruction.resolvedTarget?.accept(this) ?: true
+            }
+        }
+
+        override fun visitInstruction(instruction: Instruction): Boolean {
+            if (target != null && target != instruction) return false
+            target = instruction
+            return true
+        }
+    }
+
+    if (instructions.flatMap { it.nextInstructions }.any { !it.accept(singleSuccessorCheckingVisitor) }) return null
+    return singleSuccessorCheckingVisitor.target ?: instructions.firstOrNull()?.owner?.getSinkInstruction()
 }
 
 private fun JetType.isMeaningful(): Boolean {
@@ -295,7 +322,7 @@ private fun ExtractionData.analyzeControlFlow(
             return controlFlow.copy(outputValues = Collections.singletonList(Jump(listOf(element), element, true))) to null
         }
 
-        if (!valuedReturnExits.checkEquivalence(false)) return multipleExitsError
+        if (getCommonNonTrivialSuccessorIfAny(valuedReturnExits) == null) return multipleExitsError
         outputValues.add(ExpressionValue(true, valuedReturnExpressions, returnValueType))
     }
 
@@ -326,12 +353,14 @@ private fun ExtractionData.analyzeControlFlow(
     }
 
     if (jumpExits.isNotEmpty()) {
-        if (!jumpExits.checkEquivalence(true)) return multipleExitsError
+        val jumpTarget = getCommonNonTrivialSuccessorIfAny(jumpExits)
+        if (jumpTarget == null) return multipleExitsError
 
+        val singleExit = getCommonNonTrivialSuccessorIfAny(defaultExits) == jumpTarget
+        val conditional = !singleExit && defaultExits.isNotEmpty()
         val elements = jumpExits.map { it.element as JetExpression }
-        return controlFlow.copy(
+        val elementToInsertAfterCall = if (singleExit) null else elements.first()
-                outputValues = Collections.singletonList(Jump(elements, elements.first(), defaultExits.isNotEmpty()))
+        return controlFlow.copy(outputValues = Collections.singletonList(Jump(elements, elementToInsertAfterCall, conditional))) to null
-        ) to null
     }
 
     return controlFlow to null

idea/src/org/jetbrains/kotlin/idea/refactoring/extractFunction/extractorUtil.kt

@@ -287,7 +287,10 @@ private fun makeCall(
                 )
 
             is Jump -> {
-                if (outputValue.conditional) {
+                if (outputValue.elementToInsertAfterCall == null) {
+                    Collections.singletonList(psiFactory.createExpression(callText))
+                }
+                else if (outputValue.conditional) {
                     Collections.singletonList(
                             psiFactory.createExpression("if ($callText) ${outputValue.elementToInsertAfterCall.getText()}")
                     )

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/breakAndReturn.kt

@@ -0,0 +1,14 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    for (n in 1..b) {
+        <selection>if (a > 0) throw Exception("")
+        if (a + b > 0) break
+        println(a - b)
+        return</selection>
+    }
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/breakAndReturn.kt.after

@@ -0,0 +1,19 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    for (n in 1..b) {
+        __dummyTestFun__(a, b)
+        break
+    }
+}
+
+private fun __dummyTestFun__(a: Int, b: Int) {
+    if (a > 0) throw Exception("")
+    if (a + b > 0) return
+    println(a - b)
+    return
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/breakContinueAndDefault.kt

@@ -0,0 +1,15 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    @loop1 for (p in 1..b) {
+        @loop2 for (n in 1..b) {
+            <selection>if (a > 0) throw Exception("")
+            if (a + b > 0) break@loop2
+            if (a - b > 0) continue@loop1</selection>
+        }
+    }
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/breakContinueAndDefault.kt.after

@@ -0,0 +1,20 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    @loop1 for (p in 1..b) {
+        @loop2 for (n in 1..b) {
+            if (b(a, b)) break@loop2
+        }
+    }
+}
+
+private fun b(a: Int, b: Int): Boolean {
+    if (a > 0) throw Exception("")
+    if (a + b > 0) return true
+    if (a - b > 0) return true
+    return false
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/continueAndReturn.kt

@@ -0,0 +1,10 @@
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    for (n in 1..b) {
+        <selection>if (a > 0) throw Exception("")
+        if (a + b > 0) continue
+        println(a - b)
+        return</selection>
+    }
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/continueAndReturn.kt.conflicts

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

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndBreak.kt

@@ -0,0 +1,13 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    for (n in 1..b) {
+        <selection>if (a > 0) throw Exception("")
+        if (a + b > 0) break
+        println(a - b)</selection>
+    }
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndBreak.kt.after

@@ -0,0 +1,18 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    for (n in 1..b) {
+        if (b(a, b)) break
+    }
+}
+
+private fun b(a: Int, b: Int): Boolean {
+    if (a > 0) throw Exception("")
+    if (a + b > 0) return true
+    println(a - b)
+    return false
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndContinue.kt

@@ -0,0 +1,13 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    for (n in 1..b) {
+        <selection>if (a > 0) throw Exception("")
+        if (a + b > 0) continue
+        println(a - b)</selection>
+    }
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndContinue.kt.after

@@ -0,0 +1,17 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    for (n in 1..b) {
+        __dummyTestFun__(a, b)
+    }
+}
+
+private fun __dummyTestFun__(a: Int, b: Int) {
+    if (a > 0) throw Exception("")
+    if (a + b > 0) return
+    println(a - b)
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndReturn.kt

@@ -0,0 +1,12 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+
+    <selection>if (a > 0) throw Exception("")
+    if (b + a > 0) return
+    println(a - b)</selection>
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndReturn.kt.after

@@ -0,0 +1,16 @@
+// PARAM_TYPES: kotlin.Int
+// PARAM_TYPES: kotlin.Int
+// PARAM_DESCRIPTOR: value-parameter val a: kotlin.Int defined in foo
+// PARAM_DESCRIPTOR: val b: kotlin.Int defined in foo
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+
+    __dummyTestFun__(a, b)
+}
+
+private fun __dummyTestFun__(a: Int, b: Int) {
+    if (a > 0) throw Exception("")
+    if (b + a > 0) return
+    println(a - b)
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndReturnInWhen.kt

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+// PARAM_TYPES: kotlin.Int, Comparable<Int>
+// PARAM_DESCRIPTOR: value-parameter val i: kotlin.Int defined in example
+fun example(i: Int) {
+    when (i) {
+        1 -> {
+            <selection>if (i > 5) {
+                println("true")
+            }
+            else {
+                println("false")
+                return
+            }
+            println("!!")</selection>
+        }
+    }
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndReturnInWhen.kt.after

@@ -0,0 +1,20 @@
+// WITH_RUNTIME
+// PARAM_TYPES: kotlin.Int, Comparable<Int>
+// PARAM_DESCRIPTOR: value-parameter val i: kotlin.Int defined in example
+fun example(i: Int) {
+    when (i) {
+        1 -> {
+            __dummyTestFun__(i)
+        }
+    }
+}
+
+private fun __dummyTestFun__(i: Int) {
+    if (i > 5) {
+        println("true")
+    } else {
+        println("false")
+        return
+    }
+    println("!!")
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/multipleBreaks.kt

@@ -0,0 +1,12 @@
+// SIBLING:
+fun foo(a: Int) {
+    val b: Int = 1
+    @loop1 for (p in 1..b) {
+        @loop2 for (n in 1..b) {
+            <selection>if (a > 0) throw Exception("")
+            if (a + b > 0) break@loop1
+            println(a - b)
+            break@loop2</selection>
+        }
+    }
+}

idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/multipleBreaks.kt.conflicts

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

idea/testData/refactoring/extractFunction/controlFlow/throws/continueWithThrow.kt

@@ -9,5 +9,6 @@ fun foo(a: Int): Int {
         <selection>if (a > 0) throw Exception("")
         if (a + b > 0) continue
         println(a - b)</selection>
+        println(a + b)
     }
 }

idea/testData/refactoring/extractFunction/controlFlow/throws/continueWithThrow.kt.after

@@ -7,6 +7,7 @@ fun foo(a: Int): Int {
     val b: Int = 1
     for (n in 1..b) {
         if (b(a, b)) continue
+        println(a + b)
     }
 }
 

idea/testData/refactoring/extractFunction/controlFlow/throws/nonValuedReturnWithThrow.kt

@@ -9,4 +9,5 @@ fun foo(a: Int) {
     <selection>if (a > 0) throw Exception("")
     if (b + a > 0) return
     println(a - b)</selection>
+    println(a + b)
 }

idea/testData/refactoring/extractFunction/controlFlow/throws/nonValuedReturnWithThrow.kt.after

@@ -7,6 +7,7 @@ fun foo(a: Int) {
     val b: Int = 1
 
     if (b(a, b)) return
+    println(a + b)
 }
 
 private fun b(a: Int, b: Int): Boolean {

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

@@ -501,7 +501,7 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
 
         @TestMetadata("idea/testData/refactoring/extractFunction/controlFlow")
         @TestDataPath("$PROJECT_ROOT")
-        @InnerTestClasses({ControlFlow.ConditionalJumps.class, ControlFlow.Default.class, ControlFlow.DefiniteReturns.class, ControlFlow.EvaluateExpression.class, ControlFlow.Initializer.class, ControlFlow.OutputValues.class, ControlFlow.Throws.class, ControlFlow.Unextractable.class})
+        @InnerTestClasses({ControlFlow.ConditionalJumps.class, ControlFlow.Default.class, ControlFlow.DefiniteReturns.class, ControlFlow.EvaluateExpression.class, ControlFlow.ExitPointEquivalence.class, ControlFlow.Initializer.class, ControlFlow.OutputValues.class, ControlFlow.Throws.class, ControlFlow.Unextractable.class})
         @RunWith(JUnit3RunnerWithInners.class)
         public static class ControlFlow extends AbstractJetExtractionTest {
             public void testAllFilesPresentInControlFlow() throws Exception {
@@ -748,6 +748,63 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 }
             }
 
+            @TestMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence")
+            @TestDataPath("$PROJECT_ROOT")
+            @RunWith(JUnit3RunnerWithInners.class)
+            public static class ExitPointEquivalence extends AbstractJetExtractionTest {
+                public void testAllFilesPresentInExitPointEquivalence() throws Exception {
+                    JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence"), Pattern.compile("^(.+)\\.kt$"), true);
+                }
+
+                @TestMetadata("breakAndReturn.kt")
+                public void testBreakAndReturn() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/breakAndReturn.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("breakContinueAndDefault.kt")
+                public void testBreakContinueAndDefault() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/breakContinueAndDefault.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("continueAndReturn.kt")
+                public void testContinueAndReturn() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/continueAndReturn.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("defaultAndBreak.kt")
+                public void testDefaultAndBreak() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndBreak.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("defaultAndContinue.kt")
+                public void testDefaultAndContinue() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndContinue.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("defaultAndReturn.kt")
+                public void testDefaultAndReturn() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndReturn.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("defaultAndReturnInWhen.kt")
+                public void testDefaultAndReturnInWhen() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/defaultAndReturnInWhen.kt");
+                    doExtractFunctionTest(fileName);
+                }
+
+                @TestMetadata("multipleBreaks.kt")
+                public void testMultipleBreaks() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/exitPointEquivalence/multipleBreaks.kt");
+                    doExtractFunctionTest(fileName);
+                }
+            }
+
             @TestMetadata("idea/testData/refactoring/extractFunction/controlFlow/initializer")
             @TestDataPath("$PROJECT_ROOT")
             @RunWith(JUnit3RunnerWithInners.class)