Extract utility functions from ExtractKotlinFunctionHandler

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

@@ -31,15 +31,12 @@ import org.jetbrains.jet.plugin.refactoring.extractFunction.ui.KotlinExtractFunc
 import org.jetbrains.jet.lang.psi.JetFile
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.jet.lang.psi.JetElement
-import org.jetbrains.jet.plugin.refactoring.getAllExtractionContainers
 import com.intellij.openapi.application.ApplicationManager
 import org.jetbrains.jet.lang.psi.psiUtil.getOutermostParentContainedIn
 import org.jetbrains.jet.plugin.refactoring.checkConflictsInteractively
 import org.jetbrains.jet.plugin.refactoring.executeWriteCommand
 import org.jetbrains.jet.lang.psi.JetBlockExpression
-import org.jetbrains.jet.lang.psi.JetClassBody
 import kotlin.test.fail
-import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
 import org.jetbrains.jet.plugin.refactoring.extractFunction.AnalysisResult.Status
 import com.intellij.openapi.ui.popup.JBPopupFactory
 import com.intellij.openapi.ui.MessageType
@@ -47,27 +44,10 @@ 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
-import org.jetbrains.jet.lang.psi.JetProperty
-import org.jetbrains.jet.lang.psi.JetParameterList
-import org.jetbrains.jet.lang.psi.JetClassInitializer
 import org.jetbrains.jet.lang.psi.JetFunctionLiteral
-import com.intellij.ide.util.PsiElementListCellRenderer
-import com.intellij.openapi.util.text.StringUtil
-import javax.swing.Icon
-import org.jetbrains.jet.plugin.refactoring.getPsiElementPopup
-import com.intellij.psi.PsiNamedElement
-import org.jetbrains.jet.plugin.util.collapseSpaces
-import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache
-import org.jetbrains.jet.lang.resolve.BindingContext
-import org.jetbrains.jet.lang.descriptors.FunctionDescriptor
-import org.jetbrains.jet.renderer.DescriptorRenderer
-import org.jetbrains.jet.lang.psi.JetPropertyAccessor
-import org.jetbrains.jet.lang.psi.JetClassOrObject
 import org.jetbrains.jet.plugin.util.psi.patternMatching.toRange
