Extract Function: Improve container selection UI

annotations/com/intellij/codeInsight/unwrap/annotations.xml

@@ -0,0 +1,20 @@
+<root>
+  <item
+      name='com.intellij.codeInsight.unwrap.RangeSplitter java.util.List&lt;com.intellij.openapi.util.TextRange&gt; split(com.intellij.openapi.util.TextRange, java.util.List&lt;com.intellij.openapi.util.TextRange&gt;)'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
+  <item
+      name='com.intellij.codeInsight.unwrap.ScopeHighlighter ScopeHighlighter(com.intellij.openapi.editor.Editor, com.intellij.util.NotNullFunction&lt;com.intellij.psi.PsiElement,com.intellij.openapi.util.TextRange&gt;)'>
+    <annotation name='kotlin.jvm.KotlinSignature'>
+      <val name="value" val="&quot;fun ScopeHighlighter(editor: Editor, ranger: NotNullFunction&lt;PsiElement, TextRange&gt;?)&quot;"/>
+    </annotation>
+  </item>
+  <item
+      name='com.intellij.codeInsight.unwrap.ScopeHighlighter void highlight(com.intellij.psi.PsiElement, java.util.List&lt;com.intellij.psi.PsiElement&gt;) 0'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
+  <item
+      name='com.intellij.codeInsight.unwrap.ScopeHighlighter void highlight(com.intellij.psi.PsiElement, java.util.List&lt;com.intellij.psi.PsiElement&gt;) 1'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
+</root>

annotations/com/intellij/ide/util/annotations.xml

@@ -5,4 +5,7 @@
   <item name='com.intellij.ide.util.FileStructurePopup com.intellij.ui.treeStructure.filtered.FilteringTreeBuilder getTreeBuilder()'>
     <annotation name='org.jetbrains.annotations.NotNull'/>
   </item>
+  <item name='com.intellij.ide.util.PsiElementListCellRenderer java.lang.String getElementText(T) 0'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
 </root>

annotations/com/intellij/openapi/editor/colors/annotations.xml

@@ -0,0 +1,5 @@
+<root>
+  <item name='com.intellij.openapi.editor.colors.EditorColorsManager com.intellij.openapi.editor.colors.EditorColorsManager getInstance()'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
+</root>

idea/src/org/jetbrains/jet/plugin/JetIconProvider.java

@@ -123,6 +123,9 @@ public class JetIconProvider extends IconProvider {
                 return JetIcons.FUNCTION;
             }
         }
