Extract Function: Limit the set of allowed containers for the default action

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/ExtractKotlinFunctionHandler.kt

@@ -50,8 +50,11 @@ import javax.swing.event.HyperlinkEvent
 import com.intellij.refactoring.BaseRefactoringProcessor.ConflictsInTestsException
 import com.intellij.ui.awt.RelativePoint
 import com.intellij.openapi.ui.popup.Balloon.Position
+import org.jetbrains.jet.lang.psi.psiUtil.getParentByType
+import org.jetbrains.jet.lang.psi.JetDeclaration
+import java.util.Collections
 
-public class ExtractKotlinFunctionHandler : RefactoringActionHandler {
+public class ExtractKotlinFunctionHandler(public val allContainersEnabled: Boolean = false) : RefactoringActionHandler {
     fun doInvoke(
             editor: Editor,
             file: JetFile,
@@ -119,7 +122,7 @@ public class ExtractKotlinFunctionHandler : RefactoringActionHandler {
     override fun invoke(project: Project, editor: Editor, file: PsiFile, dataContext: DataContext?) {
         if (file !is JetFile) return
 
-        selectElements(editor, file) { (elements, targetSibling) ->
+        selectElements(editor, file, allContainersEnabled) { (elements, targetSibling) ->
             doInvoke(editor, file, elements, targetSibling)
         }
     }
@@ -142,6 +145,7 @@ private fun showErrorHintByKey(project: Project, editor: Editor, key: String) {
 fun selectElements(
         editor: Editor,
         file: PsiFile,
+        allContainersEnabled: Boolean = false,
         continuation: (elements: List<PsiElement>, targetSibling: PsiElement) -> Unit
 ) {
     fun noExpressionError() {
@@ -167,17 +171,40 @@ fun selectElements(
         continuation(elements, outermostParent)
     }
 
+    fun getContainers(element: PsiElement, strict: Boolean): List<JetElement> {
+        if (allContainersEnabled) return element.getAllExtractionContainers(strict)
+
+        val declaration = element.getParentByType(javaClass<JetDeclaration>(), strict)
+        if (declaration == null) return Collections.emptyList()
+
+        val parent = declaration.getParent()
+        return when (parent) {
+            is JetFile -> Collections.singletonList(parent)
+            is JetClassBody -> {
+                element.getAllExtractionContainers(strict)
+                        .filter {
+                            it is JetClassBody || (it is JetBlockExpression && it.getParent() is JetDeclarationWithBody)
+                        }
+                        .dropWhile { it !is JetClassBody }
+            }
+            else -> {
+                val targetContainer = parent?.getParentByType(javaClass<JetDeclarationWithBody>())?.getBodyExpression()
+                if (targetContainer is JetBlockExpression) Collections.singletonList(targetContainer) else Collections.emptyList()
+            }
+        }
+    }
+
     fun selectTargetContainer(elements: List<PsiElement>) {
         val parent = PsiTreeUtil.findCommonParent(elements)
             ?: throw AssertionError("Should have at least one parent: ${elements.makeString("\n")}")
 
-        val containers = parent.getAllExtractionContainers(elements.size == 1)
+        val containers = getContainers(parent, elements.size == 1)
         if (containers.empty) {
             noContainerError()
             return
         }
 
-        if (ApplicationManager.getApplication()!!.isUnitTestMode()) {
+        if (containers.size == 1 || ApplicationManager.getApplication()!!.isUnitTestMode()) {
             onSelectionComplete(parent, elements, containers[0])
             return
         }

idea/testData/refactoring/extractFunction/defaultContainer/anonymousObject.kt

@@ -0,0 +1,9 @@
+class A {
+    fun foo(a: Int, b: Int): Int {
+        return object: Function0<Int> {
+            override fun invoke(): Int {
+                return <selection>a + b - 1</selection>
+            }
+        }
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/anonymousObject.kt.after

@@ -0,0 +1,13 @@
+class A {
+    fun foo(a: Int, b: Int): Int {
+        return object: Function0<Int> {
+            override fun invoke(): Int {
+                return i()
+            }
+
+            fun i(): Int {
+                return a + b - 1
+            }
+        }
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/classFunction.kt

@@ -0,0 +1,7 @@
+class A {
+    class B {
+        fun foo(a: Int, b: Int): Int {
+            return <selection>a + b - 1</selection>
+        }
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/classFunction.kt.after

@@ -0,0 +1,11 @@
+class A {
+    class B {
+        fun foo(a: Int, b: Int): Int {
+            return i(a, b)
+        }
+
+        fun i(a: Int, b: Int): Int {
+            return a + b - 1
+        }
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/lambda.kt

@@ -0,0 +1,7 @@
+class A {
+    fun foo(a: Int, b: Int): Int {
+        return {
+            <selection>a + b - 1</selection>
+        }.invoke()
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/lambda.kt.after

@@ -0,0 +1,11 @@
+class A {
+    fun foo(a: Int, b: Int): Int {
+        fun i(): Int {
+            return a + b - 1
+        }
+
+        return {
+            i()
+        }.invoke()
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/localClass.kt

@@ -0,0 +1,11 @@
+class A {
+    fun foo(a: Int, b: Int): Int {
+        class L: Function0<Int> {
+            override fun invoke(): Int {
+                return <selection>a + b - 1</selection>
+            }
+        }
+
+        return L().invoke()
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/localClass.kt.after

@@ -0,0 +1,15 @@
+class A {
+    fun foo(a: Int, b: Int): Int {
+        class L: Function0<Int> {
+            override fun invoke(): Int {
+                return i()
+            }
+
+            fun i(): Int {
+                return a + b - 1
+            }
+        }
+
+        return L().invoke()
+    }
+}

idea/testData/refactoring/extractFunction/defaultContainer/localFunction.kt

@@ -0,0 +1,7 @@
+fun foo(a: Int, b: Int): Int {
+    fun bar() {
+        return <selection>a + b - 1</selection>
+    }
+
+    return bar()
+}

idea/testData/refactoring/extractFunction/defaultContainer/localFunction.kt.after

@@ -0,0 +1,11 @@
+fun foo(a: Int, b: Int): Int {
+    fun i(): Int {
+        return a + b - 1
+    }
+
+    fun bar() {
+        return i()
+    }
+
+    return bar()
+}

idea/testData/refactoring/extractFunction/defaultContainer/topLevelFunction.kt

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

idea/testData/refactoring/extractFunction/defaultContainer/topLevelFunction.kt.after

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

idea/tests/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/JetExtractionTestGenerated.java

@@ -191,7 +191,7 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
     }
     
     @TestMetadata("idea/testData/refactoring/extractFunction")
-    @InnerTestClasses({ExtractFunction.Basic.class, ExtractFunction.ControlFlow.class, ExtractFunction.Parameters.class, ExtractFunction.TypeParameters.class})
+    @InnerTestClasses({ExtractFunction.Basic.class, ExtractFunction.ControlFlow.class, ExtractFunction.DefaultContainer.class, ExtractFunction.Parameters.class, ExtractFunction.TypeParameters.class})
     public static class ExtractFunction extends AbstractJetExtractionTest {
         public void testAllFilesPresentInExtractFunction() throws Exception {
             JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.TestsPackage", new File("idea/testData/refactoring/extractFunction"), Pattern.compile("^(.+)\\.kt$"), true);
@@ -537,6 +537,44 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
             }
         }
         
+        @TestMetadata("idea/testData/refactoring/extractFunction/defaultContainer")
+        public static class DefaultContainer extends AbstractJetExtractionTest {
+            public void testAllFilesPresentInDefaultContainer() throws Exception {
+                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.TestsPackage", new File("idea/testData/refactoring/extractFunction/defaultContainer"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+            
+            @TestMetadata("anonymousObject.kt")
+            public void testAnonymousObject() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/defaultContainer/anonymousObject.kt");
+            }
+            
+            @TestMetadata("classFunction.kt")
+            public void testClassFunction() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/defaultContainer/classFunction.kt");
+            }
+            
+            @TestMetadata("lambda.kt")
+            public void testLambda() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/defaultContainer/lambda.kt");
+            }
+            
+            @TestMetadata("localClass.kt")
+            public void testLocalClass() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/defaultContainer/localClass.kt");
+            }
+            
+            @TestMetadata("localFunction.kt")
+            public void testLocalFunction() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/defaultContainer/localFunction.kt");
+            }
+            
+            @TestMetadata("topLevelFunction.kt")
+            public void testTopLevelFunction() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/defaultContainer/topLevelFunction.kt");
+            }
+            
+        }
+        
         @TestMetadata("idea/testData/refactoring/extractFunction/parameters")
         @InnerTestClasses({Parameters.ExtractSuper.class, Parameters.ExtractThis.class, Parameters.Misc.class, Parameters.NonDenotableTypes.class})
         public static class Parameters extends AbstractJetExtractionTest {
@@ -745,6 +783,7 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
             suite.addTestSuite(ExtractFunction.class);
             suite.addTestSuite(Basic.class);
             suite.addTest(ControlFlow.innerSuite());
+            suite.addTestSuite(DefaultContainer.class);
             suite.addTest(Parameters.innerSuite());
             suite.addTestSuite(TypeParameters.class);
             return suite;