-import org.jetbrains.jet.lang.psi.JetMultiDeclaration
+import org.jetbrains.jet.plugin.refactoring.chooseContainerElementIfNecessary
+import org.jetbrains.jet.plugin.refactoring.getExtractionContainers
 
 public open class ExtractKotlinFunctionHandlerHelper {
     open fun adjustExtractionData(data: ExtractionData): ExtractionData = data
@@ -212,108 +192,24 @@ 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)?.let { declaration ->
-            stream(declaration) { it.getParentByType(javaClass<JetDeclaration>(), true) }.firstOrNull { it !is JetFunctionLiteral }
-        } ?: return Collections.emptyList()
-
-        val parent = declaration.getParent()?.let {
-            when (it) {
-                is JetProperty, is JetMultiDeclaration -> it.getParent()
-                is JetParameterList -> it.getParent()?.getParent()
-                else -> it
-            }
-        }
-        return when (parent) {
-            is JetFile -> Collections.singletonList(parent)
-            is JetClassBody -> {
-                element.getAllExtractionContainers(strict).filterIsInstance(javaClass<JetClassBody>())
-            }
-            else -> {
-                val enclosingDeclaration =
-                        PsiTreeUtil.getNonStrictParentOfType(parent, javaClass<JetDeclarationWithBody>(), javaClass<JetClassInitializer>())
-                val targetContainer = when (enclosingDeclaration) {
-                    is JetDeclarationWithBody -> enclosingDeclaration.getBodyExpression()
-                    is JetClassInitializer -> enclosingDeclaration.getBody()
-                    else -> null
-                }
-                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 = getContainers(parent, elements.size == 1)
+        val containers = parent.getExtractionContainers(elements.size == 1, allContainersEnabled)
         if (containers.empty) {
             noContainerError()
             return
         }
 
-        if (containers.size == 1 || ApplicationManager.getApplication()!!.isUnitTestMode()) {
+        chooseContainerElementIfNecessary(
-            onSelectionComplete(parent, elements, containers[0])
+                containers,
-            return
-        }
-
-        getPsiElementPopup(
                 editor,
-                containers.copyToArray(),
-                object: PsiElementListCellRenderer<JetElement>() {
-                    private fun JetElement.renderName(): String? {
-                        if (this is JetPropertyAccessor) {
-                            return (getParent() as JetProperty).renderName() + if (isGetter()) ".get" else ".set"
-                        }
-                        return (this as? PsiNamedElement)?.getName() ?: "<anonymous>"
-                    }
-
-                    private fun JetElement.renderDeclaration(): String? {
-                        val name = renderName()
-                        val descriptor = AnalyzerFacadeWithCache.getContextForElement(this)[BindingContext.DECLARATION_TO_DESCRIPTOR, this]
-                        val params = (descriptor as? FunctionDescriptor)?.let { descriptor ->
-                            descriptor.getValueParameters()
-                                    .map { DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(it.getType()) }
-                                    .joinToString(", ", "(", ")")
-                        } ?: ""
-                        return "$name$params"
-                    }
-
-                    private fun JetElement.renderText(): String {
-                        return StringUtil.shortenTextWithEllipsis(getText()!!.collapseSpaces(), 53, 0)
-                    }
-
-                    private fun JetElement.getRepresentativeElement(): JetElement {
-                        return when (this) {
-                            is JetBlockExpression -> (getParent() as? JetDeclarationWithBody) ?: this
-                            is JetClassBody -> getParent() as JetClassOrObject
-                            else -> this
-                        }
-                    }
-
-                    override fun getElementText(element: JetElement): String? {
-                        val representativeElement = element.getRepresentativeElement()
-                        return when (representativeElement) {
-                            is JetFile, is JetDeclarationWithBody, is JetClassOrObject -> representativeElement.renderDeclaration()
-                            else -> representativeElement.renderText()
-                        }
-                    }
-
-                    override fun getContainerText(element: JetElement?, name: String?): String? = null
-
-                    override fun getIconFlags(): Int = 0
-
-                    override fun getIcon(element: PsiElement?): Icon? =
-                            super.getIcon((element as? JetElement)?.getRepresentativeElement())
-                },
                 "Select target code block",
-                {
+                true,
-                    onSelectionComplete(parent, elements, it)
+                { it },
-                    true
+                { onSelectionComplete(parent, elements, it) }
-                }
+        )
-        ).showInBestPositionFor(editor)
     }
 
     fun selectMultipleExpressions() {
@@ -345,4 +241,4 @@ fun selectElements(
 
     editor.getScrollingModel().scrollToCaret(ScrollType.MAKE_VISIBLE)
     selectSingleExpression()
-}
+}

idea/src/org/jetbrains/jet/plugin/refactoring/jetRefactoringUtil.kt

@@ -62,16 +62,18 @@ import com.intellij.openapi.editor.colors.EditorColorsManager
 import com.intellij.ui.components.JBList
 import com.intellij.openapi.ui.popup.JBPopupAdapter
 import com.intellij.openapi.ui.popup.LightweightWindowEvent
-import com.intellij.openapi.editor.Document
 import com.intellij.psi.PsiDocumentManager
-import com.intellij.psi.PsiElementVisitor
-import com.intellij.psi.PsiWhiteSpace
-import com.intellij.openapi.util.Computable
-import com.intellij.openapi.ui.DialogWrapper
 import org.jetbrains.jet.lang.resolve.BindingContext
 import org.jetbrains.jet.lang.psi.psiUtil.getParentByType
 import org.jetbrains.jet.lang.psi.psiUtil.isAncestor
 import org.jetbrains.jet.plugin.caches.resolve.getLazyResolveSession
+import com.intellij.psi.PsiNamedElement
+import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache
+import org.jetbrains.jet.lang.descriptors.FunctionDescriptor
+import org.jetbrains.jet.renderer.DescriptorRenderer
+import com.intellij.openapi.util.text.StringUtil
+import org.jetbrains.jet.plugin.util.collapseSpaces
+import javax.swing.Icon
 
 /**
  * Replace [[JetSimpleNameExpression]] (and its enclosing qualifier) with qualified element given by FqName
@@ -157,7 +159,7 @@ public inline fun JetFile.createTempCopy(textTransform: (String) -> String): Jet
     return tmpFile
 }
 
-public fun PsiElement.getAllExtractionContainers(strict: Boolean): List<JetElement> {
+public fun PsiElement.getAllExtractionContainers(strict: Boolean = true): List<JetElement> {
     val containers = ArrayList<JetElement>()
 
     var element: PsiElement? = if (strict) getParent() else this
@@ -172,6 +174,38 @@ public fun PsiElement.getAllExtractionContainers(strict: Boolean): List<JetEleme
     return containers
 }
 
+public fun PsiElement.getExtractionContainers(strict: Boolean = true, includeAll: Boolean = false): List<JetElement> {
+    if (includeAll) return getAllExtractionContainers(strict)
+
+    val declaration = getParentByType(javaClass<JetDeclaration>(), strict)?.let { declaration ->
+        stream(declaration) { it.getParentByType(javaClass<JetDeclaration>(), true) }.firstOrNull { it !is JetFunctionLiteral }
+    } ?: return Collections.emptyList()
+
+    val parent = declaration.getParent()?.let {
+        when (it) {
+            is JetProperty, is JetMultiDeclaration -> it.getParent()
+            is JetParameterList -> it.getParent()?.getParent()
+            else -> it
+        }
+    }
+    return when (parent) {
+        is JetFile -> Collections.singletonList(parent)
+        is JetClassBody -> {
+            getAllExtractionContainers(strict).filterIsInstance(javaClass<JetClassBody>())
+        }
+        else -> {
+            val enclosingDeclaration =
+                    PsiTreeUtil.getNonStrictParentOfType(parent, javaClass<JetDeclarationWithBody>(), javaClass<JetClassInitializer>())
+            val targetContainer = when (enclosingDeclaration) {
+                is JetDeclarationWithBody -> enclosingDeclaration.getBodyExpression()
+                is JetClassInitializer -> enclosingDeclaration.getBody()
+                else -> null
+            }
+            if (targetContainer is JetBlockExpression) Collections.singletonList(targetContainer) else Collections.emptyList()
+        }
+    }
+}
+
 public fun Project.checkConflictsInteractively(conflicts: MultiMap<PsiElement, String>, onAccept: () -> Unit) {
     if (!conflicts.isEmpty()) {
         if (ApplicationManager.getApplication()!!.isUnitTestMode()) throw ConflictsInTestsException(conflicts.values())
@@ -202,21 +236,23 @@ public fun <T: Any> Project.executeWriteCommand(name: String, command: () -> T):
     return result!!
 }
 
-public fun <T : PsiElement> getPsiElementPopup(
+public fun <T, E: PsiElement> getPsiElementPopup(
         editor: Editor,
-        elements: Array<T>,
+        elements: List<T>,
-        renderer: PsiElementListCellRenderer<T>,
+        renderer: PsiElementListCellRenderer<E>,
-        title: String? = null,
+        title: String?,
+        highlightSelection: Boolean,
+        toPsi: (T) -> E,
         processor: (T) -> Boolean): JBPopup {
-    val highlighter = SelectionAwareScopeHighlighter(editor)
+    val highlighter = if (highlightSelection) SelectionAwareScopeHighlighter(editor) else null
 
-    val list = JBList(elements.toList())
+    val list = JBList(elements.map(toPsi))
     list.setCellRenderer(renderer)
     list.addListSelectionListener { e ->
-        highlighter.dropHighlight()
+        highlighter?.dropHighlight()
         val index = list.getSelectedIndex()
         if (index >= 0) {
-            highlighter.highlight(list.getModel()!!.getElementAt(index) as PsiElement)
+            highlighter?.highlight(list.getModel()!!.getElementAt(index) as PsiElement)
         }
     }
 
@@ -224,16 +260,14 @@ public fun <T : PsiElement> getPsiElementPopup(
         title?.let { setTitle(it) }
         renderer.installSpeedSearch(this, true)
         setItemChoosenCallback {
-            for (element in list.getSelectedValues()) {
+            val index = list.getSelectedIndex()
-                element.let {
+            if (index >= 0) {
-                    [suppress("UNCHECKED_CAST")]
+                processor(elements[index])
-                    processor(it as T)
-                }
             }
         }
         addListener(object: JBPopupAdapter() {
             override fun onClosed(event: LightweightWindowEvent?) {
-                highlighter.dropHighlight();
+                highlighter?.dropHighlight();
             }
         })
 
@@ -303,4 +337,88 @@ public fun JetElement.getContextForContainingDeclarationBody(): BindingContext?
         else -> null
     }
     return bodyElement?.let { getContainingJetFile().getLazyResolveSession().resolveToElement(it) }
+}
+
+public fun chooseContainerElement<T>(
+        containers: List<T>,
+        editor: Editor,
+        title: String,
+        highlightSelection: Boolean,
+        toPsi: (T) -> JetElement,
+        onSelect: (T) -> Unit) {
+    return getPsiElementPopup(
+            editor,
+            containers,
+            object : PsiElementListCellRenderer<JetElement>() {
+                private fun JetElement.renderName(): String {
+                    if (this is JetPropertyAccessor) {
+                        return (getParent() as JetProperty).renderName() + if (isGetter()) ".get" else ".set"
+                    }
+                    if (this is JetObjectDeclaration && this.isClassObject()) {
+                        return "Class object of ${getParentByType(javaClass<JetClassOrObject>(), true)?.renderName() ?: "<anonymous>"}"
+                    }
+                    return (this as? PsiNamedElement)?.getName() ?: "<anonymous>"
+                }
+
+                private fun JetElement.renderDeclaration(): String? {
+                    val name = renderName()
+                    val descriptor = AnalyzerFacadeWithCache.getContextForElement(this)[BindingContext.DECLARATION_TO_DESCRIPTOR, this]
+                    val params = (descriptor as? FunctionDescriptor)?.let { descriptor ->
+                        descriptor.getValueParameters()
+                                .map { DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(it.getType()) }
+                                .joinToString(", ", "(", ")")
+                    } ?: ""
+                    return "$name$params"
+                }
+
+                private fun JetElement.renderText(): String {
+                    return StringUtil.shortenTextWithEllipsis(getText()!!.collapseSpaces(), 53, 0)
+                }
+
+                private fun JetElement.getRepresentativeElement(): JetElement {
+                    return when (this) {
+                        is JetBlockExpression -> (getParent() as? JetDeclarationWithBody) ?: this
+                        is JetClassBody -> getParent() as JetClassOrObject
+                        else -> this
+                    }
+                }
+
+                override fun getElementText(element: JetElement): String? {
+                    val representativeElement = element.getRepresentativeElement()
+                    return when (representativeElement) {
+                        is JetFile, is JetDeclarationWithBody, is JetClassOrObject -> representativeElement.renderDeclaration()
+                        else -> representativeElement.renderText()
+                    }
+                }
+
+                override fun getContainerText(element: JetElement?, name: String?): String? = null
+
+                override fun getIconFlags(): Int = 0
+
+                override fun getIcon(element: PsiElement?): Icon? =
+                        super.getIcon((element as? JetElement)?.getRepresentativeElement())
+            },
+            title,
+            highlightSelection,
+            toPsi,
+            {
+                onSelect(it)
+                true
+            }
+    ).showInBestPositionFor(editor)
+}
+
+public fun chooseContainerElementIfNecessary<T>(
+        containers: List<T>,
+        editor: Editor,
+        title: String,
+        highlightSelection: Boolean,
+        toPsi: (T) -> JetElement,
+        onSelect: (T) -> Unit
+) {
+    when {
+        containers.empty -> return
+        containers.size == 1 || ApplicationManager.getApplication()!!.isUnitTestMode() -> onSelect(containers.first())
+        else -> chooseContainerElement(containers, editor, title, highlightSelection, toPsi, onSelect)
+    }
 }