+
+        if (psiElement instanceof JetFunctionLiteral) return JetIcons.LAMBDA;
+
         if (psiElement instanceof JetClass) {
             JetClass jetClass = (JetClass) psiElement;
             if (jetClass.isTrait()) {

idea/src/org/jetbrains/jet/plugin/refactoring/JetRefactoringUtil.java

@@ -57,6 +57,7 @@ import org.jetbrains.jet.plugin.JetBundle;
 import org.jetbrains.jet.plugin.codeInsight.CodeInsightUtils;
 import org.jetbrains.jet.plugin.codeInsight.DescriptorToDeclarationUtil;
 import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache;
+import org.jetbrains.jet.plugin.util.UtilPackage;
 import org.jetbrains.jet.renderer.DescriptorRenderer;
 
 import javax.swing.*;
@@ -509,11 +510,7 @@ public class JetRefactoringUtil {
     }
 
     public static String getExpressionShortText(@NotNull JetElement element) { //todo: write appropriate implementation
-        String expressionText = element.getText();
+        return UtilPackage.collapseSpaces(StringUtil.shortenTextWithEllipsis(element.getText(), 53, 0));
-        if (expressionText.length() > 20) {
-            expressionText = expressionText.substring(0, 17) + "...";
-        }
-        return expressionText;
     }
 
     @Nullable

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

@@ -32,8 +32,6 @@ 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.refactoring.IntroduceTargetChooser
-import com.intellij.openapi.util.Pass
 import com.intellij.openapi.application.ApplicationManager
 import org.jetbrains.jet.lang.psi.psiUtil.getOutermostParentContainedIn
 import org.jetbrains.jet.plugin.refactoring.checkConflictsInteractively
@@ -56,6 +54,17 @@ 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
 
 public class ExtractKotlinFunctionHandler(public val allContainersEnabled: Boolean = false) : RefactoringActionHandler {
@@ -222,33 +231,61 @@ fun selectElements(
             return
         }
 
-        IntroduceTargetChooser.showChooser(
+        getPsiElementPopup(
                 editor,
-                containers,
+                containers.copyToArray(),
-                object: Pass<JetElement>() {
+                object: PsiElementListCellRenderer<JetElement>() {
-                    override fun pass(targetContainer: JetElement?) {
+                    private fun JetElement.renderName(): String? {
-                        if (targetContainer == null) {
+                        if (this is JetPropertyAccessor) {
-                            noContainerError()
+                            return (getParent() as JetProperty).renderName() + if (isGetter()) ".get" else ".set"
-                            return
                         }
+                        return (this as? PsiNamedElement)?.getName() ?: "<anonymous>"
+                    }
 
-                        onSelectionComplete(parent, elements, targetContainer)
+                    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",
                 {
-                    when (it) {
+                    onSelectionComplete(parent, elements, it)
-                        is JetFile -> "File: ${it.getName()}"
+                    true
-                        is JetBlockExpression -> {
+                }
-                            (it.getParent() as? JetDeclarationWithBody)?.getText() ?: "...${it.getStatements().first?.getText()}"
+        ).showInBestPositionFor(editor)
-                        }
-                        is JetClassBody -> {
-                            (it.getParent() as? JetClassOrObject)?.getText()
-                        }
-                        else -> it?.getText()
-                    }
-                },
-                "Select target code block"
-        )
     }
 
     fun selectMultipleExpressions() {

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

@@ -46,7 +46,22 @@ import com.intellij.refactoring.ui.ConflictsDialog
 import com.intellij.util.containers.MultiMap
 import com.intellij.openapi.command.CommandProcessor
 import org.jetbrains.jet.lang.psi.codeFragmentUtil.skipVisibilityCheck
-import com.intellij.openapi.ui.DialogWrapper
+import com.intellij.ide.util.PsiElementListCellRenderer
+import com.intellij.openapi.ui.popup.JBPopup
+import com.intellij.openapi.ui.popup.PopupChooserBuilder
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.TextRange
+import com.intellij.codeInsight.unwrap.RangeSplitter
+import com.intellij.codeInsight.unwrap.UnwrapHandler
+import com.intellij.openapi.editor.markup.TextAttributes
+import com.intellij.openapi.editor.markup.HighlighterTargetArea
+import com.intellij.openapi.editor.markup.RangeHighlighter
+import java.util.Collections
+import com.intellij.openapi.editor.colors.EditorColors
+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
 
 /**
  * Replace [[JetSimpleNameExpression]] (and its enclosing qualifier) with qualified element given by FqName
@@ -165,3 +180,73 @@ public fun Project.executeWriteCommand(name: String, command: () -> Unit) {
             this, { ApplicationManager.getApplication()!!.runWriteAction(command) }, name, null
     )
 }
+
+public fun <T : PsiElement> getPsiElementPopup(
+        editor: Editor,
+        elements: Array<T>,
+        renderer: PsiElementListCellRenderer<T>,
+        title: String? = null,
+        processor: (T) -> Boolean): JBPopup {
+    val highlighter = SelectionAwareScopeHighlighter(editor)
+
+    val list = JBList(elements.toList())
+    list.setCellRenderer(renderer)
+    list.addListSelectionListener { e ->
+        highlighter.dropHighlight()
+        val index = list.getSelectedIndex()
+        if (index >= 0) {
+            highlighter.highlight(list.getModel()!!.getElementAt(index) as PsiElement)
+        }
+    }
+
+    return with(PopupChooserBuilder(list)) {
+        title?.let { setTitle(it) }
+        renderer.installSpeedSearch(this, true)
+        setItemChoosenCallback {
+            for (element in list.getSelectedValues()) {
+                element.let {
+                    [suppress("UNCHECKED_CAST")]
+                    processor(it as T)
+                }
+            }
+        }
+        addListener(object: JBPopupAdapter() {
+            override fun onClosed(event: LightweightWindowEvent?) {
+                highlighter.dropHighlight();
+            }
+        })
+
+        createPopup()
+    }
+}
+
+public class SelectionAwareScopeHighlighter(val editor: Editor) {
+    private val highlighters = ArrayList<RangeHighlighter>()
+
+    private fun addHighlighter(r: TextRange, attr: TextAttributes) {
+        highlighters.add(
+                editor.getMarkupModel().addRangeHighlighter(
+                        r.getStartOffset(),
+                        r.getEndOffset(),
+                        UnwrapHandler.HIGHLIGHTER_LEVEL,
+                        attr,
+                        HighlighterTargetArea.EXACT_RANGE
+                )
+        )
+    }
+
+    public fun highlight(wholeAffected: PsiElement) {
+        dropHighlight()
+
+        val attributes = EditorColorsManager.getInstance().getGlobalScheme().getAttributes(EditorColors.SEARCH_RESULT_ATTRIBUTES)!!
+        val selectedRange = with(editor.getSelectionModel()) { TextRange(getSelectionStart(), getSelectionEnd()) }
+        for (r in RangeSplitter.split(wholeAffected.getTextRange()!!, Collections.singletonList(selectedRange))) {
+            addHighlighter(r, attributes)
+        }
+    }
+
+    public fun dropHighlight() {
+        highlighters.forEach { it.dispose() }
+        highlighters.clear()
+    }
+}

idea/src/org/jetbrains/jet/plugin/util/stringUtil.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.plugin.util
+
+fun String.collapseSpaces(): String {
+    val builder = StringBuilder()
+    var haveSpaces = false
+    for (c in this) {
+        if (c.isWhitespace()) {
+            haveSpaces = true
+        }
+        else {
+            if (haveSpaces) {
+                builder.append(" ")
+                haveSpaces = false
+            }
+            builder.append(c)
+        }
+    }
+    return builder.toString()
+}