Refactoring: Allow multiple ElementKinds when choosing element to refactor

idea/src/org/jetbrains/kotlin/idea/refactoring/KotlinRefactoringUtil.java

@@ -34,6 +34,8 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.ui.components.JBList;
 import com.intellij.util.Function;
 import com.intellij.util.containers.ContainerUtil;
+import kotlin.collections.CollectionsKt;
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.asJava.KtLightMethod;
@@ -333,16 +335,16 @@ public class KotlinRefactoringUtil {
     public static void selectElement(
             @NotNull Editor editor,
             @NotNull KtFile file,
-            @NotNull CodeInsightUtils.ElementKind elementKind,
+            @NotNull Collection<CodeInsightUtils.ElementKind> elementKinds,
             @NotNull SelectElementCallback callback
     ) throws IntroduceRefactoringException {
-        selectElement(editor, file, true, elementKind, callback);
+        selectElement(editor, file, true, elementKinds, callback);
     }
 
     public static void selectElement(@NotNull Editor editor,
             @NotNull KtFile file,
             boolean failOnEmptySuggestion,
-            @NotNull CodeInsightUtils.ElementKind elementKind,
+            @NotNull Collection<CodeInsightUtils.ElementKind> elementKinds,
             @NotNull SelectElementCallback callback
     ) throws IntroduceRefactoringException {
         if (editor.getSelectionModel().hasSelection()) {
@@ -351,11 +353,20 @@ public class KotlinRefactoringUtil {
             String text = file.getText();
             while (selectionStart < selectionEnd && Character.isSpaceChar(text.charAt(selectionStart))) ++selectionStart;
             while (selectionStart < selectionEnd && Character.isSpaceChar(text.charAt(selectionEnd - 1))) --selectionEnd;
-            callback.run(findElement(file, selectionStart, selectionEnd, failOnEmptySuggestion, elementKind));
+
+            for (CodeInsightUtils.ElementKind elementKind : elementKinds) {
+                PsiElement element = findElement(file, selectionStart, selectionEnd, failOnEmptySuggestion, elementKind);
+                if (element != null) {
+                    callback.run(element);
+                    return;
+                }
+            }
+
+            callback.run(null);
         }
         else {
             int offset = editor.getCaretModel().getOffset();
-            smartSelectElement(editor, file, offset, failOnEmptySuggestion, elementKind, callback);
+            smartSelectElement(editor, file, offset, failOnEmptySuggestion, elementKinds, callback);
         }
     }
 
@@ -442,13 +453,21 @@ public class KotlinRefactoringUtil {
 
     private static void smartSelectElement(
             @NotNull Editor editor,
-            @NotNull PsiFile file,
+            @NotNull final PsiFile file,
-            int offset,
+            final int offset,
             boolean failOnEmptySuggestion,
-            @NotNull CodeInsightUtils.ElementKind elementKind,
+            @NotNull Collection<CodeInsightUtils.ElementKind> elementKinds,
             @NotNull final SelectElementCallback callback
     ) throws IntroduceRefactoringException {
-        List<KtElement> elements = getSmartSelectSuggestions(file, offset, elementKind);
+        List<KtElement> elements = CollectionsKt.flatMap(
+                elementKinds,
+                new Function1<CodeInsightUtils.ElementKind, Iterable<? extends KtElement>>() {
+                    @Override
+                    public Iterable<? extends KtElement> invoke(CodeInsightUtils.ElementKind kind) {
+                        return getSmartSelectSuggestions(file, offset, kind);
+                    }
+                }
+        );
         if (elements.size() == 0) {
             if (failOnEmptySuggestion) throw new IntroduceRefactoringException(
                     KotlinRefactoringBundle.message("cannot.refactor.not.expression"));
@@ -542,7 +561,7 @@ public class KotlinRefactoringUtil {
         return element;
     }
 
-    public static class IntroduceRefactoringException extends Exception {
+    public static class IntroduceRefactoringException extends RuntimeException {
         private final String myMessage;
 
         public IntroduceRefactoringException(String message) {

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractFunction/ExtractKotlinFunctionHandler.kt

@@ -68,7 +68,7 @@ class ExtractKotlinFunctionHandler(
                 editor,
                 file,
                 "Select target code block",
-                CodeInsightUtils.ElementKind.EXPRESSION,
+                listOf(CodeInsightUtils.ElementKind.EXPRESSION),
                 { elements, parent -> parent.getExtractionContainers(elements.size == 1, allContainersEnabled) },
                 continuation
         )

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceParameter/KotlinIntroduceParameterHandler.kt

@@ -191,7 +191,7 @@ fun selectNewParameterContext(
             editor = editor,
             file = file,
             title = "Introduce parameter to declaration",
-            elementKind = CodeInsightUtils.ElementKind.EXPRESSION,
+            elementKinds = listOf(CodeInsightUtils.ElementKind.EXPRESSION),
             getContainers = { elements, parent ->
                 val parents = parent.parents
                 val stopAt = (parent.parents.zip(parent.parents.drop(1)))

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceProperty/KotlinIntroducePropertyHandler.kt

@@ -72,7 +72,7 @@ class KotlinIntroducePropertyHandler(
                 editor,
                 file,
                 "Select target code block",
-                CodeInsightUtils.ElementKind.EXPRESSION,
+                listOf(CodeInsightUtils.ElementKind.EXPRESSION),
                 { elements, parent ->
                     parent.getExtractionContainers(strict = true, includeAll = true).filter { it is KtClassBody || (it is KtFile && !it.isScript) }
                 },

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceTypeAlias/IntroduceTypeAliasHandler.kt

@@ -24,7 +24,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.refactoring.RefactoringActionHandler
-import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils
+import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils.ElementKind.TYPE_ELEMENT
 import org.jetbrains.kotlin.idea.refactoring.checkConflictsInteractively
 import org.jetbrains.kotlin.idea.refactoring.getExtractionContainers
 import org.jetbrains.kotlin.idea.refactoring.introduce.AbstractIntroduceAction
@@ -47,7 +47,7 @@ object KotlinIntroduceTypeAliasHandler : RefactoringActionHandler {
                 editor,
                 file,
                 "Select target code block",
-                CodeInsightUtils.ElementKind.TYPE_ELEMENT,
+                listOf(TYPE_ELEMENT),
                 { elements, parent -> parent.getExtractionContainers(strict = true, includeAll = true) },
                 continuation
         )

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtil.kt

@@ -45,7 +45,7 @@ fun selectElementsWithTargetSibling(
         editor: Editor,
         file: KtFile,
         title: String,
-        elementKind: CodeInsightUtils.ElementKind,
+        elementKinds: Collection<CodeInsightUtils.ElementKind>,
         getContainers: (elements: List<PsiElement>, commonParent: PsiElement) -> List<PsiElement>,
         continuation: (elements: List<PsiElement>, targetSibling: PsiElement) -> Unit
 ) {
@@ -68,7 +68,7 @@ fun selectElementsWithTargetSibling(
         continuation(elements, outermostParent)
     }
 
-    selectElementsWithTargetParent(operationName, editor, file, title, elementKind, getContainers, ::onSelectionComplete)
+    selectElementsWithTargetParent(operationName, editor, file, title, elementKinds, getContainers, ::onSelectionComplete)
 }
 
 fun selectElementsWithTargetParent(
@@ -76,7 +76,7 @@ fun selectElementsWithTargetParent(
         editor: Editor,
         file: KtFile,
         title: String,
-        elementKind: CodeInsightUtils.ElementKind,
+        elementKinds: Collection<CodeInsightUtils.ElementKind>,
         getContainers: (elements: List<PsiElement>, commonParent: PsiElement) -> List<PsiElement>,
         continuation: (elements: List<PsiElement>, targetParent: PsiElement) -> Unit
 ) {
@@ -109,27 +109,26 @@ fun selectElementsWithTargetParent(
         val startOffset = editor.selectionModel.selectionStart
         val endOffset = editor.selectionModel.selectionEnd
 
-        val elements = CodeInsightUtils.findElements(file, startOffset, endOffset, elementKind)
+        val elements = elementKinds.flatMap { CodeInsightUtils.findElements(file, startOffset, endOffset, it).toList() }
         if (elements.isEmpty()) {
-            val messageKey = when (elementKind) {
+            return when (elementKinds.singleOrNull()) {
-                CodeInsightUtils.ElementKind.EXPRESSION -> "cannot.refactor.no.expression"
+                CodeInsightUtils.ElementKind.EXPRESSION -> showErrorHintByKey("cannot.refactor.no.expression")
-                CodeInsightUtils.ElementKind.TYPE_ELEMENT -> "cannot.refactor.no.type"
+                CodeInsightUtils.ElementKind.TYPE_ELEMENT -> showErrorHintByKey("cannot.refactor.no.type")
+                else -> showErrorHint(file.project, editor, "Refactoring can't be performed on the selected code element", title)
             }
-            showErrorHintByKey(messageKey)
-            return
         }
 
-        selectTargetContainer(elements.toList())
+        selectTargetContainer(elements)
     }
 
     fun selectSingleElement() {
-        KotlinRefactoringUtil.selectElement(editor, file, false, elementKind) { expr ->
+        KotlinRefactoringUtil.selectElement(editor, file, false, elementKinds) { expr ->
             if (expr != null) {
                 selectTargetContainer(listOf(expr))
             }
             else {
                 if (!editor.selectionModel.hasSelection()) {
-                    if (elementKind == CodeInsightUtils.ElementKind.EXPRESSION) {
+                    if (elementKinds.singleOrNull() == CodeInsightUtils.ElementKind.EXPRESSION) {
                         val elementAtCaret = file.findElementAt(editor.caretModel.offset)
                         elementAtCaret?.getParentOfTypeAndBranch<KtProperty> { nameIdentifier }?.let {
                             return@selectElement selectTargetContainer(listOf(it))

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.kt

@@ -719,7 +719,9 @@ object KotlinIntroduceVariableHandler : RefactoringActionHandler {
         if (file !is KtFile) return
 
         try {
-            KotlinRefactoringUtil.selectElement(editor, file, CodeInsightUtils.ElementKind.EXPRESSION) { doRefactoring(project, editor, it as KtExpression?, null, null) }
+            KotlinRefactoringUtil.selectElement(editor, file, listOf(CodeInsightUtils.ElementKind.EXPRESSION)) {
+                doRefactoring(project, editor, it as KtExpression?, null, null)
+            }
         }
         catch (e: KotlinRefactoringUtil.IntroduceRefactoringException) {
             showErrorHint(project, editor, e.message!!)

idea/testData/refactoring/introduceTypeAlias/noTypeElement.kt.conflicts

@@ -1 +1 @@
-Cannot perform refactoring without a type
+Refactoring can't be performed on the selected code element

idea/tests/org/jetbrains/kotlin/idea/AbstractExpressionSelectionTest.java

@@ -25,6 +25,8 @@ import org.jetbrains.kotlin.idea.refactoring.KotlinRefactoringUtil;
 import org.jetbrains.kotlin.psi.KtFile;
 import org.jetbrains.kotlin.test.KotlinTestUtils;
 
+import java.util.Collections;
+
 public abstract class AbstractExpressionSelectionTest extends LightCodeInsightTestCase {
 
     public void doTestExpressionSelection(@NotNull String path) throws Exception {
@@ -32,13 +34,18 @@ public abstract class AbstractExpressionSelectionTest extends LightCodeInsightTe
         final String expectedExpression = KotlinTestUtils.getLastCommentInFile((KtFile) getFile());
 
         try {
-            KotlinRefactoringUtil.selectElement(getEditor(), (KtFile) getFile(), CodeInsightUtils.ElementKind.EXPRESSION, new KotlinRefactoringUtil.SelectElementCallback() {
+            KotlinRefactoringUtil.selectElement(
-                @Override
+                    getEditor(),
-                public void run(@Nullable PsiElement element) {
+                    (KtFile) getFile(),
-                    assertNotNull("Selected expression mustn't be null", element);
+                    Collections.singletonList(CodeInsightUtils.ElementKind.EXPRESSION),
-                    assertEquals(expectedExpression, element.getText());
+                    new KotlinRefactoringUtil.SelectElementCallback() {
-                }
+                        @Override
-            });
+                        public void run(@Nullable PsiElement element) {
+                            assertNotNull("Selected expression mustn't be null", element);
+                            assertEquals(expectedExpression, element.getText());
+                        }
+                    }
+            );
         }
         catch (KotlinRefactoringUtil.IntroduceRefactoringException e) {
             assertEquals(expectedExpression, "");

idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/AbstractExtractionTest.kt

@@ -114,7 +114,7 @@ abstract class AbstractExtractionTest() : KotlinLightCodeInsightFixtureTestCase(
             with (handler) {
                 val target = (file as KtFile).findElementByCommentPrefix("// TARGET:") as? KtNamedDeclaration
                 if (target != null) {
-                    KotlinRefactoringUtil.selectElement(fixture.getEditor(), file, true, CodeInsightUtils.ElementKind.EXPRESSION) { element ->
+                    KotlinRefactoringUtil.selectElement(fixture.getEditor(), file, true, listOf(CodeInsightUtils.ElementKind.EXPRESSION)) { element ->
                         invoke(fixture.getProject(), fixture.getEditor(), element as KtExpression, target)
                     }
                 }

idea/tests/org/jetbrains/kotlin/idea/refactoring/nameSuggester/KotlinNameSuggesterTest.kt

@@ -90,7 +90,7 @@ class KotlinNameSuggesterTest : LightCodeInsightFixtureTestCase() {
             if (withRuntime) {
                 ConfigLibraryUtil.configureKotlinRuntimeAndSdk(myModule, PluginTestCaseBase.mockJdk())
             }
-            KotlinRefactoringUtil.selectElement(myFixture.editor, file, CodeInsightUtils.ElementKind.EXPRESSION, object : KotlinRefactoringUtil.SelectElementCallback {
+            KotlinRefactoringUtil.selectElement(myFixture.editor, file, listOf(CodeInsightUtils.ElementKind.EXPRESSION), object : KotlinRefactoringUtil.SelectElementCallback {
                 override fun run(element: PsiElement?) {
                     val names = KotlinNameSuggester
                             .suggestNamesByExpressionAndType(element as KtExpression,